Source code for torx.analysis.lineouts.curve_lineouts_m

"""Provides functionality of curved lineouts."""
import xarray as xr
from pathlib import Path
from warnings import warn
import matplotlib.pyplot as plt

from torx.grid.grid_2d_m import Grid2D
from torx.analysis.lineouts.lineout_m import Lineout
from torx.analysis.contour_finding_m import find_contours
from torx.equilibrium.equilibrium_m import EquilibriumBaseClass
from torx.equilibrium.numerical_m import NumericalEquilibrium
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def penalization_contour(
    grid: Grid2D,
    equi: EquilibriumBaseClass,
    penalization_characteristic: xr.DataArray,
    level: float,
    contour_index: int = None,
    epsilon: float = 1e-10,
    smoothing: float = 0.1,
    n_samples: int = 500,
    grid_epsilon: float = 1e-6,
    shift_inward: int = 0,
):
    """
    Return a lineout along a contour of the penalization characteristic.

    Should use filled_penalization_characteristic to get the
    penalization_characteristic filled to include ghost points.

    Smoothing is recommended, since this removes the grid staircase pattern
    of the contour.
    """
    assert (
        "R" in penalization_characteristic.dims
        and "Z" in penalization_characteristic.dims
    ), "Call vector_to_matrix on\
        penalization_characteristic before calling penalization_contour"

    # Shift the level slightly in from [0, 1] to ensure a clear contour
    level = max(level, epsilon)
    level = min(level, 1.0 - epsilon)

    contours = find_contours(
        grid.r_s, grid.z_s, penalization_characteristic, level=level
    )

    if isinstance(equi, NumericalEquilibrium) and equi._flipped_Z:
        # Reverse the contour list is using flipped_Z, since contours are returned from bottom to top
        contours.reverse()

    if contour_index is None:
        plt.pcolormesh(
            grid.r_s, grid.z_s, penalization_characteristic, shading="nearest"
        )
        for i, contour in enumerate(contours):
            r_points, z_points = contour.T

            plt.plot(r_points, z_points, label=f"contour_index = {i}")
        plt.legend()
        plt.show()

        warn("Select a contour index for penalization_contour")

    else:
        r_points, z_points = contours[contour_index].T

    if shift_inward != 0:
        r_points = r_points[shift_inward : -shift_inward]
        z_points = z_points[shift_inward : -shift_inward]

    lineout = Lineout(r_points, z_points, smoothing=smoothing)
    lineout.find_points_on_grid(grid, n_samples=n_samples, epsilon=grid_epsilon)

    return lineout