"""Contains numerical tokamak equilibrium."""
import dask
import warnings
import numpy as np
from pathlib import Path
import xarray as xr
from xarray import Dataset
from scipy.interpolate import RectBivariateSpline, interp1d
from collections import defaultdict
from numba import njit, cfunc
from numbalsoda import lsoda_sig, lsoda
from .equilibrium_m import EquilibriumBaseClass
from torx.analysis import bspline_basis_and_deriv1
from torx.units import Quantity
from torx.units import check_units
from torx.arrays import make_xarray
from torx.units import convert_xarray_to_quantity
from torx.decorators import autogrid_method
from torx.geometry import Polygon2D
from torx.decorators import autodoc_class
from torx.analysis import MagneticFieldTracer
@njit(cache=True)
def _eval_psi_derivs(r, z, tx, ty, c, nx, ny, kx, ky):
"""
Evaluate dpsi/dr and dpsi/dz at (r, z) from bivariate B-spline tck.
c is 1-D, C-order: c[i * n_cy + j] for basis pair (i, j).
"""
n_cy = ny - ky - 1
Br, dBr, ir = bspline_basis_and_deriv1(r, kx, tx)
Bz, dBz, iz = bspline_basis_and_deriv1(z, ky, ty)
psi_dr = 0.0
psi_dz = 0.0
for a in range(kx + 1):
ci = ir - kx + a
row = ci * n_cy
for b in range(ky + 1):
cj = iz - ky + b
coef = c[row + cj]
psi_dr += coef * dBr[a] * Bz[b]
psi_dz += coef * Br[a] * dBz[b]
return psi_dr, psi_dz
[docs]
@autodoc_class
class NumericalEquilibrium(EquilibriumBaseClass):
"""Equilibrium with a divertor defined by numerical poloidal flux."""
[docs]
@classmethod
def initialize_from_params(cls, filepath: Path, params: defaultdict):
"""Create an equilibrium from a parameter file."""
path_to_netcdf = Path(params["equi_numerical_params"]["path_to_netcdf"])
equilibrium_case = params["equi_numerical_params"]["equilibrium_case"]
equi_file = Path(str(filepath / path_to_netcdf / equilibrium_case)
+ ".nc")
assert (
equi_file.exists()
), f"Error: equilibrium file {equi_file.absolute()} not found"
return cls.initialize_from_equi_file(equi_file)
[docs]
@classmethod
def initialize_from_equi_file(cls, equi_file: Path):
"""Create an equilibrium from an equilibrium file."""
from .io import read_netcdf_equilibrium
assert equi_file.exists() and equi_file.suffix == ".nc"
return cls(read_netcdf_equilibrium(equi_file), equi_file)
[docs]
def write_to_equi_file(
self,
file_path: Path,
description: str,
comment: str = "",
allow_overwrite: bool = False,
ordering: str = "Fortran",
):
"""
Write the equilibrium to a NetCDF file.
Parameters
----------
file_path : Path
Output path, must have .nc suffix.
description : str
Human-readable description written to the file header.
comment : str, optional
Additional comment written to the file header.
allow_overwrite : bool, optional
If False, raises FileExistsError if file_path exists.
ordering : str, optional
Array ordering, "Fortran" or "C". Default is "Fortran".
"""
from .io import write_netcdf_equilibrium
write_netcdf_equilibrium(
file_path=file_path,
equi=self,
description=description,
comment=comment,
allow_overwrite=allow_overwrite,
ordering=ordering,
)
[docs]
def __init__(self, equi_dict: dict, equi_file: Path, B0=Quantity("1T")):
"""
Initialize the equilibrium from a standardized dictionary.
The dictionary is expected to contain data in normalized units for
lengths and Weber for poloidal flux, as returned by
read_netcdf_equilibrium(). The normalization length R0 is fixed to
the magnetic axis radial position. B0 is set separately after
initialization via the B0 setter.
Parameters
----------
equi_dict : dict
Dictionary as returned by read_netcdf_equilibrium() with keys:
- axis_r: float, magnetic axis R in normalized units
- axis_z: float, magnetic axis Z in normalized units
- axis_Btor: float, toroidal field on axis in Tesla
- axis_Btor_units: str, units of axis_Btor
- psi_axis: float, poloidal flux on axis in Weber
- psi_separatrix: float, poloidal flux on separatrix in Weber
- rho_min: float, minimum flux surface label
- rho_max: float, maximum flux surface label
- spline_basis_r: np.ndarray, R grid in normalized units
- spline_basis_z: np.ndarray, Z grid in normalized units
- psi_data: np.ndarray, poloidal flux on grid in Weber
- norm_length: float, R0 in meters
- poloidal_field_factor: float
- x_point_r: float or np.ndarray or None
- x_point_z: float or np.ndarray or None
- divertor_polygon: Polygon2D in normalized units
- exclusion_polygon: Polygon2D or None in normalized units
- flux_limit_polygons: dict of Polygon2D in normalized units
- flux_limit_rho_min: dict of float or None
- flux_limit_rho_max: dict of float or None
equi_file : Path
Path to the source equilibrium file.
B0 : Quantity, optional
Magnetic field normalization. Default is 1 T.
"""
self.B0 = B0
self.filepath = equi_file
self.poloidal_field_factor = float(equi_dict.get("poloidal_field_factor", 1.0))
self._flipped_Z = False
axis_r_units = equi_dict["axis_r_units"]
self.axis_r = make_xarray(equi_dict["axis_r"], norm=Quantity(axis_r_units))
self.axis_z = make_xarray(equi_dict["axis_z"], norm=Quantity(axis_r_units))
x_point_units = equi_dict["x_point_units"]
self.x_point_r = (
make_xarray(equi_dict["x_point_r"], norm=Quantity(x_point_units))
if equi_dict.get("x_point_r") is not None else None
)
self.x_point_z = (
make_xarray(equi_dict["x_point_z"], norm=Quantity(x_point_units))
if equi_dict.get("x_point_z") is not None else None
)
self._axis_Btor = float(equi_dict["axis_Btor"])
self._axis_Btor_units = equi_dict["axis_Btor_units"]
norm_length = Quantity(equi_dict["axis_r"], axis_r_units)
self._norm_length = norm_length
self._spline_basis_r = make_xarray(
equi_dict["spline_basis_r"], norm=norm_length, dims=["R"]
)
self._spline_basis_z = make_xarray(
equi_dict["spline_basis_z"], norm=norm_length, dims=["Z"]
)
self._psi_data = equi_dict["psi_data"]
self.psi_interpolator = RectBivariateSpline(
self.spline_basis_r,
self.spline_basis_z,
self.psi_data.T,
)
self._b_pol_interpolator = None
self.divertor_polygon = equi_dict["divertor_polygon"]
self.exclusion_polygon = equi_dict.get("exclusion_polygon")
self.flux_limit_polygons = equi_dict.get("flux_limit_polygons", {})
self.flux_limit_rho_min = equi_dict.get("flux_limit_rho_min", {})
self.flux_limit_rho_max = equi_dict.get("flux_limit_rho_max", {})
self._psi_axis = float(equi_dict["psi_axis"])
self._psi_separatrix = float(equi_dict["psi_separatrix"])
self._rho_min = float(equi_dict["rho_min"])
self._rho_max = float(equi_dict["rho_max"])
[docs]
def flip_Z(self, preserve_helicity: bool = True):
"""
Flip the vertical direction of the equilibrium.
Affects toroidal field reversal.
"""
self._flipped_Z = not self._flipped_Z
if preserve_helicity:
self.poloidal_field_factor *= -1.0
self.axis_z *= -1.0
if self.x_point_z is not None:
self.x_point_z *= -1.0
self._spline_basis_z = -self._spline_basis_z[::-1]
self._psi_data = self._psi_data[::-1, :]
self.psi_interpolator = RectBivariateSpline(
self.spline_basis_r, self.spline_basis_z, self.psi_data.T
)
self._b_pol_interpolator = None
self.divertor_polygon.flip_Z()
if self.exclusion_polygon is not None:
self.exclusion_polygon.flip_Z()
def _build_b_pol_interpolator(self) -> RectBivariateSpline:
"""Build interpolator of the poloidal magnetic field magnitude."""
r = np.asarray(self.spline_basis_r)
z = np.asarray(self.spline_basis_z)
psi_dr = self.psi_interpolator(r, z, grid=True, dx=1).T
psi_dz = self.psi_interpolator(r, z, grid=True, dy=1).T
r_grid = r[np.newaxis, :]
b_pol = np.sqrt(psi_dr**2 + psi_dz**2) / r_grid
return RectBivariateSpline(r, z, b_pol.T)
@property
def poloidal_magnetic_field_interpolator(self) -> RectBivariateSpline:
"""Interpolator of the poloidal magnetic field magnitude."""
if not hasattr(self, '_b_pol_interpolator') or \
self._b_pol_interpolator is None:
self._b_pol_interpolator = self._build_b_pol_interpolator()
return self._b_pol_interpolator
@property
def spline_basis_r(self):
"""R grid vector."""
return self._spline_basis_r
@property
def spline_basis_z(self):
"""Z grid vector."""
return self._spline_basis_z
@property
def psi_data(self) -> xr.DataArray:
"""Raw poloidal flux data on the grid."""
return make_xarray(
self._psi_data,
norm=Quantity(self.psi_units),
)
@property
def psi_units(self):
"""Units of the poloidal flux."""
return "Wb"
@property
def axis_Btor(self) -> xr.DataArray:
"""On-axis toroidal magnetic field."""
return make_xarray(self._axis_Btor, norm=Quantity(self._axis_Btor_units))
@property
def axis_Btor_norm(self) -> xr.DataArray:
"""On-axis toroidal field normalized to B0."""
return make_xarray(self._axis_Btor / self.B0.magnitude, norm=self.B0)
@property
def rho_min(self) -> xr.DataArray:
"""Minimum flux surface label."""
return make_xarray(self._rho_min, norm=Quantity(1))
@property
def rho_max(self) -> xr.DataArray:
"""Maximum flux surface label."""
return make_xarray(self._rho_max, norm=Quantity(1))
@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(self.axis_z / self.axis_r)
@property
def x_point_r_norm(self):
"""Magnetic X-point radial position in normalized units."""
return make_xarray(self.x_point_r / self.axis_r)
@property
def x_point_z_norm(self):
"""Magnetic X-point vertical position in normalized units."""
return make_xarray(self.x_point_z / self.axis_r)
@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
@property
def psi_axis(self) -> float:
"""Poloidal flux on axis in Weber."""
return self._psi_axis
@property
def psi_separatrix(self) -> float:
"""Poloidal flux on separatrix in Weber."""
return self._psi_separatrix
[docs]
@autogrid_method
def poloidal_flux(self, r_norm, z_norm, **kwargs):
"""Return the poloidal flux at a point."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
is_structured = kwargs["is_structured"]
interp_result = self.psi_interpolator(
x=r_norm, y=z_norm, grid=is_structured
).T
if not kwargs["dims"]:
interp_result = interp_result.item()
psi = make_xarray(
interp_result,
norm=Quantity(1, self.psi_units),
dims=kwargs["dims"],
coords=kwargs["coords"]
)
return psi
[docs]
@autogrid_method
def normalized_flux_surface_label(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""Return the normalized flux surface label at a point."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
rho = self._poloidal_flux_to_rho(self.poloidal_flux(r_norm, z_norm))
return xr.DataArray(rho,
attrs={"norm": Quantity(1)},
dims=kwargs["dims"],
coords=kwargs["coords"])
def _poloidal_flux_to_rho(self, psi_value):
"""Convert poloidal flux (psi) to normalized flux coordinate (rho)."""
psi_ax = self.psi_axis
psi_sep = self.psi_separatrix
rho_squared = (psi_value - psi_ax) / (psi_sep - psi_ax)
rho_squared = np.atleast_1d(rho_squared)
# NOTE: Because the evaluation precision of the poloidal flux might be
# limited, it is not guaranteed that the flux at the axis (which
# is stored at init) is larger than the evaluated flux at axis.
# Thus we set a floor of zero for rho squared.
rho_squared[np.asarray(rho_squared < 0.0)] = 0.0
if np.size(rho_squared) == 1:
rho_squared = rho_squared[0]
return np.sqrt(rho_squared)
[docs]
@autogrid_method
def magfield_component_r(self, r_norm, z_norm, **kwargs):
"""Return the radial (R) magnetic field at a point."""
is_structured = kwargs["is_structured"]
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
# Evaluate the poloidal flux, taking the 0th derivative in x (R)
# and the 1st derivative in y(Z). Then calculate the magnetic field
# according to eq. 2.25 of [1].
# [1] H. Zohm, MHD stability of Tokamaks, p24
interp_dz = self.psi_interpolator(
x=r_norm, y=z_norm, dx=0, dy=1, grid=is_structured
).T
if not kwargs["dims"]:
interp_dz = interp_dz.item()
psi_dz = xr.DataArray(
interp_dz, dims=kwargs["dims"], coords=kwargs["coords"]
)
# NOTE: In the normalization, one factor of R0 comes from the gradient
# of psi and one from the normalized R used in the formula.
norm = self.axis_r**2
if not np.logical_and(
isinstance(r_norm, xr.DataArray),
isinstance(z_norm, xr.DataArray)
):
norm = norm.values
b_r = psi_dz / (2 * np.pi * r_norm * norm)
b_r = -self.poloidal_field_factor * b_r
return xr.DataArray(b_r / self.B0.magnitude,
attrs={"norm": self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_z(self, r_norm, z_norm, **kwargs):
"""Return the vertical (Z) magnetic field at a point."""
is_structured = kwargs["is_structured"]
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
# Evaluate the poloidal flux, taking the 1st derivative in x (R)
# and the 0th derivative in y(Z). Then calculate the magnetic field
# according to eq. 2.24 of [1].
# [1] H. Zohm, MHD stability of Tokamaks, p24
interp_dr = self.psi_interpolator(
x=r_norm, y=z_norm, dx=1, dy=0, grid=is_structured
).T
if not kwargs["dims"]:
interp_dr = interp_dr.item()
psi_dr = xr.DataArray(
interp_dr, dims=kwargs["dims"], coords=kwargs["coords"]
)
norm = self.axis_r**2
if not np.logical_and(
isinstance(r_norm, xr.DataArray),
isinstance(z_norm, xr.DataArray)
):
norm = norm.values
b_z = psi_dr / (2 * np.pi * r_norm * norm)
b_z = self.poloidal_field_factor * b_z
return xr.DataArray(b_z / self.B0.magnitude,
attrs={"norm": self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_toroidal(self, r_norm, z_norm, **kwargs):
"""Return the toroidal (Phi) magnetic field at a point."""
if kwargs["is_structured"]:
r_mesh, _ = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
r_mesh = self.convert_length_to_normalized(r_mesh)
return xr.DataArray(
self.axis_Btor_norm.item() / r_mesh,
attrs={"norm": self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"]
)
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."""
if not hasattr(self, '_denom_factor'):
self._denom_factor = (
2.0 * np.pi * self.axis_r.item() ** 2
* self.B0.magnitude
)
self._btor_norm = self.axis_Btor_norm.item()
tck = self.psi_interpolator.tck
self._tx = tck[0]
self._ty = tck[1]
self._c = tck[2]
self._nx = len(self._tx)
self._ny = len(self._ty)
self._kx = self.psi_interpolator.degrees[0]
self._ky = self.psi_interpolator.degrees[1]
psi_dr, psi_dz = _eval_psi_derivs(
r, z, self._tx, self._ty, self._c,
self._nx, self._ny, self._kx, self._ky
)
factor = self.poloidal_field_factor / (r * self._denom_factor)
b_r = -factor * psi_dz
b_z = factor * psi_dr
b_t = self._btor_norm / r
return b_r, b_z, b_t, r
[docs]
def make_trace_eq(self):
"""Return cfunc address for numbalsoda field-line tracing."""
tx = self.psi_interpolator.tck[0]
ty = self.psi_interpolator.tck[1]
c = self.psi_interpolator.tck[2]
nx = len(tx)
ny = len(ty)
kx = self.psi_interpolator.degrees[0]
ky = self.psi_interpolator.degrees[1]
denom_factor = (
2.0 * np.pi * self.axis_r.item()**2 * self.B0.magnitude
)
btor_norm = self.axis_Btor_norm.item()
pol_factor = self.poloidal_field_factor
@cfunc(lsoda_sig)
def trace_eq(t, y, du, p): # pragma: no cover
"""Define trace equation for fast field-line tracer."""
r = y[0]
z = y[1]
psi_dr, psi_dz = _eval_psi_derivs(r, z, tx, ty, c, nx, ny, kx, ky)
factor = pol_factor / (r * denom_factor)
b_r = -factor * psi_dz
b_z = factor * psi_dr
r2_bt = r * r / btor_norm
du[0] = b_r * r2_bt
du[1] = b_z * r2_bt
du[2] = np.sqrt(du[0] * du[0] + du[1] * du[1] + r * r)
return trace_eq.address
[docs]
def make_poloidal_trace_eq(self):
r"""
Return cfunc address for numbalsoda poloidal field-line integration.
The ODE integrates over theta (0 to 2*pi); y[2] accumulates the
symmetry angle so that $q = y[2] / (2\pi)$ after one full poloidal
turn.
The compiled cfunc address is cached on the instance so repeated
calls to safety_factor do not trigger recompilation.
"""
if hasattr(self, '_poloidal_trace_addr'):
return self._poloidal_trace_addr
tx = self.psi_interpolator.tck[0]
ty = self.psi_interpolator.tck[1]
c = self.psi_interpolator.tck[2]
nx = len(tx)
ny = len(ty)
kx = self.psi_interpolator.degrees[0]
ky = self.psi_interpolator.degrees[1]
denom_factor = (
2.0 * np.pi * self.axis_r.item()**2 * self.B0.magnitude
)
btor_norm = self.axis_Btor_norm.item()
pol_factor = self.poloidal_field_factor
r_ax = self.axis_r_norm.item()
z_ax = self.axis_z_norm.item()
@cfunc(lsoda_sig)
def poloidal_trace_eq(t, y, du, p): # pragma: no cover
r = y[0]
z = y[1]
psi_dr, psi_dz = _eval_psi_derivs(r, z, tx, ty, c, nx, ny, kx, ky)
factor = pol_factor / (r * denom_factor)
b_r = -factor * psi_dz
b_z = factor * psi_dr
dz = z - z_ax
dr = r - r_ax
theta = np.arctan2(dz, dr)
radius = np.sqrt(dr * dr + dz * dz)
b_along_fs = -b_r * np.sin(theta) + b_z * np.cos(theta)
du[0] = radius * b_r / np.abs(b_along_fs)
du[1] = radius * b_z / np.abs(b_along_fs)
du[2] = btor_norm * np.abs(radius / b_along_fs / (r * r))
self._poloidal_trace_addr = poloidal_trace_eq.address
return self._poloidal_trace_addr
[docs]
def get_boundary_polygon(self):
"""Return the divertor polygon."""
return self.divertor_polygon
[docs]
def get_limiter_polygon(self, quantile=0.999, npoints=1000,
tolerance=1e-4):
"""
Return the limiter polygon.
Defined as the last closed non-separatrix surface. With the quantile
argument, one can set how close in psi the limiter should be located.
"""
# NOTE: We currently ignore flux surface wander
integrator = MagneticFieldTracer(self, max_flux_surface_wander=1.0)
R_0 = np.array(self.axis_r_norm)
Z_0 = np.array(self.axis_z_norm)
R_max = np.max(self.get_boundary_polygon().x_points)
R_u = np.linspace(R_0, R_max, npoints)
psi_u = self.psi_interpolator(x=R_u, y=Z_0,
dx=0, dy=0, grid=False)
psi_start = self._psi_axis \
+ (self._psi_separatrix - self._psi_axis) * quantile
R_start = interp1d(psi_u, R_u, kind="linear")(psi_start)
poloidal_trace = integrator.poloidal_integration(
r_initial=float(R_start),
z_initial=float(Z_0),
theta_max=2.0*np.pi,
tolerance=tolerance)
theta = np.linspace(0.0, 2.0*np.pi, npoints, endpoint=False)
R, Z, _ = poloidal_trace.sol(theta)
limiter = Polygon2D(R, Z)
return limiter
[docs]
def safety_factor(self, rloc, zloc, tolerance=1e-4, return_sol=False,
theta_eval=None):
"""
Return the safety factor calculated at the given location.
Uses a fast poloidal field-line trace via a compiled cfunc. The
per-point traces are distributed over threads with dask; the lsoda
cfunc releases the GIL, so this gives a real speed-up for a full
q-profile. Optionally returns the raw lsoda solution array.
Parameters
----------
rloc, zloc : array_like
Normalized locations at which to evaluate the safety factor.
tolerance : float, optional
Relative and absolute tolerance passed to lsoda.
return_sol : bool, optional
If True, also return the interpolated field-line solutions.
theta_eval : array_like, optional
Poloidal angles at which to sample the trace when
``return_sol`` is True. Must be a strictly increasing 1D array
spanning a full poloidal turn [0, 2*pi] (so q stays correct).
Defaults to 200 points over [0, 2*pi].
Returns
-------
numpy.ndarray or (numpy.ndarray, list)
The safety factors, plus the solution callables if
``return_sol`` is True.
"""
rloc = np.atleast_1d(rloc)
zloc = np.atleast_1d(zloc)
assert rloc.size == zloc.size, "R and Z arrays must have same size!"
if theta_eval is not None:
theta_eval = np.asarray(theta_eval, dtype=float)
if theta_eval.ndim != 1 or theta_eval.size < 2:
raise ValueError(
"theta_eval must be a 1D array-like with >= 2 points.")
if np.any(np.diff(theta_eval) <= 0.0):
raise ValueError("theta_eval must be strictly increasing.")
if not (np.isclose(theta_eval[0], 0.0)
and np.isclose(theta_eval[-1], 2.0 * np.pi)):
raise ValueError(
"theta_eval must span a full poloidal turn [0, 2*pi].")
addr = self.make_poloidal_trace_eq()
results = dask.compute(*[
dask.delayed(self._single_safety_factor)(
addr, r, z, tolerance, return_sol, theta_eval)
for r, z in zip(rloc, zloc)
], scheduler='threads')
if return_sol:
q_factors, sols = zip(*results)
return np.array(q_factors), list(sols)
return np.array(results)
def _single_safety_factor(self, addr, r, z, tolerance=1e-4,
return_sol=False, theta_eval=None):
"""
Calculate the safety factor for a single location.
Near the magnetic axis (rho < axis_rho_tol) the poloidal trace is
singular, so q is evaluated on a small offset surface (at the
tighter axis_trace_tol) to approximate the on-axis limit q0.
Optionally returns a callable sol(theta) -> (3, N) array matching
the scipy OdeSolution interface used by eqdsk export. When
return_sol is True the trace is sampled at theta_eval (defaulting
to 200 points over [0, 2*pi]).
"""
# The poloidal trace is singular at the magnetic axis (0/0
# right-hand side). Below axis_rho_tol we offset the start radially
# by axis_offset in R (landing near rho ~= 4e-3, clear of the
# rho <~ 2e-4 breakdown region) to approximate the on-axis limit q0.
# The near-axis trace uses the tighter axis_trace_tol so the
# offset's O(rho^2) bias (~1e-5), not integration error, floors q0.
axis_rho_tol = 1e-3
axis_offset = 1e-3
axis_trace_tol = 1e-8
rho = self.normalized_flux_surface_label(r, z).item()
if rho >= 1.0:
return (np.nan, None) if return_sol else np.nan
if rho < axis_rho_tol:
warnings.warn(
f"safety_factor evaluated at rho={rho:.2e} < {axis_rho_tol:.0e}"
f" (near magnetic axis); tracing from an offset of "
f"{axis_offset:.0e} in R to approximate the on-axis limit q0.",
stacklevel=2,
)
r = r + axis_offset
tolerance = min(tolerance, axis_trace_tol)
if return_sol:
if theta_eval is None:
theta_eval = np.linspace(0.0, 2.0 * np.pi, 200)
usol, _ = lsoda(addr, np.array([r, z, 0.0]),
theta_eval, rtol=tolerance, atol=tolerance)
q_factor = usol[-1, 2] / (2.0 * np.pi)
sol = interp1d(theta_eval, usol.T, kind='cubic')
return (q_factor, sol)
usol, _ = lsoda(addr, np.array([r, z, 0.0]),
np.array([0.0, 2.0 * np.pi]),
rtol=tolerance, atol=tolerance)
return usol[-1, 2] / (2.0 * np.pi)
[docs]
def gradshaf_star(self, r_norm, z_norm, component=None):
"""
Evaluate the Laplace star operator of the Grad-Shafranov equation.
One can specify to evaluate the radial (component="R"),
vertical (component="Z") or the total operator (None).
"""
assert any([component == y for y in (None, "R", "Z")])
if (component is None) or (component == "R"):
# We use the product rule in the formulation to be able to utilize
# the derivative from the interpolator
dpsi = self.psi_interpolator(x=r_norm, y=z_norm,
dx=1, dy=0, grid=False)
d2psi = self.psi_interpolator(x=r_norm, y=z_norm,
dx=2, dy=0, grid=False)
res_R = d2psi - dpsi / r_norm
if component == "R":
return res_R
res_Z = self.psi_interpolator(x=r_norm, y=z_norm,
dx=0, dy=2, grid=False)
if component == "Z":
return res_Z
return res_R + res_Z
