Source code for torx.analysis.lineouts.omp_map_m

"""Functions for parameterizing a lineout curve."""
import numpy as np
from scipy.interpolate import make_interp_spline

from torx import make_xarray
from torx.grid.grid_2d_m import Grid2D
from torx.equilibrium.equilibrium_m import EquilibriumBaseClass
from torx.normalization.normalization_m import Normalization
from torx.analysis.lineouts.chord_lineouts_m import outboard_zaxis_chord
from torx import Quantity, Dimensionless
from torx.autodoc_decorators_m import autodoc_class



@autodoc_class
class OutboardMidplaneMap:
    """
    Interpolates rho from an arbitrary line to the equivalent distance on omp.

    This is useful for accounting for flux expansion when comparing fall-off
    lengths at different positions.
    """



    def __init__(
        self,
        grid: Grid2D,
        equi: EquilibriumBaseClass,
        norm: Normalization,
        separatrix_index: int = 0,
    ):
        """Initialize the omp map interpolator."""
        lineout = outboard_zaxis_chord(grid, equi)
        self.rho_points = equi.normalized_flux_surface_label(
            lineout.r_points, lineout.z_points, grid=False
        )
        self.arc_points = lineout.poloidal_arc_length(norm=norm)

        omp_map = make_interp_spline(self.rho_points, self.arc_points)
        # Call again to make the separatrix have a value of '0.0'
        self.separatrix_arc = omp_map(1.0)
        omp_map = make_interp_spline(
            self.rho_points, self.arc_points - self.separatrix_arc
        )

        self.equi = equi
        self.omp_map = omp_map




    def __call__(self, r_points: np.array, z_points: np.array):
        """
        Map points back to omp and calculate the equiv. distance to the sep.

        If the curve includes points in the private flux region, the
        rho < 1.0 (private flux region) or equivalently
        omp_mapped_distance < 0.0 is mapped to the core, which may have
        significantly different flux-surface expansion.
        """
        rho_points = self.equi.normalized_flux_surface_label(
            r_points, z_points, grid=False
        )
        return self.convert_rho_to_distance(rho_points=rho_points)




    def convert_rho_to_distance(self, rho_points: np.array):
        """Map an array of rho to omp equivalent distance to the separatrix."""
        return make_xarray(
            self.omp_map(rho_points),
            norm=self.arc_points.norm,
            name="OMP-mapped distance to sep.",
        )




    def convert_distance_to_rho(self, arc_points: Quantity):
        """Invert the omp map."""
        omp_inv = make_interp_spline(
            self.arc_points - self.separatrix_arc, self.rho_points
        )

        arc_norm = (arc_points / self.arc_points.norm).to("")

        return make_xarray(
            omp_inv(arc_norm), norm=Dimensionless,
            name="Normalized flux-surface label"
        )