Source code for torx.analysis.lineouts.xpoint_lineouts_m

"""Provides lineouts starting from the X-point."""
from torx.analysis.vector_field_tracer_m import VectorFieldTracer
from torx.analysis.lineouts import Lineout
from torx.autodoc_decorators_m import autodoc_class



@autodoc_class
class XPointLineout(Lineout):
    """Interpolates to a line perpendicular to Bpol close to the X-Point."""



    def __init__(
        self,
        grid,
        equi,
        direction: str="lfs",
        xpoint_dist: float=1e-3,
        t_max: float=1000.0,
        find_points_on_grid: bool=True,
        n_samples: int=200,
        tolerance: float=1e-4,
        **integrator_kwargs
    ):
        """
        Initialize the X-point lineout.

        Uses the X-Point information stored in the equilibrium to perform a
        RZ-trace perpendicular to the poloidal magnetic field. Via the
        'direction' keyword four different lineout types are available:

        1) 'lfs':  The lineout originates slightly right of the X-Point and
                   extends into the low-field side.
        2) 'hfs':  The lineout originates slightly left of the X-Point and
                   extends into the high-field side.
        3) 'up':   The lineout originates slightly above the X-Point and
                   extends to lower values of rho_poloidal.
        3) 'down': The lineout originates slightly below the X-Point and
                   extends to lower values of rho_poloidal.

        Note
        ----
        The above mentioned initialization methods are well defined only
        for a standard single null equilibrium.

        Note
        ----
        Due to possible additional flip in Z-direction it is not possible
        to automatically distinguish the X-Point lineout going into the private
        flux and core region. Hence only 'up' and 'down' are available.
        """
        self.xpt_r_norm = float(equi.x_point_r_norm.values)
        self.xpt_z_norm = float(equi.x_point_z_norm.values)

        self.boundary_polygon = equi.get_boundary_polygon()

        if direction in ["lfs", "hfs"]:
            self.z_initial = self.xpt_z_norm
            vecfield = equi.magfield_vector_radial(grid.r_s, grid.z_s)
            if direction == "lfs":
                self.r_initial = self.xpt_r_norm + xpoint_dist
            if direction == "hfs":
                self.r_initial = self.xpt_r_norm - xpoint_dist

        elif direction in ["up", "down"]:
            self.r_initial = self.xpt_r_norm
            vecfield = -equi.magfield_vector_radial(grid.r_s, grid.z_s)
            if direction == "up":
                self.z_initial = self.xpt_z_norm + xpoint_dist
            elif direction == "down":
                self.z_initial = self.xpt_z_norm - xpoint_dist
        else:
            raise ValueError("Invalid direction indication, see doc-string!")

        self._integrator = VectorFieldTracer(vecfield)
        self._tolerance = tolerance

        r, z = self._find_points_from_rz_trace(t_max=t_max, **integrator_kwargs)
        self.r_points = r
        self.z_points = z

        super().__init__(self.r_points, self.z_points)

        if find_points_on_grid:
            self.find_points_on_grid(grid, n_samples=n_samples)


    def _find_points_from_rz_trace(
        self,
        t_max: float=1000.0,
        **integrator_kwargs
    ):
        """
        Find points along the radial unit vector field.

        Uses parametric trace in the RZ-plane.
        """
        def in_vessel(t, state):
            x, y = state[0], state[1]
            return +1.0 if self.boundary_polygon.point_inside(x, y) else -1.0

        in_vessel.terminal = True
        assert t_max > 0.0, \
            f"Should always use a positive value for timeout_trace"

        sol = self._integrator.rz_integration(
            self.r_initial,
            self.z_initial,
            t_max=t_max,
            events=in_vessel,
            tolerance=self._tolerance,
            **integrator_kwargs
        )

        return sol.y[0], sol.y[1]