"""Contains s-alpha equilibrium."""
import numpy as np
import xarray as xr
from numba import njit, cfunc
from numbalsoda import lsoda_sig
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
[docs]
@autodoc_class
class SalphaEquilibrium(EquilibriumBaseClass):
"""
Represents a Salpha equilibrium as defined in PARALLAX.
Similar to CircularToroidalEquilibrium but mostly used for legacy access
to such simulations. Please consider using CircularToroidalEquilibrium
instead.
"""
[docs]
def __init__(
self,
rhomin: float=0.1,
rhomax: float=0.7,
minor_r: float=0.365,
b_ref: float=1.0,
l_ref: float=1.0,
q_ref: float=0.854,
shear: float=2.184,
theta_limiter: float=3.1415,
dtheta_limiter: float=0.4,
rho_limiter: float=0.3235,
B0: Quantity=Quantity("1T"),
L_ref: float=1.0,
):
"""
Initialize the SalphaEquilibrium.
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).
The L_ref parameter sets the length normalization in meters. By
definition, the magnetic axis is located at R = L_ref meters.
"""
self.B0 = B0
self._norm_length = Quantity(L_ref, "m")
# Box-fraction
boxfac = 1.25
self.x0 = 1.0
self.y0 = 0.0
self.minor_r = minor_r
self.rhomin = rhomin
self.rhomax = rhomax
self.xmin = 1.0 - minor_r * boxfac * rhomax
self.xmax = 1.0 + minor_r * boxfac * rhomax
self.ymin = -minor_r * boxfac * rhomax
self.ymax = minor_r * boxfac * rhomax
self.L_ref = l_ref
self.B_ref = b_ref
self.q_ref = q_ref
self.shear = shear
self.theta_limiter = theta_limiter
self.dtheta_limiter = dtheta_limiter
self.rho_limiter = rho_limiter
# compute position of limiter plates
self.thlim1 = theta_limiter + dtheta_limiter
self.thlim1 = np.mod(self.thlim1, 2.0 * np.pi)
self.thlim2 = theta_limiter - dtheta_limiter
self.thlim2 = np.mod(self.thlim2, 2.0 * np.pi)
# if no limiter set, set rho_limiter large
if dtheta_limiter <= 0.0:
self.rho_limiter = 10.0 * rhomax
assert(dtheta_limiter < np.pi), \
"Error: dtheta_limiter cannot be larger than pi"
@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 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 normalized_flux_surface_label(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""Return the normalized flux surface label (rho_pol)."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
if kwargs["is_structured"]:
r_mesh, z_mesh = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
z_mesh = z_norm
rho = type(self)._rho_base(
np.asarray(r_mesh), np.asarray(z_mesh), self.minor_r
)
return xr.DataArray(rho,
attrs={"norm":Quantity(1.0)},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_r(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return the radial magnetic field component."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
if kwargs["is_structured"]:
r_mesh, z_mesh = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
z_mesh = z_norm
Br = type(self)._br_base(np.asarray(r_mesh), np.asarray(z_mesh),
self.q_ref, self.shear, self.minor_r)
return xr.DataArray(Br,
attrs={"norm":self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_z(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return the vertical magnetic field component."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
if kwargs["is_structured"]:
r_mesh, z_mesh = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
z_mesh = z_norm
Bz = type(self)._bz_base(np.asarray(r_mesh), np.asarray(z_mesh),
self.q_ref, self.shear, self.minor_r)
return xr.DataArray(Bz,
attrs={"norm":self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_toroidal(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""Return the toroidal magnetic field component."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
if kwargs["is_structured"]:
r_mesh, _ = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
btor = type(self)._bt_base(np.asarray(r_mesh))
return xr.DataArray(btor,
attrs={"norm":self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def jacobian(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return the Jacobian of the s-alpha geometry."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
if kwargs["is_structured"]:
r_mesh, _ = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
jac = type(self)._jac_base(np.asarray(r_mesh))
return xr.DataArray(jac,
attrs={"norm":self.R0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def theta(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return the geometric poloidal angle in the range [0, 2*pi]."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
if kwargs["is_structured"]:
r_mesh, z_mesh = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
z_mesh = z_norm
thval = type(self)._theta_base(np.asarray(r_mesh), np.asarray(z_mesh))
return xr.DataArray(thval,
attrs={"norm":Quantity(1.0)},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def qval(self, r_norm, z_norm, **kwargs):
"""Return the safety factor."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
if kwargs["is_structured"]:
r_mesh, z_mesh = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
z_mesh = z_norm
qval = type(self)._qval_base(np.asarray(r_mesh), np.asarray(z_mesh),
self.q_ref, self.shear, self.minor_r)
return xr.DataArray(qval,
attrs={"norm":Quantity(1.0)},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
def theta_centered(self, theta, limiter_edge):
"""Return the signed poloidal distance to limiter edge."""
theta_centered = np.mod(theta - limiter_edge, 2.0 * np.pi)
if theta_centered.ndim > 0:
LL = np.where(theta_centered > np.pi - self.dtheta_limiter)
theta_centered[LL] -= 2.0 * np.pi
elif theta_centered > np.pi - self.dtheta_limiter:
theta_centered -= 2.0 * np.pi
return theta_centered
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."""
rho = np.sqrt((r - 1.0)**2 + z**2) / self.minor_r
theta = np.arctan2(z, r - 1.0) % (2.0 * np.pi)
q = self.q_ref + self.shear * rho**2
b_r = -(rho * self.minor_r / r) * np.sin(theta) / q
b_z = (rho * self.minor_r / r) * np.cos(theta) / q
return b_r, b_z, 1.0 / r, r
[docs]
def get_boundary_polygon(self, angular_resolution=100, rhomax_factor=1.1):
"""
Return a polygon which approximates the position of the limiter.
Can be used for plotting and to control fieldline tracing.
"""
from torx.geometry import Polygon2D
theta_points = np.linspace(0, 2 * np.pi, num=angular_resolution)
rho_points = np.ones_like(theta_points) * self.rhomax * rhomax_factor
in_limiter = np.logical_and(
self.theta_centered(theta_points, self.thlim1) < 0.0,
self.theta_centered(theta_points, self.thlim2) > 0.0,
)
rho_points[in_limiter] = min(self.rho_limiter, \
self.rhomax * rhomax_factor)
x_points = rho_points * np.cos(theta_points)
y_points = rho_points * np.sin(theta_points)
return Polygon2D(x_points=x_points, y_points=y_points)
@staticmethod
@njit
def _rho_base(r, z, minor_r):
return np.sqrt((r - 1.0)**2 + z**2) / minor_r
@staticmethod
@njit
def _theta_base(r, z):
th = np.arctan2(z, r - 1.0)
return np.mod(th, 2.0 * np.pi)
@staticmethod
@njit
def _qval_base(r, z, qref, shear, minor_r):
rho = np.sqrt((r - 1.0)**2 + z**2) / minor_r
return qref + shear * rho**2
@staticmethod
@njit
def _br_base(r, z, qref, shear, minor_r):
rho = np.sqrt((r - 1.0)**2 + z**2) / minor_r
theta = np.mod(np.arctan2(z, r - 1.0), 2.0 * np.pi)
qval = qref + shear * rho**2
return -(rho * minor_r / r) * np.sin(theta) / qval
@staticmethod
@njit
def _bz_base(r, z, qref, shear, minor_r):
rho = np.sqrt((r - 1.0)**2 + z**2) / minor_r
theta = np.mod(np.arctan2(z, r - 1.0), 2.0 * np.pi)
qval = qref + shear * rho**2
return (rho * minor_r / r) * np.cos(theta) / qval
@staticmethod
@njit
def _bt_base(r):
return 1.0 / r
@staticmethod
@njit
def _jac_base(r):
return r
[docs]
def make_trace_eq(self):
"""Return cfunc address for numbalsoda field-line tracing."""
minor_r = self.minor_r
qref = self.q_ref
shear = self.shear
@cfunc(lsoda_sig)
def trace_eq(t, y, du, p): # pragma: no cover
r_norm = y[0]
z_norm = y[1]
rho = np.sqrt((r_norm - 1.0)**2 + z_norm**2) / minor_r
theta = np.mod(np.arctan2(z_norm, r_norm - 1.0), 2.0 * np.pi)
qval = qref + shear * rho**2
factor = rho * minor_r * r_norm / qval
du[0] = -factor * np.sin(theta)
du[1] = factor * np.cos(theta)
du[2] = np.sqrt(du[0] * du[0] + du[1] * du[1] + r_norm * r_norm)
return trace_eq.address
