"""Defines the basic lineout class."""
import numpy as np
import xarray as xr
from torx import make_xarray
from torx.geometry import compute_arc_length
from torx.analysis.magnetic_field_tracer_m import MagneticFieldTracer
from torx.normalization.normalization_m import Normalization
from torx.equilibrium.equilibrium_m import EquilibriumBaseClass
from torx.analysis.lineouts._parametric_spline_m import ParametricSpline
from torx.grid.grid_2d_m import Grid2D
from torx import Quantity
from torx.analysis.lineouts.lineout_base_m import LineoutBase
from .on_grid_m import on_grid
from torx.autodoc_decorators_m import autodoc_class
[docs]
@autodoc_class
class Lineout(LineoutBase):
"""Allows to map from grid points to points on lines or curves."""
[docs]
def __init__(
self,
r_points: np.ndarray,
z_points: np.ndarray,
periodic: bool=False,
smoothing: float=0.0,
):
"""
Initialize the Lineout object from a specified set of r and z points.
If smoothing if > 0.0, then the input points will be smoothed.
Typically smoothing required << 1.
It is recommended that the fit is checked.
"""
r_points, z_points = np.array(r_points), np.array(z_points)
assert r_points.ndim == 1
assert r_points.shape == z_points.shape
self.r_source, self.z_source = r_points, z_points
# Default use a third order spline, unless not enough points are given in
# which case drop the order
order = min(3, r_points.size - 1)
self.curve = ParametricSpline(
sample_points=[r_points, z_points],
periodic=periodic,
smoothing=smoothing,
order=order,
)
[docs]
def find_points_on_grid(
self,
grid: Grid2D,
n_samples: int=500,
epsilon: float=1e-6,
):
"""
Find equally spaced (R, Z) points along the lineout.
Note that 'curve' must have a single segment on the grid -- this
algorithm will fail if there is an off grid segment between two
on-grid segments.
n_samples * precision test points will be used to check which values
of the parameter are on the grid epsilon shifts the endpoints slightly
inwards to ensure that the returned points are actually on the grid.
"""
# Make many sample points, and check which ones are on the grid
t_tests = np.linspace(
self.curve.t_min, self.curve.t_max, num=n_samples, endpoint=True
)
r_tests, z_tests = self.curve(t_tests)
t_on_grid = on_grid(grid.r_u.values, grid.z_u.values,
r_tests, z_tests, grid.spacing)
self.r_points, self.z_points = r_tests[t_on_grid], z_tests[t_on_grid]
self.make_interpolation_matrix(grid)
self.setup_perpendicular_gradient(grid)
[docs]
def omp_mapped_distance(self, omp_map) -> Quantity:
"""Calculate the OMP-mapped radial distance."""
values = omp_map(self.r_points, self.z_points)
return values.values * values.norm
[docs]
def target_mapped_distance(
self,
equi: EquilibriumBaseClass,
direction: str="fwd",
phi_initial: float=0.0,
phi_to_max: float=100,
max_flux_surface_wander: float=1e-3,
find_sep_dist: bool=False,
**integrator_kwargs,
):
"""
Return a dataset with coordinates of the lineout mapped to the target.
Additionally returns the distance of those points from
the strike line. The tracing is performed either 'fwd' or 'bwd' along
the magnetic field, controlled by the 'direction' parameter. The default
is 'both', tracing forward and backward separately.
"""
r_target, z_target = self.map_to_target(
equi,
direction=direction,
phi_initial=phi_initial,
phi_to_max=phi_to_max,
max_flux_surface_wander=max_flux_surface_wander,
**integrator_kwargs
)
x_mapped = compute_arc_length(r_target, z_target, skipna=True)
s_mapped = make_xarray(x_mapped,
dims="interp_points",
name=f"s_mapped_{direction}")
# Find the separatrix index to measure s_mapped from strike line
if find_sep_dist:
rho = self.rho(equi)
sep_idx = np.argmin(np.abs(rho.values - 1.0))
s_mapped -= s_mapped[sep_idx]
return s_mapped
[docs]
def map_to_target(
self,
equi: EquilibriumBaseClass,
direction: str="fwd",
phi_initial: float=0.0,
phi_to_max: float=100,
max_flux_surface_wander: float=1e-3,
**integrator_kwargs,
) -> np.ndarray:
"""Map points of the lineout to the target."""
assert direction in ["fwd", "bwd"],\
"Can only trace 'fwd' or 'bwd' along magnetic field!"
if not (
hasattr(self, f"r_target_{direction}") and
hasattr(self, f"r_target_{direction}")
):
tracer = MagneticFieldTracer(
equi, max_flux_surface_wander=max_flux_surface_wander
)
r_target = np.full(self.curve_size, np.nan)
z_target = np.full(self.curve_size, np.nan)
for i, [r,z] in enumerate(zip(self.r_points, self.z_points)):
sol = tracer.toroidal_integration(r, z,
direction=direction,
phi_initial=phi_initial,
phi_to_max=phi_to_max,
check_out_of_bounds=True,
**integrator_kwargs)
if (sol.status == 1):
r_target[i] = sol.y[0][-1]
z_target[i] = sol.y[1][-1]
else:
continue
r_target = make_xarray(r_target,
dims="interp_points",
name=f"r_target_{direction}")
z_target = make_xarray(z_target,
dims="interp_points",
name=f"z_target_{direction}")
setattr(self, f"r_target_{direction}", r_target)
setattr(self, f"z_target_{direction}", z_target)
return (
getattr(self, f"r_target_{direction}"),
getattr(self, f"z_target_{direction}")
)
