"""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 autogrid_method
from torx.decorators import autodoc_class
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
[docs]
@autodoc_class
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.
"""
from torx.specializations.flare import import_flare
self._model, self._bfield, self._analysis = import_flare()
# 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 _read_phi(self, **kwargs):
"""Read, scale and convert phi to radians for FLARE."""
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
[docs]
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]
[docs]
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]
[docs]
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"]},
)
[docs]
@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)
[docs]
@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)
[docs]
@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)
[docs]
def get_boundary_polygon(self, **kwargs):
"""
Return the boundary polygon.
Todo
----
Not yet implemented.
"""
raise NotImplementedError
[docs]
def normalized_flux_surface_label(self, r_norm, z_norm, **kwargs):
"""
Return the normalized flux surface label.
Todo
----
Not yet implemented.
"""
raise NotImplementedError
[docs]
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."
)
