"""Contains Dommaschk equilibrium."""
from pathlib import Path
import numpy as np
import xarray as xr
from numba import njit, cfunc
from numbalsoda import lsoda_sig
from ._dommaschk_initialization_m import _init_CD_CN, _init_Imn
from ._dommaschk_fitting_coefficients_m import DommaschkFittingCoefficients
from .equilibrium_m import EquilibriumBaseClass
from torx.units import Quantity
from torx.arrays import make_xarray
from torx import genex_domm_resources_dir
from torx.units import check_units
from torx.geometry import Polygon3D
from torx.decorators import autogrid_method
from torx.decorators import autodoc_class
@njit(cache=True)
def _eval_dommaschk_fields(
r, z, phi,
m_arr, fc,
I_mn_c, I_mn_p,
dIdZ_c, dIdZ_p,
CD_r_c, C_r_p, CD_ln_c, C_ln_p,
CN_r_c, CN_ln_c,
dCDdR_r_c, dCdR_r_p,
dCDdR_ln_c, dCdR_ln_p,
dCNdR_r_c, dCNdR_ln_c,
B_norm):
"""Compute (b_r, b_z, b_t) for a Dommaschk equilibrium."""
logR = np.log(r)
cos_m = np.cos(m_arr * phi)
sin_m = np.sin(m_arr * phi)
I_mn = I_mn_c * z ** I_mn_p
dIdZ = dIdZ_c * z ** dIdZ_p
CD_mk = (np.sum(CD_r_c * r ** C_r_p, axis=2)
+ np.sum(CD_ln_c * r ** C_ln_p * logR, axis=2))
CN_mk = (np.sum(CN_r_c * r ** C_r_p, axis=2)
+ np.sum(CN_ln_c * r ** C_ln_p * logR, axis=2))
dCDdR_mk = (np.sum(dCDdR_r_c * r ** dCdR_r_p, axis=2)
+ np.sum(dCDdR_ln_c * r ** dCdR_ln_p * logR, axis=2))
dCNdR_mk = (np.sum(dCNdR_r_c * r ** dCdR_r_p, axis=2)
+ np.sum(dCNdR_ln_c * r ** dCdR_ln_p * logR, axis=2))
D_ml = np.sum(I_mn[:, 1:, :] * CD_mk[:, np.newaxis, :], axis=2)
N_ml1 = np.sum(I_mn[:, :-1, :] * CN_mk[:, np.newaxis, :], axis=2)
dD_dR = np.sum(I_mn[:, 1:, :] * dCDdR_mk[:, np.newaxis, :], axis=2)
dN_dR = np.sum(I_mn[:, :-1, :] * dCNdR_mk[:, np.newaxis, :], axis=2)
dD_dZ = np.sum(dIdZ[:, 1:, :] * CD_mk[:, np.newaxis, :], axis=2)
dN_dZ = np.sum(dIdZ[:, :-1, :] * CN_mk[:, np.newaxis, :], axis=2)
cos_ml = cos_m[:, np.newaxis]
sin_ml = sin_m[:, np.newaxis]
b_r = np.sum(
(fc[0] * cos_ml + fc[1] * sin_ml) * dD_dR
+ (fc[2] * cos_ml + fc[3] * sin_ml) * dN_dR
) / B_norm
b_z = np.sum(
(fc[0] * cos_ml + fc[1] * sin_ml) * dD_dZ
+ (fc[2] * cos_ml + fc[3] * sin_ml) * dN_dZ
) / B_norm
dcos_dphi = -m_arr[:, np.newaxis] * sin_ml
dsin_dphi = m_arr[:, np.newaxis] * cos_ml
b_t = (1.0 + np.sum(
(fc[0] * dcos_dphi + fc[1] * dsin_dphi) * D_ml
+ (fc[2] * dcos_dphi + fc[3] * dsin_dphi) * N_ml1
)) / (r * B_norm)
return b_r, b_z, b_t
[docs]
@autodoc_class
class DommaschkEquilibrium(EquilibriumBaseClass):
"""
Dommaschk type stellarator equilibrium.
Analytic representation of the vacuum magnetic field of stellarator
configurations, as described in [1]_.
References
----------
.. [1] W. Dommaschk, "Representations for vacuum potentials in
stellarators", Computer Physics Communications 40, pg. 203 (1986).
"""
[docs]
def __init__(
self,
filepath: str="",
x0: float=1.0,
y0: float=0.0,
phi0: float=0.0,
num_field_periods: int=5,
l_pol: int=4,
m_tor_consecutive: int=1,
fitting_coef= None,
bndry_from_file: bool=False,
outer_bndry_file: str="",
outer_bndry_var: str="",
inner_bndry_file: str="",
inner_bndry_var: str="",
xmin: float= 0.80,
xmax: float= 1.20,
ymin: float=-0.20,
ymax: float= 0.20,
B0: Quantity=Quantity("1T"),
L_ref: float=1.0,
**kwargs
):
"""
Initialize the DommaschkEquilibrium.
Arguments to the initializer match the parameters supplied in PARALLAX.
You can either manually specify the arguments,
or pass a **kwargs_dictionary generated from a parameter file.
The B0 parameter sets the value of the toroidal magnetic field at the
normalized axis location (x0, y0, phi0).
The L_ref parameter sets the length normalization in meters. Note that
in 3D, the magnetic axis might not always be located at
R = L_ref meters.
"""
self.filepath = filepath
self.x0 = x0
self.y0 = y0
self.phi0 = phi0
self.num_field_periods = num_field_periods
self.l_pol = l_pol
self.m_tor_consecutive = m_tor_consecutive
self.bndry_from_file = bndry_from_file
self.outer_bndry_file = outer_bndry_file
self.outer_bndry_var = outer_bndry_var
self.inner_bndry_file = inner_bndry_file
self.inner_bndry_var = inner_bndry_var
self.xmin = xmin
self.xmax = xmax
self.ymin = ymin
self.ymax = ymax
self.B0 = B0
self._norm_length = Quantity(L_ref, "m")
self.m = xr.DataArray(np.arange(self.m_tor_consecutive + 1) \
* self.num_field_periods, dims=("m"))
if type(fitting_coef) == type(None):
self.fitting_coef = DommaschkFittingCoefficients(l_pol,
m_tor_consecutive)
elif type(fitting_coef) == DommaschkFittingCoefficients:
assert fitting_coef.l_pol == l_pol
assert fitting_coef.m_tor_consecutive == m_tor_consecutive
self.fitting_coef = fitting_coef
else:
self.fitting_coef = DommaschkFittingCoefficients(l_pol,
m_tor_consecutive,
fitting_coef)
_init_CD_CN(self)
_init_Imn(self)
self._init_B_norm()
self._init_numpy_cache()
self._init_boundary()
def _init_numpy_cache(self):
"""Extract coefficients to numpy for fast scalar ODE evaluation."""
self._np_m = self.m.values
self._np_fc = self.fitting_coef._parsed_coef.values
self._np_Imn_c = self.I_mn_coef.values
self._np_Imn_p = self.I_mn_power.values
self._np_dIdZ_c = self.dI_dZ_coef.values
self._np_dIdZ_p = self.dI_dZ_power.values
self._np_CD_r_c = self.CD_r_coef.values
self._np_CN_r_c = self.CN_r_coef.values
self._np_C_r_p = self.C_r_power.values
self._np_CD_lnr_c = self.CD_lnr_coef.values
self._np_CN_lnr_c = self.CN_lnr_coef.values
self._np_C_lnr_p = self.C_lnr_power.values
self._np_dCDdR_r_c = self.dCD_dR_r_coef.values
self._np_dCNdR_r_c = self.dCN_dR_r_coef.values
self._np_dCdR_r_p = self.dC_dR_r_power.values
self._np_dCDdR_ln_c = self.dCD_dR_lnr_coef.values
self._np_dCNdR_ln_c = self.dCN_dR_lnr_coef.values
self._np_dCdR_ln_p = self.dC_dR_lnr_power.values
self._np_B_norm = float(self._B_norm_internal)
def _init_B_norm(self):
"""
Determine the magnetic field normalization.
Uses the non-normalized value of btor evaluated at the location of the
magnetic axis. This is used to normalize all subsequent magnetic
field calculations.
Set the normalization factor to 1 to calculate the non-normalized value
of btor.
"""
self._B_norm_internal = 1.0
btor_unnormalized = self.magfield_component_toroidal(self.x0, self.y0,
phi=self.phi0)
self._B_norm_internal = btor_unnormalized
def _init_boundary(self):
"""
Load the inner and/or outer boundary from file, if specified.
If no boundary is given, construct one from the box limits.
"""
default_polygon_file = genex_domm_resources_dir \
/ "flux_polygons_dommaschk_default.nc"
if self.outer_bndry_file == "default":
self.outer_bndry_file = default_polygon_file
if self.inner_bndry_file == "default":
self.inner_bndry_file = default_polygon_file
if self.bndry_from_file:
outer_boundary_filepath = (
Path(self.filepath) / Path(self.outer_bndry_file)
)
self.outer_bndry = Polygon3D.initialize_from_file(
outer_boundary_filepath,
self.outer_bndry_var,
num_field_periods=self.num_field_periods)
inner_boundary_filepath = (
Path(self.filepath) / Path(self.inner_bndry_file)
)
self.inner_bndry = Polygon3D.initialize_from_file(
inner_boundary_filepath,
self.inner_bndry_var,
num_field_periods=self.num_field_periods) \
if self.inner_bndry_file != "" else None
else:
# Construct the outer boundary polygon using the box limits in a
# counterclockwise direction
vertices = np.array([
[[self.xmin, self.xmax, self.xmax, self.xmin],
[self.ymin, self.ymin, self.ymax, self.ymax]]
])
self.outer_bndry = Polygon3D(
polygon_vertices=vertices,
phi_array=np.array(0.0),
num_field_periods=self.num_field_periods)
self.inner_bndry = None
@property
def axis_r_norm(self):
"""Magnetic axis radial position in normalized units."""
return make_xarray(self.x0)
@property
def axis_z_norm(self):
"""Magnetic axis vertical position in normalized units."""
return make_xarray(self.y0)
@property
def axis_phi(self):
"""Magnetic axis toroidal position."""
return make_xarray(self.phi0)
@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 the length normalization."""
check_units(value, {"[length]":1}, "R0")
self._norm_length = value
[docs]
@autogrid_method
def magfield_component_r(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
r"""
Return the radial component of the vacuum magnetic field B.
Calculated according to $B_R = \partial V / \partial R$.
"""
if "phi" not in kwargs:
raise TypeError("Missing required keyword argument: 'phi'")
phi = kwargs["phi"]
R = self.convert_length_to_normalized(r_norm)
Z = self.convert_length_to_normalized(z_norm)
logR = np.log(R)
cos_m_phi = np.cos(self.m * phi)
sin_m_phi = np.sin(self.m * phi)
I_mn = self.I_mn_coef * Z**self.I_mn_power
dCD_mk_dR = (self.dCD_dR_r_coef \
* R**self.dC_dR_r_power).sum(dim="terms") \
+ (self.dCD_dR_lnr_coef \
* R**self.dC_dR_lnr_power * logR).sum(dim="terms")
dD_ml_dR = (I_mn.isel(l=slice(1, None)) * dCD_mk_dR).sum(dim="k")
dCN_mk_dR = (self.dCN_dR_r_coef \
* R**self.dC_dR_r_power).sum(dim="terms") \
+ (self.dCN_dR_lnr_coef \
* R**self.dC_dR_lnr_power * logR).sum(dim="terms")
dN_ml_1_dR = (I_mn.isel(l=slice(0, -1)) * dCN_mk_dR).sum(dim="k")
# Eq. 12 (differentiated with respect to R)
dV_ml_dR = ( self.fitting_coef[0, :, :] * cos_m_phi \
+ self.fitting_coef[1, :, :] * sin_m_phi \
) * dD_ml_dR \
+ ( self.fitting_coef[2, :, :] * cos_m_phi \
+ self.fitting_coef[3, :, :] * sin_m_phi \
) * dN_ml_1_dR
Br = dV_ml_dR.sum(("m", "l")) / self._B_norm_internal
return Br.assign_attrs(norm=self.B0)
[docs]
@autogrid_method
def magfield_component_z(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
r"""
Return the vertical component of the vacuum magnetic field B.
Calculated according to $B_Z = \partial V / \partial Z$.
"""
if "phi" not in kwargs:
raise TypeError("Missing required keyword argument: 'phi'")
phi = kwargs["phi"]
R = self.convert_length_to_normalized(r_norm)
Z = self.convert_length_to_normalized(z_norm)
logR = np.log(R)
cos_m_phi = np.cos(self.m * phi)
sin_m_phi = np.sin(self.m * phi)
dI_mn_dZ = self.dI_dZ_coef * Z**self.dI_dZ_power
CD_mk = (
(self.CD_r_coef * R**self.C_r_power).sum(dim="terms")
+ (self.CD_lnr_coef * R**self.C_lnr_power * logR).sum(dim="terms")
)
dD_ml_dZ = (dI_mn_dZ.isel(l=slice(1, None)) * CD_mk).sum(dim="k")
CN_mk = (
(self.CN_r_coef * R**self.C_r_power).sum(dim="terms")
+ (self.CN_lnr_coef * R**self.C_lnr_power * logR).sum(dim="terms")
)
dN_ml_1_dZ = (dI_mn_dZ.isel(l=slice(0, -1)) * CN_mk).sum(dim="k")
# Eq. 12 (differentiated with respect to Z)
dV_ml_dZ = ( self.fitting_coef[0, :, :] * cos_m_phi \
+ self.fitting_coef[1, :, :] * sin_m_phi \
) * dD_ml_dZ \
+ ( self.fitting_coef[2, :, :] * cos_m_phi \
+ self.fitting_coef[3, :, :] * sin_m_phi \
) * dN_ml_1_dZ
Bz = dV_ml_dZ.sum(("m", "l")) / self._B_norm_internal
return Bz.assign_attrs(norm=self.B0)
[docs]
@autogrid_method
def magfield_component_toroidal(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
r"""
Return the poloidal component of the vacuum magnetic field B.
Calculated according to $B_\phi = (1/R)\, \partial V / \partial \phi$.
"""
if "phi" not in kwargs:
raise TypeError("Missing required keyword argument: 'phi'")
phi = kwargs["phi"]
R = self.convert_length_to_normalized(r_norm)
Z = self.convert_length_to_normalized(z_norm)
logR = np.log(R)
dsin_m_phi_dphi = self.m * np.cos(self.m * phi)
dcos_m_phi_dphi = -self.m * np.sin(self.m * phi)
I_mn = self.I_mn_coef * Z**self.I_mn_power
CD_mk = (
(self.CD_r_coef * R**self.C_r_power).sum(dim="terms")
+ (self.CD_lnr_coef * R**self.C_lnr_power * logR).sum(dim="terms")
)
D_ml = (I_mn.isel(l=slice(1, None)) * CD_mk).sum(dim="k")
CN_mk = (
(self.CN_r_coef * R**self.C_r_power).sum(dim="terms")
+ (self.CN_lnr_coef * R**self.C_lnr_power * logR).sum(dim="terms")
)
N_ml_1 = (I_mn.isel(l=slice(0, -1)) * CN_mk).sum(dim="k")
# Eq. 12 (differentiated with respect to phi)
dV_ml_dphi = ( self.fitting_coef[0, :, :] * dcos_m_phi_dphi \
+ self.fitting_coef[1, :, :] * dsin_m_phi_dphi \
) * D_ml \
+ ( self.fitting_coef[2, :, :] * dcos_m_phi_dphi \
+ self.fitting_coef[3, :, :] * dsin_m_phi_dphi \
) * N_ml_1
# The leading 1.0 is the derivative of the phi term in Eq. 1
Btor = (
(1.0 + dV_ml_dphi.sum(("m", "l"))) / (R * self._B_norm_internal)
)
return Btor.assign_attrs(norm=self.B0)
def _magfield_scalars(self, r: float, z: float, phi: float) -> tuple:
"""Return (b_r, b_z, b_t, jacobian) as plain floats for ODE hot path."""
logR = np.log(r)
cos_m = np.cos(self._np_m * phi)
sin_m = np.sin(self._np_m * phi)
cos_me = cos_m[:, np.newaxis]
sin_me = sin_m[:, np.newaxis]
# $I_{m,n}$: (n_m, n_l_ext, n_k) - shared between $B_r$ and $B_\phi$
I_mn = self._np_Imn_c * z ** self._np_Imn_p
# $C^D_{m,k}$, $C^N_{m,k}$: (n_m, n_k) - shared between $B_z$, $B_\phi$
CD_mk = ((self._np_CD_r_c * r ** self._np_C_r_p ).sum(-1)
+ (self._np_CD_lnr_c * r ** self._np_C_lnr_p * logR).sum(-1))
CN_mk = ((self._np_CN_r_c * r ** self._np_C_r_p ).sum(-1)
+ (self._np_CN_lnr_c * r ** self._np_C_lnr_p * logR).sum(-1))
# $B_r = \partial V / \partial R$
dCDdR_mk = (
(self._np_dCDdR_r_c * r ** self._np_dCdR_r_p).sum(-1)
+ (self._np_dCDdR_ln_c * r ** self._np_dCdR_ln_p * logR).sum(-1)
)
dCNdR_mk = (
(self._np_dCNdR_r_c * r ** self._np_dCdR_r_p).sum(-1)
+ (self._np_dCNdR_ln_c * r ** self._np_dCdR_ln_p * logR).sum(-1)
)
dD_dR = (I_mn[:, 1:, :] * dCDdR_mk[:, np.newaxis, :]).sum(-1)
dN_dR = (I_mn[:, :-1,:] * dCNdR_mk[:, np.newaxis, :]).sum(-1)
fc = self._np_fc
b_r = (
(fc[0] * cos_me + fc[1] * sin_me) * dD_dR
+ (fc[2] * cos_me + fc[3] * sin_me) * dN_dR
).sum() / self._np_B_norm
# $B_z = \partial V / \partial Z$
dIdZ = self._np_dIdZ_c * z ** self._np_dIdZ_p
dD_dZ = (dIdZ[:, 1:, :] * CD_mk[:, np.newaxis, :]).sum(-1)
dN_dZ = (dIdZ[:, :-1,:] * CN_mk[:, np.newaxis, :]).sum(-1)
b_z = (
(fc[0] * cos_me + fc[1] * sin_me) * dD_dZ
+ (fc[2] * cos_me + fc[3] * sin_me) * dN_dZ
).sum() / self._np_B_norm
# $B_\phi = (1/R)\, \partial V / \partial \phi$
dcos_dphi = -self._np_m * sin_m
dsin_dphi = self._np_m * cos_m
D_ml = (I_mn[:, 1:, :] * CD_mk[:, np.newaxis, :]).sum(-1)
N_ml1 = (I_mn[:, :-1,:] * CN_mk[:, np.newaxis, :]).sum(-1)
dphi_cos = dcos_dphi[:, np.newaxis]
dphi_sin = dsin_dphi[:, np.newaxis]
b_t = (
1.0 + (
(fc[0] * dphi_cos + fc[1] * dphi_sin) * D_ml
+ (fc[2] * dphi_cos + fc[3] * dphi_sin) * N_ml1
).sum()
) / (r * self._np_B_norm)
return float(b_r), float(b_z), float(b_t), r
[docs]
def get_boundary_polygon(self, **kwargs):
"""Return the boundary polygon."""
return {"outer":self.outer_bndry, "inner":self.inner_bndry}
[docs]
def normalized_flux_surface_label(self, r_norm, z_norm, **kwargs):
"""Return the normalized flux surface label."""
raise NotImplementedError
[docs]
def make_trace_eq(self):
"""Return cfunc address for numbalsoda field-line tracing."""
m_arr = self._np_m
fc = self._np_fc
I_mn_c = self._np_Imn_c
I_mn_p = self._np_Imn_p
dIdZ_c = self._np_dIdZ_c
dIdZ_p = self._np_dIdZ_p
CD_r_c = self._np_CD_r_c
C_r_p = self._np_C_r_p
CD_ln_c = self._np_CD_lnr_c
C_ln_p = self._np_C_lnr_p
CN_r_c = self._np_CN_r_c
CN_ln_c = self._np_CN_lnr_c
dCDdR_r_c = self._np_dCDdR_r_c
dCdR_r_p = self._np_dCdR_r_p
dCDdR_ln_c = self._np_dCDdR_ln_c
dCdR_ln_p = self._np_dCdR_ln_p
dCNdR_r_c = self._np_dCNdR_r_c
dCNdR_ln_c = self._np_dCNdR_ln_c
B_norm = self._np_B_norm
@cfunc(lsoda_sig)
def trace_eq(phi, y, du, p): # pragma: no cover
"""Define trace equation for fast field-line tracer."""
r = y[0]
z = y[1]
b_r, b_z, b_t = _eval_dommaschk_fields(
r, z, phi,
m_arr, fc,
I_mn_c, I_mn_p,
dIdZ_c, dIdZ_p,
CD_r_c, C_r_p, CD_ln_c, C_ln_p,
CN_r_c, CN_ln_c,
dCDdR_r_c, dCdR_r_p,
dCDdR_ln_c, dCdR_ln_p,
dCNdR_r_c, dCNdR_ln_c,
B_norm
)
du[0] = b_r / b_t * r
du[1] = b_z / b_t * r
du[2] = np.sqrt(du[0] * du[0] + du[1] * du[1] + r * r)
return trace_eq.address
