"""Provides lineouts that are along flux-surfaces."""
import numpy as np
import xarray as xr
from torx.analysis.lineouts.lineout_base_m import LineoutBase
from torx.analysis.magnetic_field_tracer_m import (
MagneticFieldTracer,
polar_to_cart_chord
)
from torx.equilibrium.equilibrium_m import EquilibriumBaseClass
from torx.grid.grid_2d_m import Grid2D
from torx.autodoc_decorators_m import autodoc_class
[docs]
@autodoc_class
class FluxSurfaceLineout(LineoutBase):
"""
Interpolates to a flux surface.
Provides properties which assist in analysis like the symmetry angle.
"""
[docs]
@classmethod
def initialize_from_rho_value(
cls,
grid: Grid2D,
equi: EquilibriumBaseClass,
rho: float,
theta_chord: float = 0.0,
tolerance: float=1e-4,
max_flux_surface_wander: float=1e-3,
**integrator_kwargs,
):
"""
Initialize the flux surface lineout from a given rho poloidal value.
Uses a chord from the magnetic axis to search for (R, Z) points which
have the correct rho.
This initialization will not work for the private-flux region, since it
assumes that rho < 1 implies closed fieldline region.
"""
r_initial, z_initial, _ = polar_to_cart_chord(equi, rho=rho, theta=theta_chord)
# Make r_initial and z_initial into scalars
assert r_initial.size == 1 and z_initial.size == 1
r_initial, z_initial = r_initial[0], z_initial[0]
flux_surface_label = equi.normalized_flux_surface_label(r_initial, z_initial)
closed = flux_surface_label < 1.0
return cls(
grid,
equi,
r_initial,
z_initial,
closed_fieldlines=closed,
tolerance=tolerance,
max_flux_surface_wander=max_flux_surface_wander,
**integrator_kwargs,
)
[docs]
def __init__(
self,
grid: Grid2D,
equi: EquilibriumBaseClass,
r_initial: float,
z_initial: float,
timeout_trace: float = 100.0,
closed_fieldlines: bool = False,
curve_size: int = 1000,
tolerance: float=1e-4,
max_flux_surface_wander: float=1e-3,
**integrator_kwargs,
):
"""Initialize the flux surface lineout from a fieldline tracing."""
self.integrator = MagneticFieldTracer(equi, max_flux_surface_wander)
self.curve_size = curve_size
self._closed = closed_fieldlines
self._tolerance = tolerance
self.r_initial = np.squeeze(r_initial)
self.z_initial = np.squeeze(z_initial)
self.flux_surface_label = equi.normalized_flux_surface_label(
self.r_initial, self.z_initial
)
if self._closed:
# Use the poloidal integration, which returns the symmetry angle
# as a parameter
(
self.poloidal_angle,
solution,
max_symmetry_angle,
) = self.find_points_from_poloidal_trace(**integrator_kwargs)
# Shift the poloidal angle so that the outboard midplane is always
# at poloidal_angle = 0.0
self.poloidal_angle = np.mod(
self.poloidal_angle - self.integrator.theta(r_initial, z_initial),
2.0 * np.pi,
)
self.r_points, self.z_points, self.symmetry_angle = solution
# Divide the symmetry angle by the q-factor, so that it has values
# between 0 and 2*np.pi
self.q_factor = max_symmetry_angle / (2.0 * np.pi)
self.symmetry_angle /= self.q_factor
else:
# Use the toroidal integration, which returns the parallel trace length
# as a parameter
self.toroidal_angle, solution = self.find_points_from_toroidal_trace(
timeout_trace, **integrator_kwargs
)
# Note that the toroidal angle has an arbitrary added scalar.
# Typically toroidal_angle = 0 corresponds to the outboard midplane
# (if you use initialize_from_rho_value and don't set theta_chord)
# If you'd prefer to always have a divertor plate at zero, you can use
# self.toroidal_angle -= np.min(self.toroidal_angle)
self.r_points, self.z_points, self.parallel_distance = solution
self.make_interpolation_matrix(grid)
self.setup_perpendicular_gradient(grid)
from torx.analysis.lineouts._helicity_shift_m import (
helicity_shift,
upscale_toroidal_resolution,
)
[docs]
def find_points_from_toroidal_trace(
self, phi_max: float = 100.0, **integrator_kwargs
):
"""
Find points along the flux surface using toroidal integration.
The result will be equally spaced in the toroidal angle.
Note that this returns the parallel trace length as a result
The toroidal trace is performed only inside the divertor polygon.
"""
assert (
phi_max > 0.0
), f"Should always use a positive value for timeout_trace"
forward_trace = self.integrator.toroidal_integration(
self.r_initial,
self.z_initial,
0.0,
phi_max,
direction="fwd",
tolerance=self._tolerance,
check_out_of_bounds=True,
**integrator_kwargs
)
backward_trace = self.integrator.toroidal_integration(
self.r_initial,
self.z_initial,
0.0,
phi_max,
direction="bwd",
tolerance=self._tolerance,
check_out_of_bounds=True,
**integrator_kwargs
)
assert (
forward_trace.status == 1 and backward_trace.status == 1
), \
(
f"Toroidal trace did not terminate at the divertor polygon. \
Consider increasing timeout_trace. \nStatus forward:\n \
{forward_trace.status}\nStatus backward:\n{backward_trace.status}"
)
return self._combine_solutions(forward_trace, backward_trace)
def _combine_solutions(self, forward_trace, backward_trace):
"""
Combine two trace solutions.
Assumes that backward_trace is for negative trace angles and
forward trace is for positive trace angles.
The use of positive and negative angles is actually not strictly
necessary, but this prevents use from possibly hitting something
weird by extrapolating a long way with the dense output.
"""
assert (
np.array(forward_trace.t_events).size == 1
and np.array(backward_trace.t_events).size == 1
)
assert (
backward_trace.t_events[0][0] < 0.0 and forward_trace.t_events[0][0] > 0.0
)
toroidal_angle = np.linspace(
backward_trace.t_events[0][0],
forward_trace.t_events[0][0],
num=self.curve_size,
)
positive_angle = np.maximum(toroidal_angle, np.zeros_like(toroidal_angle))
negative_angle = np.minimum(toroidal_angle, np.zeros_like(toroidal_angle))
assert np.min(positive_angle) >= 0.0 and np.max(negative_angle) <= 0.0
combined_solution = np.where(
np.broadcast_to(toroidal_angle < 0.0, (3, toroidal_angle.size)),
backward_trace.sol(negative_angle),
forward_trace.sol(positive_angle),
)
return toroidal_angle, combined_solution
[docs]
def find_points_from_poloidal_trace(self, **integrator_kwargs):
"""
Use poloidal integration to find points along the flux surface.
Traces for exactly 2*pi, to cover one full closed flux surface.
"""
poloidal_trace = self.integrator.poloidal_integration(
r_initial=self.r_initial,
z_initial=self.z_initial,
theta_max=2.0 * np.pi,
tolerance=self._tolerance,
**integrator_kwargs
)
assert (
poloidal_trace.status == 0
), f"Poloidal trace did not complete a full 2*pi trace before terminating"
poloidal_angle = np.linspace(
0.0, 2.0 * np.pi, num=self.curve_size, endpoint=False
)
return (
poloidal_angle,
poloidal_trace.sol(poloidal_angle),
poloidal_trace.sol(2.0 * np.pi)[2],
)
