"""Contains the RawNumericalEquilibrium class for preprocessing equilibria."""
import numpy as np
from scipy.interpolate import RectBivariateSpline
from pathlib import Path
import warnings
from .numerical_m import NumericalEquilibrium
from torx.units import Quantity
from torx.arrays import make_xarray, pad_array_with_coords
from torx.geometry import Polygon2D
from torx.decorators import autodoc_class
[docs]
@autodoc_class
class RawNumericalEquilibrium(NumericalEquilibrium):
"""
A NumericalEquilibrium in an uninitialized state for preprocessing.
Allows setting raw SI data before normalization. Once normalize() is
called the equilibrium is fully initialized and can be used as a
NumericalEquilibrium. The equilibrium cannot be used for field
evaluations before normalization.
"""
[docs]
@classmethod
def create_empty(cls):
"""
Create an empty uninitialized equilibrium for preprocessing.
Returns
-------
RawNumericalEquilibrium
An empty uninitialized equilibrium.
"""
obj = cls.__new__(cls)
obj._is_normalized = False
obj._is_filtered = False
obj._is_padded = False
obj._flipped_Z = False
obj._norm_length = None
obj._norm_magfield = None
obj.poloidal_field_factor = +1.0
obj.axis_r = None
obj.axis_z = None
obj.x_point_r = None
obj.x_point_z = None
obj._axis_Btor = None
obj._axis_Btor_units = None
obj._spline_basis_r = None
obj._spline_basis_z = None
obj._psi_axis = None
obj._psi_separatrix = None
obj._psi_data = None
obj.psi_interpolator = None
obj._b_pol_interpolator = None
obj._rho_min = None
obj._rho_max = None
obj.divertor_polygon = None
obj.exclusion_polygon = None
obj.flux_limit_polygons = {}
obj.flux_limit_rho_min = {}
obj.flux_limit_rho_max = {}
obj.raw = None
return obj
[docs]
@classmethod
def from_dict(cls, equi_dict: dict):
"""
Create a RawNumericalEquilibrium from a standardized dictionary.
Parameters
----------
equi_dict : dict
Dictionary with standardized keys:
- axis_r: float, magnetic axis R in meters
- axis_z: float, magnetic axis Z in meters
- axis_Btor: float, toroidal field on axis in Tesla
- axis_Btor_units: str, units of axis_Btor
- psi_axis: float, poloidal flux on axis in Weber
- psi_separatrix: float, poloidal flux on separatrix in Weber
- spline_basis_r: np.ndarray, R grid vector in meters
- spline_basis_z: np.ndarray, Z grid vector in meters
- psi_data: np.ndarray, poloidal flux on grid in Weber
- x_point_r: float, optional, x-point R in meters
- x_point_z: float, optional, x-point Z in meters
Returns
-------
RawNumericalEquilibrium
"""
obj = cls.create_empty()
axis_r_units = equi_dict.get("axis_r_units", "m")
x_point_units = equi_dict.get("x_point_units", axis_r_units)
obj.poloidal_field_factor = float(equi_dict.get("poloidal_field_factor", 1.0))
obj._norm_length = Quantity(equi_dict["axis_r"], axis_r_units)
obj._axis_Btor_units = equi_dict.get("axis_Btor_units", "T")
obj.axis_r = make_xarray(equi_dict["axis_r"], norm=Quantity(axis_r_units))
obj.axis_z = make_xarray(equi_dict["axis_z"], norm=Quantity(axis_r_units))
obj.x_point_r = (
make_xarray(equi_dict["x_point_r"], norm=Quantity(x_point_units))
if equi_dict.get("x_point_r") is not None else None
)
obj.x_point_z = (
make_xarray(equi_dict["x_point_z"], norm=Quantity(x_point_units))
if equi_dict.get("x_point_z") is not None else None
)
obj._axis_Btor = float(equi_dict["axis_Btor"])
obj.psi_axis = equi_dict["psi_axis"]
obj.psi_separatrix = equi_dict["psi_separatrix"]
obj.spline_basis_r = equi_dict["spline_basis_r"]
obj.spline_basis_z = equi_dict["spline_basis_z"]
obj.psi_data = equi_dict["psi_data"]
obj._rho_min = float(equi_dict["rho_min"]) \
if equi_dict.get("rho_min") is not None else None
obj._rho_max = float(equi_dict["rho_max"]) \
if equi_dict.get("rho_max") is not None else None
obj.divertor_polygon = equi_dict.get("divertor_polygon")
obj.exclusion_polygon = equi_dict.get("exclusion_polygon")
obj.flux_limit_polygons = equi_dict.get("flux_limit_polygons", {})
obj.flux_limit_rho_min = equi_dict.get("flux_limit_rho_min", {})
obj.flux_limit_rho_max = equi_dict.get("flux_limit_rho_max", {})
obj.raw = equi_dict.get("raw")
return obj
[docs]
@classmethod
def initialize_from_equi_file(cls, equi_file: Path):
"""
Create a RawNumericalEquilibrium from a NetCDF equilibrium file.
Parameters
----------
equi_file : Path
Path to the equilibrium NetCDF file.
Returns
-------
RawNumericalEquilibrium
"""
from .io import read_netcdf_equilibrium
assert equi_file.exists() and equi_file.suffix == ".nc"
obj = cls(read_netcdf_equilibrium(equi_file), equi_file)
obj._is_normalized = True
obj._is_filtered = True
obj._is_padded = True
return obj
[docs]
def write_to_equi_file(
self,
file_path: Path,
description: str,
comment: str = "",
allow_overwrite: bool = False,
ordering: str = "Fortran",
):
"""
Write the equilibrium to a NetCDF file.
The equilibrium must be normalized before writing. Call
normalize(B0) first.
Parameters
----------
file_path : Path
Output path, must have .nc suffix.
description : str
Human-readable description written to the file header.
comment : str, optional
Additional comment written to the file header.
allow_overwrite : bool, optional
If False, raises FileExistsError if file_path exists.
ordering : str, optional
Array ordering, "Fortran" or "C". Default is "Fortran".
"""
self._check_initialized()
super().write_to_equi_file(
file_path=file_path,
description=description,
comment=comment,
allow_overwrite=allow_overwrite,
ordering=ordering,
)
[docs]
def to_numerical(self) -> NumericalEquilibrium:
"""
Convert to a NumericalEquilibrium instance.
The equilibrium must be normalized before conversion.
"""
from .numerical_m import NumericalEquilibrium
self._check_initialized()
equi_dict = {
"axis_r": float(np.asarray(self.axis_r)),
"axis_r_units": str(self.R0.units),
"axis_z": float(np.asarray(self.axis_z)),
"axis_Btor": self._axis_Btor,
"axis_Btor_units": self._axis_Btor_units,
"x_point_r": float(np.asarray(self.x_point_r)) if self.x_point_r is not None else None,
"x_point_z": float(np.asarray(self.x_point_z)) if self.x_point_z is not None else None,
"x_point_units": str(self.R0.units),
"psi_axis": self._psi_axis,
"psi_separatrix": self._psi_separatrix,
"rho_min": self._rho_min,
"rho_max": self._rho_max,
"spline_basis_r": np.asarray(self.spline_basis_r),
"spline_basis_z": np.asarray(self.spline_basis_z),
"psi_data": self._psi_data,
"poloidal_field_factor": self.poloidal_field_factor,
"divertor_polygon": self.divertor_polygon,
"exclusion_polygon": self.exclusion_polygon,
"flux_limit_polygons": self.flux_limit_polygons,
"flux_limit_rho_min": self.flux_limit_rho_min,
"flux_limit_rho_max": self.flux_limit_rho_max,
}
return NumericalEquilibrium(equi_dict, self.filepath if hasattr(self, "filepath") else Path("."))
def _check_initialized(self):
"""Check that the equilibrium is normalized before use."""
assert self._is_normalized, (
"RawNumericalEquilibrium must be normalized before use. "
"Call normalize() first."
)
if not self._is_filtered:
warnings.warn(
"RawNumericalEquilibrium has not been filtered. "
"Consider calling apply_psi_filter() before use.",
UserWarning,
stacklevel=2,
)
if not self._is_padded:
warnings.warn(
"RawNumericalEquilibrium has not been padded."
"Consider calling pad_psi() before use.",
UserWarning,
stacklevel=2,
)
@property
def _not_initialized(self):
"""
Return if the equilibrium has not been initialized.
This means that neither functions normalize, filter and pad have been
called. Useful to determine if the equilibrium is still "raw" and not
"heavily" processed in any means.
"""
return (not self._is_normalized) \
and (not self._is_filtered) \
and (not self._is_padded)
[docs]
def add_flux_limit_polygon(
self,
name: str,
polygon: Polygon2D,
rho_min: float = None,
rho_max: float = None,
):
"""
Add a flux limiting polygon with optional rho limits.
Exactly one of rho_min or rho_max must be supplied.
Parameters
----------
name : str
Name of the polygon.
polygon : Polygon2D
The polygon to add.
rho_min : float, optional
Minimum flux surface label inside the polygon.
rho_max : float, optional
Maximum flux surface label inside the polygon.
"""
assert not (rho_min is None and rho_max is None), \
"Must supply at least one flux limit per polygon"
assert rho_min is None or rho_max is None, \
"Should only apply one flux limit per polygon"
self.flux_limit_polygons[name] = polygon
self.flux_limit_rho_min[name] = rho_min
self.flux_limit_rho_max[name] = rho_max
[docs]
def pad_psi(
self,
left: int = 0,
right: int = 0,
bottom: int = 0,
top: int = 0,
) -> None:
"""
Pad the psi grid and its coordinate vectors.
Parameters
----------
left, right, bottom, top:
Number of points to add on each side.
"""
if self._is_padded:
warnings.warn(
"Cannot pad an already padded equilibrium - skipping.",
UserWarning,
stacklevel=2,
)
return
self.spline_basis_r, self.spline_basis_z, self.psi_data = pad_array_with_coords(
x_coords=self.spline_basis_r,
y_coords=self.spline_basis_z,
array=self.psi_data,
left=left,
right=right,
bottom=bottom,
top=top,
)
self._is_padded = True
[docs]
def upsample_psi(self, n_samples: int | tuple):
"""
Upsample the psi grid to a higher resolution.
Interpolates the existing psi profile to a higher resolution
grid. Can help to localize numerical noise.
Parameters
----------
n_samples : int or tuple of (int, int)
Number of sample points. If an int, the same number of points
is used in both R and Z directions. If a tuple, the first
element is the number of R points and the second is the number
of Z points.
"""
if isinstance(n_samples, tuple):
n_r, n_z = n_samples
else:
n_r = n_samples
n_z = n_samples
# Skip resampling if grid already matches requested size
current_n_r = len(self.spline_basis_r)
current_n_z = len(self.spline_basis_z)
if n_r == current_n_r and n_z == current_n_z:
return
r_samples = np.linspace(
float(np.min(self.spline_basis_r)),
float(np.max(self.spline_basis_r)),
num=n_r,
)
z_samples = np.linspace(
float(np.min(self.spline_basis_z)),
float(np.max(self.spline_basis_z)),
num=n_z,
)
new_psi = self.psi_interpolator(r_samples, z_samples, grid=True).T
self.spline_basis_r = r_samples
self.spline_basis_z = z_samples
self.psi_data = new_psi
[docs]
def apply_psi_filter(self, filter_func: callable, **filter_kwargs):
"""
Apply a filter function to psi_data.
Parameters
----------
filter_func : callable
Filter function with signature ``filter_func(psi_data, **kwargs) -> np.ndarray``.
**filter_kwargs
Keyword arguments passed directly to ``filter_func``.
Examples
--------
>>> equi.apply_psi_filter(fourier_gaussian, filter_strength=10.0, edge_pad=100)
>>> equi.apply_psi_filter(gaussian_filter, sigma=2.0)
"""
if self._is_filtered:
warnings.warn(
"Cannot filter an already filtered equilibrium - skipping.",
UserWarning,
stacklevel=2,
)
return
self._psi_data_unfiltered = self._psi_data.copy()
self.psi_interpolator_unfiltered = RectBivariateSpline(
self.spline_basis_r, self.spline_basis_z, self._psi_data_unfiltered.T
)
self.psi_data = filter_func(self.psi_data, **filter_kwargs)
self._is_filtered = True
[docs]
def normalize(self, R0: Quantity = None, B0: Quantity = None):
"""
Normalize the equilibrium from physical to normalized units.
Divides all length quantities by R0 and sets the magnetic field
normalization to B0. If not provided, R0 defaults to the magnetic
axis R position and B0 to the on-axis toroidal field.
Parameters
----------
R0 : Quantity, optional
Normalization length. Defaults to magnetic axis R position.
If a plain float is given, meters are assumed.
B0 : Quantity, optional
Normalization magnetic field. Defaults to on-axis toroidal field.
If a plain float is given, Tesla are assumed.
"""
if self._is_normalized:
warnings.warn(
"Cannot normalize an already normalized equilibrium - skipping.",
UserWarning,
stacklevel=2,
)
return
if R0 is None:
R0 = Quantity(self.axis_r.values, self.axis_r.norm)
elif not isinstance(R0, Quantity):
R0 = Quantity(R0, "m")
if B0 is None:
B0 = Quantity(self._axis_Btor, self._axis_Btor_units or "T")
elif not isinstance(B0, Quantity):
B0 = Quantity(B0, "T")
self.R0 = R0
self.B0 = B0
self._norm_length = R0
self._norm_magfield = B0
R0_val = R0.m
# NOTE: No need to scale axis or x point coords since those
# have auto normalized properties stored anyway. Only
# spline basis and polygons are required.
self.spline_basis_r = self.spline_basis_r / R0_val
self.spline_basis_z = self.spline_basis_z / R0_val
if self.divertor_polygon is not None:
self.divertor_polygon.x_points = self.divertor_polygon.x_points / R0_val
self.divertor_polygon.y_points = self.divertor_polygon.y_points / R0_val
if self.exclusion_polygon is not None:
self.exclusion_polygon.x_points = self.exclusion_polygon.x_points / R0_val
self.exclusion_polygon.y_points = self.exclusion_polygon.y_points / R0_val
for polygon in self.flux_limit_polygons.values():
polygon.x_points = polygon.x_points / R0_val
polygon.y_points = polygon.y_points / R0_val
self.psi_interpolator = RectBivariateSpline(
self.spline_basis_r, self.spline_basis_z, self._psi_data.T
)
self._b_pol_interpolator = self._build_b_pol_interpolator()
self._is_normalized = True
@property
def poloidal_magnetic_field_interpolator(self) -> RectBivariateSpline:
"""Interpolator of the poloidal magnetic field magnitude."""
return self._b_pol_interpolator
@NumericalEquilibrium.psi_axis.setter
def psi_axis(self, value):
"""Poloidal flux at magnetic axis."""
self._psi_axis = None if value is None else float(value)
@NumericalEquilibrium.psi_separatrix.setter
def psi_separatrix(self, value):
"""Poloidal flux at separatrix."""
self._psi_separatrix = None if value is None else float(value)
@property
def rho_min(self):
"""Minimum flux surface label."""
if self._is_normalized:
return super().rho_min
return self._rho_min
@rho_min.setter
def rho_min(self, value):
self._rho_min = None if value is None else float(value)
@property
def rho_max(self):
"""Maximum flux surface label."""
if self._is_normalized:
return super().rho_max
return self._rho_max
@rho_max.setter
def rho_max(self, value):
self._rho_max = None if value is None else float(value)
@property
def axis_Btor(self):
"""On-axis toroidal magnetic field."""
if self._is_normalized:
return super().axis_Btor
return self._axis_Btor
@axis_Btor.setter
def axis_Btor(self, value):
"""Set the on-axis toroidal field as a plain float in Tesla."""
self._axis_Btor = None if value is None else float(value)
@property
def axis_r_norm(self):
"""Magnetic axis radial position in normalized units."""
self._check_initialized()
return super().axis_r_norm
@property
def axis_z_norm(self):
"""Magnetic axis vertical position in normalized units."""
self._check_initialized()
return super().axis_z_norm
@property
def psi_data(self):
"""Raw poloidal flux data on the grid in Weber."""
return self._psi_data
@psi_data.setter
def psi_data(self, value: np.ndarray):
if self._spline_basis_r is None or self._spline_basis_z is None:
raise RuntimeError(
"Cannot set psi_data before spline_basis_r and "
"spline_basis_z are set."
)
self._psi_data = np.asarray(value, dtype=float)
self.psi_interpolator = RectBivariateSpline(
self.spline_basis_r, self.spline_basis_z, self._psi_data.T
)
self._b_pol_interpolator = self._build_b_pol_interpolator()
@property
def R0(self) -> Quantity:
"""Length normalization."""
return self._norm_length
@R0.setter
def R0(self, value: Quantity):
self._norm_length = value
@property
def B0(self) -> Quantity:
"""Magnetic field normalization."""
return self._norm_magfield
@B0.setter
def B0(self, value: Quantity):
self._norm_magfield = value
@property
def psi_units(self):
"""Units of the poloidal flux."""
return "Wb"
@property
def spline_basis_r(self):
"""R grid vector in meters (before normalization) or normalized units."""
return self._spline_basis_r
@spline_basis_r.setter
def spline_basis_r(self, value: np.ndarray):
value = np.asarray(value, dtype=float)
assert np.all(np.diff(value) > 0), \
"spline_basis_r must be strictly increasing"
self._spline_basis_r = value
@property
def spline_basis_z(self):
"""Z grid vector in meters (before normalization) or normalized units."""
return self._spline_basis_z
@spline_basis_z.setter
def spline_basis_z(self, value: np.ndarray):
value = np.asarray(value, dtype=float)
assert np.all(np.diff(value) > 0), \
"spline_basis_z must be strictly increasing"
self._spline_basis_z = value
