"""Contains FLARE equilibrium."""
from pathlib import Path
import numpy as np
import xarray as xr
from typing import Sequence
from ._flare_coilset_loader_m import load_coilsets
from .check_equilibrium_m import FlareEquiType
from .equilibrium_m import EquilibriumBaseClass
from torx.units import Quantity
from torx.arrays import make_xarray
from torx.units import check_units
from torx.decorators import autodoc_function, autogrid_method
from torx.grid import Grid2D
from torx.fileio import filepath_resolver
def _add_phi_doc(func):
"""Append phi documentation to a function's docstring."""
phi_doc = """
Other Parameters
----------------
phi_index : int or float, optional
Plane index. Converted to a toroidal angle for FLARE in radians
based on the grid resolution and whether only the first field period
is used.
phi : float, optional
Physical toroidal angle in radians. If only the first field
period is simulated, the angle is mapped into
[0, 2*pi/num_field_periods).
Exactly one of phi_index or phi must be provided.
"""
if func.__doc__:
func.__doc__ += phi_doc
return func
_active_flare_equilibrium = None
@autodoc_function
def free_active_flare_equilibrium():
"""
Free the shared FLARE model held by the active FlareEquilibrium.
Helper function which is safe to call even if no FlareEquilibrium is
currently loaded. Checks if there is an active FLARE equilibrium loaded in
memory and if so, frees the model.
"""
if _active_flare_equilibrium is not None:
_active_flare_equilibrium.free()
[docs]
class FlareEquilibrium(EquilibriumBaseClass):
"""
Class for equilibria generated by the FLARE tool [2]_.
The FLARE code is a magnetic mesh generator that is integrated within a
suite of tools for the analysis of the magnetic geometry in toroidal
fusion devices [1]_.
References
----------
.. [1] H. Frerichs 2024 Nucl. Fusion 64 106034,
https://doi.org/10.1088/1741-4326/ad7303
.. [2] https://hfrerichs.gitlab.io/flare
"""
[docs]
def __init__(
self,
filepath: Path=Path("./"),
rhomin: float=0.0,
rhomax: float=1.0,
B0: Quantity=Quantity("1T"),
L_ref: float=1.0,
equi_flare_type: str="",
boundary_flare_type: str="",
rho_flare_type: str="",
flare_rho_dirpath: str="",
flare_rho_prefix: str="",
rho_symms: int=1,
flare_equi3d_path_mgrid: Path=Path(""),
flare_equi3d_path_axis: Path=Path(""),
flare_equi3d_path_bfield: Path=Path(""),
flare_equi3d_path_hint: Path=Path(""),
Rmin: float=0.0,
Rmax: float=1.0,
Zmin: float=-1.0,
Zmax: float=1.0,
n_points_phi: int=1,
only_first_field_period: bool=False,
amplitudes: list=[15000, 15000, 15000, 15000, 15000],
bspline_data: str="magnetic_field",
spline_order: int=3,
num_field_periods: int=5,
group: int=-1,
bmax: float=5.0,
bnd_paths: list=[],
**kwargs
):
"""
Initialize the FlareEquilibrium.
Compatible with mgrid, coilset and HINT FLARE equilibrium types.
"""
global _active_flare_equilibrium
if (
_active_flare_equilibrium is not None
and not _active_flare_equilibrium._freed
):
raise RuntimeError(
"Another FlareEquilibrium is already loaded from the "
"shared FLARE model. Free up the loaded model before "
"allocating a new one."
)
from torx.specializations.flare import import_flare
self._model, self._bfield, self._analysis = import_flare()
self._freed = False
_active_flare_equilibrium = self
# Legacy correction
if equi_flare_type == "equi3d":
equi_flare_type = "equi3d_mgrid"
filepath = Path(filepath)
paths = {
"axis": flare_equi3d_path_axis,
"bfield": flare_equi3d_path_bfield,
"hint": flare_equi3d_path_hint,
"mgrid": flare_equi3d_path_mgrid,
}
paths = {
key: p if (p := Path(path)).is_absolute() else filepath / p
for key, path in paths.items()
}
self.__equi_name__ = "FlareEquilibrium"
self.filepath = filepath
self.rhomin = rhomin
self.rhomax = rhomax
self.B0 = B0
self.equi_flare_type = equi_flare_type
self.boundary_flare_type = boundary_flare_type
self.rho_flare_type = rho_flare_type
self.flare_rho_dirpath = flare_rho_dirpath
self.flare_rho_prefix = flare_rho_prefix
self.rho_symms = rho_symms
self._norm_length = Quantity(L_ref, "m")
self.flare_equi3d_path_mgrid = paths["mgrid"]
self.flare_equi3d_path_axis = paths["axis"]
self.flare_equi3d_path_bfield = paths["bfield"]
self.flare_equi3d_path_hint = paths["hint"]
self.Rmin = Rmin
self.Rmax = Rmax
self.Zmin = Zmin
self.Zmax = Zmax
self.n_points_phi = n_points_phi
self.only_first_field_period = only_first_field_period
self.num_field_periods = num_field_periods
self.amplitudes = amplitudes
self.bspline_data = bspline_data
self.spline_order = spline_order
self.group = group
self.bmax = bmax
self.bnd_paths = bnd_paths
if self.boundary_flare_type == "kisslinger":
boundary_path = Path(self.bnd_paths[0]).parents[0]
if not boundary_path.is_absolute():
boundary_path = self.filepath / boundary_path
boundary = self._model.boundary_config(str(boundary_path))
if FlareEquiType(equi_flare_type) == FlareEquiType.equi3d_mgrid:
flare_model = self._model.Equi3d_Mgrid(str(paths["mgrid"]),
amplitudes,
bspline_data,
spline_order
)
elif FlareEquiType(equi_flare_type) == FlareEquiType.equi3d_coilset:
flare_model = load_coilsets(paths["bfield"])
elif FlareEquiType(equi_flare_type) == FlareEquiType.equi3d_hint:
flare_model = self._model.Equi3d_Hint(str(paths["hint"]),
group, bmax)
else:
raise NotImplementedError(f"{equi_flare_type} not implemented.")
if self.boundary_flare_type == "kisslinger":
self._model.init_equi3d_parallax(flare_model,
str(paths["axis"]),
boundary=boundary)
else:
self._model.init_equi3d_parallax(flare_model,
str(paths["axis"]))
self._norm_length = Quantity(self.axis_r_flare, "m")
self.Z0 = Quantity(self.axis_z_flare, "m")
self._norm_magfield = Quantity(self.Btor0, "T")
def free(self):
"""
Free the underlying FLARE model memory.
Checks if there is an existing model first, so is safe to call multiple
times. The FLARE model is a single global resource, so this must not be
called while another `FlareEquilibrium` built from the same model
is still in use.
"""
global _active_flare_equilibrium
if not self._freed:
self._model.free()
self._freed = True
if _active_flare_equilibrium is self:
_active_flare_equilibrium = None
def __enter__(self):
"""Return self for use as a context manager."""
return self
def __exit__(self, exc_type, exc_value, traceback):
"""Free the underlying FLARE model memory on context exit."""
self.free()
def _read_phi(self, **kwargs):
"""Read, scale and convert phi to radians for FLARE."""
if self._freed:
raise RuntimeError(
"This FlareEquilibrium's FLARE model has already been "
"freed and can no longer be used."
)
phi_index = kwargs.get("phi_index", None)
phi = kwargs.get("phi", None)
if (phi_index is None) == (phi is None):
raise TypeError(
"Exactly one of 'phi_index' or 'phi' must be provided."
)
if phi_index is not None:
if self.only_first_field_period:
phi = 2 * np.pi / (self.num_field_periods * self.n_points_phi) \
* phi_index
else:
phi = 2 * np.pi / self.n_points_phi * phi_index
elif phi is not None:
if self.only_first_field_period:
phi = phi % (2 * np.pi / self.num_field_periods)
else:
phi = phi % (2 * np.pi)
return phi
def _denormalize_flare(self, r_norm, z_norm):
"""Denormalize lengths for FLARE functions."""
r_values = r_norm * self.axis_r_flare
z_values = z_norm * self.axis_r_flare + self.axis_z_flare
return r_values, z_values
def _normalize_flare(self, r_values, z_values):
"""Normalize lengths from FLARE function inputs."""
r_norm = r_values / self.axis_r_flare
z_norm = (z_values - self.axis_z_flare) / self.axis_r_flare
return r_norm, z_norm
def axis_r(self, **kwargs):
"""Magnetic axis radial position in non-normalized units."""
phi = self._read_phi(**kwargs)
axis = self._bfield.f2py.analysis.magnetic_axis(phi)
return axis[0]
def axis_z(self, **kwargs):
"""Magnetic axis vertical position in non-normalized units."""
phi = self._read_phi(**kwargs)
axis = self._bfield.f2py.analysis.magnetic_axis(phi)
return axis[1]
def Btor_axis(self, **kwargs):
"""Return the toroidal magnetic field strength on axis."""
phi = self._read_phi(**kwargs)
return self._bfield.eval(self.axis_r_flare, 0.0, phi)[2]
@property
def axis_r_flare(self):
"""Unnormalized magnetic axis radial position on first plane."""
return self.axis_r(phi_index = 0)
@property
def axis_z_flare(self):
"""Unnormalized magnetic axis vertical position on first plane."""
return self.axis_z(phi_index = 0)
@property
def Btor0(self):
"""Toroidal magnetic field strength on axis for first plane."""
return self.Btor_axis(phi_index = 0)
@property
def axis_r_norm(self):
"""Magnetic axis radial position in normalized units."""
return make_xarray(1.0)
@property
def axis_z_norm(self):
"""Magnetic axis vertical position in normalized units."""
return make_xarray(0.0)
@property
def axis_phi(self):
"""Magnetic axis toroidal position."""
return make_xarray(0.0)
@property
def B0(self) -> Quantity:
"""Magnetic field normalization."""
return self._norm_magfield
@B0.setter
def B0(self, value: Quantity):
"""Set the magnetic field normalization."""
check_units(value, {"[mass]":1, "[time]":-2, "[current]":-1}, "B0")
self._norm_magfield = value
@property
def R0(self) -> Quantity:
"""Length normalization."""
return self._norm_length
@R0.setter
def R0(self, value: Quantity):
"""Set length normalization."""
check_units(value, {"[length]":1}, "R0")
self._norm_length = value
@autogrid_method
def _magfield_all_components_flare(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""
Evaluate the magnetic field using FLARE on the given points.
Returns all three magnetic field components from FLARE.
Compatible with single values, structured grids and unstructured grids.
"""
phi = self._read_phi(**kwargs)
if (isinstance(r_norm, xr.DataArray) and "points" in r_norm.dims) or kwargs.get("is_unstructured"):
return self._magfield_unstructured(r_norm, z_norm, phi=phi)
else:
return self._magfield_structured(r_norm, z_norm, phi=phi)
@autogrid_method
def _magfield_unstructured(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""B components evaluated for unstructured arrays from FLARE."""
phi = kwargs["phi"]
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
R_size = r_norm.size
R_flare, Z_flare = self._denormalize_flare(r_norm, z_norm)
phi_vec = np.full(R_size, phi)
B_flare = self._bfield.eval(R_flare, Z_flare, phi_vec) / self.Btor0
B_flare = np.asarray(B_flare)
B_flare = Quantity(B_flare, "T")
return xr.DataArray(
B_flare,
attrs={"norm": self.B0},
dims=("points", "component"),
coords={
"component": ["Br", "Bz", "Bphi"]
},
)
@autogrid_method
def _magfield_structured(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""B components evaluated for structured arrays or values from FLARE."""
phi = kwargs["phi"]
if isinstance(r_norm, xr.DataArray):
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
r_norm = np.atleast_1d(r_norm)
z_norm = np.atleast_1d(z_norm)
Nr, Nz = r_norm.size, z_norm.size
R_vals, Z_vals = self._denormalize_flare(r_norm, z_norm)
R_flare = R_vals[:, None]
Z_flare = Z_vals[None, :]
R_flare, Z_flare = np.broadcast_arrays(R_flare, Z_flare)
phi_grid = np.full(R_flare.shape, phi)
B_flare = self._bfield.eval(R_flare.ravel(),
Z_flare.ravel(),
phi_grid.ravel()
) \
/ self.Btor0
B_flare = np.asarray(B_flare)
if B_flare.ndim == 1:
B_flare = B_flare[None, :]
B_flare = B_flare.reshape(Nr, Nz, B_flare.shape[-1]).transpose(1,0,2)
B_flare = Quantity(B_flare, "T")
return xr.DataArray(
B_flare,
attrs={"norm": self.B0},
dims=("Z", "R", "component"),
coords={"component": ["Br", "Bz", "Bphi"]},
)
@autogrid_method
@_add_phi_doc
def magfield_component_r(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""
Return the radial component of the FLARE magnetic field B.
Parameters
----------
r_norm : xr.DataArray or float
Normalized radial coordinates. Can be a float, structured data or
unstructured data.
z_norm : xr.DataArray or float
Normalized vertical coordinates. Can be a float, structured data or
unstructured data.
"""
all_components = self._magfield_all_components_flare(
r_norm, z_norm, **kwargs)
component = all_components.sel(component="Br").drop_vars("component")
return component.squeeze(drop=True)
@autogrid_method
@_add_phi_doc
def magfield_component_z(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""
Return the vertical component of the FLARE magnetic field B.
Parameters
----------
r_norm : xr.DataArray or float
Normalized radial coordinates. Can be a float, structured data or
unstructured data.
z_norm : xr.DataArray or float
Normalized vertical coordinates. Can be a float, structured data or
unstructured data.
"""
all_components = self._magfield_all_components_flare(
r_norm, z_norm, **kwargs)
component = all_components.sel(component="Bz").drop_vars("component")
return component.squeeze(drop=True)
@autogrid_method
@_add_phi_doc
def magfield_component_toroidal(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""
Return the toroidal component of the FLARE magnetic field B.
Parameters
----------
r_norm : xr.DataArray or float
Normalized radial coordinates. Can be a float, structured data or
unstructured data.
z_norm : xr.DataArray or float
Normalized vertical coordinates. Can be a float, structured data or
unstructured data.
"""
all_components = self._magfield_all_components_flare(
r_norm, z_norm, **kwargs)
component = all_components.sel(component="Bphi").drop_vars("component")
return component.squeeze(drop=True)
def get_boundary_polygon(self, **kwargs):
"""
Return the boundary polygon.
Todo
----
Not yet implemented.
"""
raise NotImplementedError
def normalized_flux_surface_label(self, r_norm, z_norm, **kwargs):
"""
Return the normalized flux surface label.
Todo
----
Not yet implemented.
"""
raise NotImplementedError
def make_trace_eq(self):
"""
Not supported for FlareEquilibrium.
Raises
------
NotImplementedError
Always. Use ``torx.analysis.flare_fieldline_trace`` instead,
which calls the FLARE backend directly.
"""
raise NotImplementedError(
"FastTracer is not supported for FlareEquilibrium. "
"Use torx.analysis.flare_fieldline_trace to trace field lines "
"via the FLARE backend."
)
