Source code for torx.analysis.lineouts._helicity_shift_m

"""
Allows to use the symmetry angle to correct for the helicity.

Additionally allows to perform an upscaling and shift back to the
original coordinates.
"""
from scipy.interpolate import interp1d, RectBivariateSpline
import numpy as np
import xarray as xr


def helicity_shift(
    self, values: xr.DataArray, toroidal_angle: xr.DataArray, reverse: bool=False
):
    """
    Shift helicity based on angles such that it is reversed or not.

    Given a 2D xarray of values, which is of shape(poloidal_angle, toroidal_angle),
    and the corresponding toroidal angles and symmetry angle
    (shape(poloidal_angle)), return the values such that the helicity is
    either subtracted (reverse=False) or added (reverse=True).

    toroidal_angle = values.phi * 2.0 * np.pi / values.phi.size.
    """
    shifted_values = np.zeros_like(values)
    shift_direction = +1.0 if not reverse else -1.0

    # Iterate over the planes
    for i, phi in enumerate(toroidal_angle.values):
        # Shifted the symmetry angle to account for the helicity
        shifted_symmetry_angle = np.mod(
            self.symmetry_angle + shift_direction * phi / self.q_factor, 2.0 * np.pi
        )

        # Interpolate the original values to the new shifted values
        shifted_values[i, :] = interp1d(
            self.symmetry_angle,
            values.isel(phi=i).load(),
            kind="cubic",
            fill_value="extrapolate",
            bounds_error=False,
        )(shifted_symmetry_angle)

    return xr.DataArray(
        shifted_values, dims=values.dims, coords=values.coords, attrs=values.attrs
    )


def upscale_toroidal_resolution(
    self,
    values: xr.DataArray,
    original_toroidal_angle: xr.DataArray,
    output_toroidal_resolution: int = 500,
):
    """
    Scale toroidal resolution to given value.

    Uses cubic spline interpolation along the field line to find the value
    of "values" at inter-plane positions. "values" should be the interpolated
    values, not the full field (since the q-profile is not constant).

    This is made possible by helicity-shifting the input array such that the
    fieldlines are vertical, using a bicubic spline
    interpolation to fill the values, and then reversing the helicity-shift.
    """
    values_shifted = self.helicity_shift(values, original_toroidal_angle)

    interpolator = RectBivariateSpline(
        self.symmetry_angle, original_toroidal_angle, values_shifted.T
    )
    toroidal_angle_upscaled = xr.DataArray(
        np.linspace(0, 2.0 * np.pi,
                    num=output_toroidal_resolution,
                    endpoint=False),
        dims=original_toroidal_angle.dims,
    )

    values_upscaled = xr.DataArray(
        interpolator(self.symmetry_angle, toroidal_angle_upscaled).T,
        dims=["phi", "interp_points"],
        coords={"phi": toroidal_angle_upscaled, "tau": values_shifted.tau},
        attrs=values_shifted.attrs,
    )

    return self.helicity_shift(values_upscaled, toroidal_angle_upscaled, reverse=True)