"""Contains circular toroidal equilibrium."""
import numpy as np
import xarray as xr
from numba import cfunc
from numbalsoda import lsoda_sig
from .equilibrium_m import EquilibriumBaseClass
from torx.units import Quantity
from torx.units import check_units
from torx.arrays import make_xarray
from torx.decorators import autogrid_method
from torx.decorators import autodoc_class
[docs]
@autodoc_class
class CircularToroidalEquilibrium(EquilibriumBaseClass):
"""
Represents a circular toroidal equilibrium as defined in PARALLAX.
Has concentric circles as flux surfaces and is located at R=1, the
toroidal field is defined as 1/R.
"""
[docs]
def __init__(
self,
rhomin: float=0.0,
rhomax: float=0.3,
q_0: float=1.0,
q_quad_param: float=12.5,
hel_amp: float=0.0,
hel_m: int=3,
hel_n: int=2,
hel_rho: float=0.2,
hel_sigma: float=0.01,
B0: Quantity=Quantity("1T"),
R0: Quantity=Quantity("1m"),
kiss_boundary_on: bool=False,
):
"""
Initialize the CircularToroidalEquilibrium.
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.
"""
self.B0 = B0
self.R0 = R0
self.x0 = 1.0
self.y0 = 0.0
self.rhomin = rhomin
self.rhomax = rhomax
self.xmin = 0.0 + np.finfo(float).eps
self.xmax = 2.0 - np.finfo(float).eps
self.ymin = -1.0 + np.finfo(float).eps
self.ymax = 1.0 - np.finfo(float).eps
self.q_0 = q_0
self.q_quad_param = q_quad_param
self.hel_amp = hel_amp
self.hel_m = hel_m
self.hel_n = hel_n
self.hel_rho = hel_rho
self.hel_sigma = hel_sigma
@property
def axis_r_norm(self):
"""Magnetic axis radial position in normalized units."""
return make_xarray(self.x0, attrs={"norm":self.R0})
@property
def axis_z_norm(self):
"""Magnetic axis vertical position in normalized units."""
return make_xarray(self.y0, attrs={"norm":self.R0})
@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).
It is calculated only for the unperturbed case, so no effects of
the magnetic islands by the perturbation are taken into account.
"""
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 = np.sqrt((r_mesh - 1.0)**2 + z_mesh**2)
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, phi, **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)
rho = self.normalized_flux_surface_label(r_norm, z_norm)
theta = self.theta(r_norm, z_norm)
qval = self.qval(rho)
Br_zero = -rho * np.sin(theta) / (r_norm * qval * np.sqrt(1.0 - rho**2))
Br_perturb = self.magfield_component_perturbation_r(r_norm, z_norm, phi)
Br = Br_zero + Br_perturb
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, phi, **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)
rho = self.normalized_flux_surface_label(r_norm, z_norm)
theta = self.theta(r_norm, z_norm)
qval = self.qval(rho)
Bz_zero = rho * np.cos(theta) / (r_norm * qval * np.sqrt(1.0 - rho**2))
Bz_perturb = self.magfield_component_perturbation_z(r_norm, z_norm, phi)
Bz = Bz_zero + Bz_perturb
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
return xr.DataArray(1.0 / r_mesh,
attrs={"norm":self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_poloidal(self, r_norm, z_norm, phi, **kwargs) \
-> xr.DataArray:
"""Return the Jacobian of the circular geometry."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
B_r = self.magfield_component_r(r_norm, z_norm, phi)
B_z = self.magfield_component_z(r_norm, z_norm, phi)
return xr.DataArray(np.sqrt(B_r**2 + B_z**2),
attrs={"norm":self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_absolute(self, r_norm, z_norm, phi, **kwargs) -> xr.DataArray:
"""Return the absolute value of the magnetic field."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
B_x = self.magfield_component_r(r_norm, z_norm, phi, **kwargs)
B_y = self.magfield_component_z(r_norm, z_norm, phi, **kwargs)
B_t = self.magfield_component_toroidal(r_norm, z_norm, **kwargs)
return np.sqrt(B_x**2 + B_y**2 + B_t**2)
[docs]
@autogrid_method
def magfield_component_perturbation_r(self, r_norm, z_norm, phi, **kwargs) \
-> xr.DataArray:
"""Return the radial magfield component of the helical perturbation."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
rho = self.normalized_flux_surface_label(r_norm, z_norm)
theta = self.theta(r_norm, z_norm)
hel_amp = self.hel_amp
hel_n = self.hel_n
hel_m = self.hel_m
eta = self.eta(rho, theta)
deta_drho = self.deta_drho(rho, theta)
deta_dtheta = self.deta_dtheta(rho, theta)
radenv = self.radenv(rho)
dradenv_drho = self.dradenv_drho(rho)
drho_dz = z_norm / rho
dtheta_dz = (r_norm - 1.0) / rho**2
deta_dz = deta_drho * drho_dz + deta_dtheta * dtheta_dz
phase = hel_m * eta - hel_n * phi
Br_perturb = hel_amp / r_norm * (
- dradenv_drho * drho_dz * np.sin(phase)
- hel_m * radenv * deta_dz * np.cos(phase)
)
return xr.DataArray(Br_perturb,
attrs={"norm":self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_perturbation_z(self, r_norm, z_norm, phi, **kwargs) \
-> xr.DataArray:
"""Return vertical magfield component of the helical perturbation."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
rho = self.normalized_flux_surface_label(r_norm, z_norm)
theta = self.theta(r_norm, z_norm)
hel_amp = self.hel_amp
hel_n = self.hel_n
hel_m = self.hel_m
eta = self.eta(rho, theta)
deta_drho = self.deta_drho(rho, theta)
deta_dtheta = self.deta_dtheta(rho, theta)
radenv = self.radenv(rho)
dradenv_drho = self.dradenv_drho(rho)
drho_dr = (r_norm - 1.0) / rho
dtheta_dr = -z_norm / rho**2
deta_dr = deta_drho * drho_dr + deta_dtheta * dtheta_dr
phase = hel_m * eta - hel_n * phi
Bz_perturb = hel_amp / r_norm * (
dradenv_drho * drho_dr * np.sin(phase)
+ hel_m * radenv * deta_dr * np.cos(phase)
)
return xr.DataArray(Bz_perturb,
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 circular toroidal 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
return xr.DataArray(r_mesh,
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
theta = np.arctan2(z_mesh, r_mesh - 1.0)
theta = np.mod(theta, 2.0 * np.pi)
return xr.DataArray(theta,
attrs={"norm":Quantity(1.0)},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
def qval(self, rho):
"""Return the safety factor for a given rho poloidal."""
return self.q_0 + self.q_quad_param * rho**2
[docs]
def eta(self, rho, theta, **kwarg) -> xr.DataArray:
"""Return the straight field line angle in the range [0, 2*pi]."""
eta = -2.0 * np.arctan((rho - 1.0) * np.tan(theta / 2.0) / \
np.sqrt(1.0 - rho**2))
eta = np.mod(eta, 2.0 * np.pi)
return eta
[docs]
def radenv(self, rho):
"""
Return the radial envelope of the perturbation to the magnetic field.
It is a Gaussian at position hel_rho and with the variance hel_sigma.
"""
hel_rho = self.hel_rho
hel_sigma = self.hel_sigma
radenv = np.exp(-(rho - hel_rho)**2 / (2.0 * hel_sigma**2))
return radenv
[docs]
def dradenv_drho(self, rho):
"""Return the derivative of the radial envelope with respect to rho."""
hel_rho = self.hel_rho
hel_sigma = self.hel_sigma
dradenv_drho = -(rho - hel_rho) / hel_sigma**2 * \
np.exp(-(rho - hel_rho)**2 / (2.0 * hel_sigma**2))
return dradenv_drho
[docs]
def deta_drho(self, rho, theta):
"""Return the derivative of eta with respect to rho."""
deta_drho = -np.sin(theta) / (np.sqrt(1.0 - rho**2) * \
(1.0 + rho * np.cos(theta)))
return deta_drho
[docs]
def deta_dtheta(self, rho, theta):
"""Return the derivative of eta with respect to theta."""
deta_dtheta = np.sqrt(1.0 - rho**2) / (1.0 + rho * np.cos(theta))
return deta_dtheta
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)
theta = np.arctan2(z, r - 1.0) % (2.0 * np.pi)
q = self.q_0 + self.q_quad_param * rho**2
denom = r * q * np.sqrt(1.0 - rho**2)
b_r = -rho * np.sin(theta) / denom
b_z = rho * np.cos(theta) / denom
if self.hel_amp != 0.0:
eta = (
-2.0 * np.arctan(
(rho - 1.0) * np.tan(theta / 2.0)
/ np.sqrt(1.0 - rho**2)
)
) % (2.0 * np.pi)
radenv = np.exp(
-(rho - self.hel_rho)**2 / (2.0 * self.hel_sigma**2)
)
dradenv_drho = (
-(rho - self.hel_rho) / self.hel_sigma**2 * radenv
)
deta_drho = (
-np.sin(theta)
/ (np.sqrt(1.0 - rho**2) * (1.0 + rho * np.cos(theta)))
)
deta_dtheta = (
np.sqrt(1.0 - rho**2) / (1.0 + rho * np.cos(theta))
)
phase = self.hel_m * eta - self.hel_n * phi
sin_phase = np.sin(phase)
cos_phase = np.cos(phase)
drho_dz = z / rho
dtheta_dz = (r - 1.0) / rho**2
deta_dz = deta_drho * drho_dz + deta_dtheta * dtheta_dz
b_r += (self.hel_amp / r) * (
-dradenv_drho * drho_dz * sin_phase
- self.hel_m * radenv * deta_dz * cos_phase
)
drho_dr = (r - 1.0) / rho
dtheta_dr = -z / rho**2
deta_dr = deta_drho * drho_dr + deta_dtheta * dtheta_dr
b_z += (self.hel_amp / r) * (
dradenv_drho * drho_dr * sin_phase
+ self.hel_m * radenv * deta_dr * cos_phase
)
return b_r, b_z, 1.0 / r, r
[docs]
def make_trace_eq(self):
"""Return cfunc address for numbalsoda field-line tracing."""
q_0 = self.q_0
q_quad_param = self.q_quad_param
hel_amp = self.hel_amp
hel_m = self.hel_m
hel_n = self.hel_n
hel_rho = self.hel_rho
hel_sigma = self.hel_sigma
@cfunc(lsoda_sig)
def trace_eq(t, y, du, p): # pragma: no cover
r = y[0]
z = y[1]
rho = np.sqrt((r - 1.0)**2 + z**2)
theta = np.mod(np.arctan2(z, r - 1.0), 2.0 * np.pi)
q = q_0 + q_quad_param * rho**2
denom = r * q * np.sqrt(1.0 - rho**2)
b_r = -rho * np.sin(theta) / denom
b_z = rho * np.cos(theta) / denom
if hel_amp != 0.0:
arg = (
(rho - 1.0) * np.tan(theta / 2.0)
/ np.sqrt(1.0 - rho**2)
)
eta = np.mod(-2.0 * np.arctan(arg), 2.0 * np.pi)
radenv = np.exp(
-(rho - hel_rho)**2 / (2.0 * hel_sigma**2)
)
dradenv_drho = -(rho - hel_rho) / hel_sigma**2 * radenv
_denom = np.sqrt(1.0 - rho**2) * (1.0 + rho * np.cos(theta))
deta_drho = -np.sin(theta) / _denom
deta_dtheta = np.sqrt(1.0 - rho**2) / (
1.0 + rho * np.cos(theta)
)
phase = hel_m * eta - hel_n * t
sin_phase = np.sin(phase)
cos_phase = np.cos(phase)
drho_dz = z / rho
dtheta_dz = (r - 1.0) / rho**2
deta_dz = deta_drho * drho_dz + deta_dtheta * dtheta_dz
b_r += (hel_amp / r) * (
-dradenv_drho * drho_dz * sin_phase
- hel_m * radenv * deta_dz * cos_phase
)
drho_dr = (r - 1.0) / rho
dtheta_dr = -z / rho**2
deta_dr = deta_drho * drho_dr + deta_dtheta * dtheta_dr
b_z += (hel_amp / r) * (
dradenv_drho * drho_dr * sin_phase
+ hel_m * radenv * deta_dr * cos_phase
)
du[0] = b_r * r * r
du[1] = b_z * r * r
du[2] = np.sqrt(du[0] * du[0] + du[1] * du[1] + r * r)
return trace_eq.address
[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
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)
