"""Polygon sets forming Poincare structures."""
import numpy as np
import xarray as xr
from scipy.interpolate import LinearNDInterpolator, CloughTocher2DInterpolator
from torx.decorators import AutoDocClass
from .ordered_polygon_2d_m import OrderedPolygon2D
[docs]
class PoincarePolygonSet(AutoDocClass):
"""
Set of ordered polygon objects constructed from Poincare data.
Allows computing a radial coordinate based on the polygon areas.
"""
[docs]
def __init__(self, poincare_data: xr.DataArray,
phi_array: xr.DataArray, start_idx: int=0, k: int=20,
n_interp: int=0, resample_method: str="spline",
make_simple: bool=False, sort: bool=True,
is_separatrix: int=None):
"""
Initialize the set from Poincare data from PARALLAX.
The polygons are used to define flux surfaces and compute
a radial coordinate based on their average areas.
Parameters
----------
poincare_data : xarray.DataArray
4D array with dimensions (dim_surf, dim_turn, dim_plane, dim_xy)
containing the Poincare points from PARALLAX.
phi_array : xarray.DataArray
Array of toroidal angles corresponding to each plane.
start_idx : int, optional
Index of the first point to include.
k : int, optional
Number of nearest neighbors for polygon construction.
n_interp : int, optional
Number of points to interpolate along polygon edges.
resample_method : {'spline', 'pchip'}, optional
Polygon resampling scheme (see OrderedPolygon2D.resample).
'pchip' is recommended for scrape-off layer surfaces, as it
adds no self-intersections at thin fingers. Neither smooths;
use smooth_surfaces after construction.
make_simple : bool, optional
Remove self-intersection loops from each polygon (outer
scrape-off layer fingers doubling back).
sort : bool, optional
Sort surfaces by increasing mean area. Disable to preserve
the input order (for example when surface 0 is the magnetic
axis and is_separatrix refers to a fixed input index).
is_separatrix : int, optional
Surface index (into the current, possibly sorted, ordering)
whose mean area normalizes rho to 1. Defaults to the
outermost surface.
"""
assert(type(poincare_data) == xr.DataArray), \
"Poincare data must be Xarray DataArray."
required_dims = {"dim_surf", "dim_turn", "dim_plane", "dim_xy"}
assert required_dims.issubset(poincare_data.dims), \
f"Expected dims {required_dims}, got {poincare_data.dims}"
self.n_surfaces = poincare_data.sizes["dim_surf"]
self.n_turns = poincare_data.sizes["dim_turn"]
self.n_planes = poincare_data.sizes["dim_plane"]
assert(poincare_data.sizes["dim_xy"] == 2), \
"Poincare data must have dim_xy of size 2."
# TODO: Store the vertices as a numeric DataArray with dims
# (dim_surf, dim_plane, dim_vertex, dim_xy) instead of an object
# array of OrderedPolygon2D. Area, centroid, mean area, rho
# normalization, orientation and smoothing would then vectorize,
# and the set could be written to netCDF/zarr and chunked with
# dask. This needs a uniform vertex count, so require n_interp > 0
# when make_simple is True (or NaN-pad with a mask). Ordering,
# self-intersection splitting and resampling stay per-polygon
# loops. Keep a polygon(s, p) accessor for existing callers.
# Create OrderedPolygon2D objects from Poincare data
polygon_data = xr.apply_ufunc(
OrderedPolygon2D.from_2d_array,
poincare_data,
start_idx, k, n_interp, resample_method,
make_simple,
input_core_dims=[['dim_turn', 'dim_xy'],
[], [], [], [], []],
output_core_dims=[[]],
vectorize=True,
output_dtypes=[object]
)
mean_area = self._mean_area(polygon_data)
# Sort everything according to mean area
if sort:
sort_indices = mean_area.argsort()
polygon_data = polygon_data.isel(dim_surf=sort_indices)
poincare_data = poincare_data.isel(dim_surf=sort_indices)
mean_area = mean_area.isel(dim_surf=sort_indices)
self.polygon_data = polygon_data
self.poincare_data = poincare_data
self.mean_area = mean_area
self._is_separatrix = is_separatrix
self.phi_array = phi_array
# Compute normalized radial coordinate rho
self._normalize()
self._rho_interpolators = None
@staticmethod
def _mean_area(polygon_data: xr.DataArray) -> xr.DataArray:
"""Return the plane-averaged signed area of each polygon."""
polygon_area = xr.apply_ufunc(
lambda p: p.signed_area(),
polygon_data,
input_core_dims=[[]],
output_core_dims=[[]],
vectorize=True,
output_dtypes=[float]
)
return polygon_area.mean(dim='dim_plane')
def _normalize(self):
"""Set rho_data as the mean area normalized at the separatrix."""
if self._is_separatrix is None:
norm = self.mean_area.isel(dim_surf=-1)
else:
norm = self.mean_area.isel(dim_surf=self._is_separatrix)
self.rho_data = self.mean_area / norm
def smooth_surfaces(self, smooth: float, rho_threshold: float=1.0,
surfaces=None):
"""
Spline-smooth the ragged outer (scrape-off layer) surfaces.
Replaces each selected surface's polygon in every plane by a
smoothing periodic spline with the same number of points (see
OrderedPolygon2D.interpolate), then recomputes the mean-area rho
label. This corrects spiky field-line-traced surfaces while
preserving their area. Inner/closed surfaces are left untouched.
Parameters
----------
smooth : float
Dimensionless smoothing strength: the allowed RMS distance of
the spline from the points as a fraction of the polygon's
equivalent radius sqrt(|A| / pi). Typically 0.005 to 0.05;
a value of 0 or less does nothing.
rho_threshold : float, optional
Smooth every surface with rho greater than this (default 1.0,
i.e. the scrape-off layer only). Ignored when ``surfaces`` is
given.
surfaces : sequence of int, optional
Explicit surface indices to smooth, overriding rho_threshold.
Returns
-------
PoincarePolygonSet
``self``, for chaining.
"""
if smooth <= 0.0:
return self
if surfaces is None:
rho = self.rho_data.values
surfaces = [s for s in range(self.n_surfaces)
if rho[s] > rho_threshold]
for s in surfaces:
for ip in range(self.n_planes):
polygon = self.polygon_data.isel(
dim_surf=s, dim_plane=ip).item()
polygon.resample(len(polygon.x_points), "spline", smooth)
self.mean_area = self._mean_area(self.polygon_data)
self._normalize()
self.reset_interpolator()
return self
[docs]
def __getitem__(self, key: int | tuple) -> object:
"""
Get item(s) from the polygon set using indices.
[i_plane, i_surf, i_point, i_coord].
"""
# Normalize key to tuple and check that it has 1-4 elements
if not isinstance(key, tuple):
key = (key,)
if not 1 <= len(key) <= 4:
raise IndexError(
"Indexing requires 1-4 indices: "
"[i_plane, i_surf, i_point, i_coord]."
)
# Normalize the key to a 4-tuple (pad with slices as needed)
# Think of slice(None) as ":", it works as a "pass-through" operator
# The alternative would be to do many else if statements
key = key + (slice(None),) * (4 - len(key))
i_plane, i_surf, i_point, i_coord = key
sliced = self.polygon_data.isel(dim_plane=i_plane, dim_surf=i_surf)
# If fully indexed return a scalar
if all(isinstance(i, (int, np.integer)) for i in key):
return sliced.item()[i_point, i_coord]
# Determine output core dimensions based on key
output_core_dims = []
# If i_point is a slice, the first output dimension is 'dim_point'
if isinstance(i_point, slice):
output_core_dims.append('dim_turn')
# If i_coord is a slice, the second output dimension is 'dim_coord'
if isinstance(i_coord, slice):
output_core_dims.append('dim_xy')
def _extract_polygon_slice(poly: OrderedPolygon2D, i_point, i_coord):
if i_coord == slice(None):
return poly[i_point]
else:
return poly[i_point, i_coord]
# when there are slices, use apply_ufunc to extract the data
return xr.apply_ufunc(
_extract_polygon_slice,
sliced,
i_point,
i_coord,
input_core_dims=[[], [], []],
output_core_dims=[output_core_dims],
vectorize=True,
output_dtypes=[np.float64]
)
[docs]
@classmethod
def from_file(cls, filename:str, start_idx: int=0, k: int=20,
n_interp=0, resample_method: str="spline",
make_simple: bool=False, sort: bool=True,
is_separatrix: int=None):
"""Initialize from a file produced with PARALLAX diagnose_poincare."""
poincare = xr.open_dataset(filename)
return cls(poincare.poincare_data, poincare.phi_array,
start_idx, k, n_interp, resample_method, make_simple,
sort, is_separatrix)
[docs]
@classmethod
def from_torx_poincare(cls, starting_points, n_turns):
"""Initialize from torx field line tracing."""
# Placeholder - implement field line tracing
raise NotImplementedError("Field line tracing not implemented yet.")
@property
def rho_interp(self):
"""
Lazy-loaded interpolator for rho(R, Z, phi).
Returns an xarray of interpolators for each plane.
"""
if self._rho_interpolators is None:
# Create interpolators for each plane
def _plane_interpolator(poincare_data, phi, rho_data):
R_data = poincare_data[:,:,0].flatten()
Z_data = poincare_data[:,:,1].flatten()
rho_data = np.broadcast_to(rho_data[:, np.newaxis],
poincare_data[:,:,0].shape).flatten()
coords = np.column_stack([R_data, Z_data])
return CloughTocher2DInterpolator(coords, rho_data)
self._rho_interpolators = xr.apply_ufunc(
_plane_interpolator,
self.poincare_data,
self.phi_array,
self.rho_data,
# Separate along dim_plane
input_core_dims=[["dim_surf", "dim_turn", "dim_xy"],
[], ["dim_surf"]],
output_core_dims=[[]],
vectorize=True,
output_dtypes=[object],
)
return self._rho_interpolators
def rho(self, R, Z, phi):
"""Return interpolated rho(R, Z, phi)."""
# Find closest phi
phi_closest = abs(self.phi_array-phi).argmin(dim='dim_plane')
# Get interpolator for that plane
interp = self.rho_interp.isel(dim_plane=phi_closest).item()
# Return interpolated value
return interp(R, Z)
def reset_interpolator(self):
"""Reset the interpolator."""
self._rho_interpolators = None
def to_lineout_set(self, rho_i):
"""Convert the polygon set with radial index rho_i to a LineoutSet."""
from torx.analysis.lineouts import Lineout, LineoutSet
lineout_x = self[:,rho_i,:,0].T
lineout_y = self[:,rho_i,:,1].T
num_planes = lineout_x.sizes['dim_plane']
lineout_objects = np.empty(num_planes, dtype=object)
for i in range(num_planes):
x_data = lineout_x.isel(dim_plane=i).values
y_data = lineout_y.isel(dim_plane=i).values
lineout = Lineout(x_data, y_data)
lineout_objects[i] = lineout
ds_lineouts = xr.DataArray(
lineout_objects,
dims="phi",
)
return LineoutSet(ds_lineouts)
