"""Defines routines for calculating integral quantities."""
import numpy as np
import xarray as xr
from scipy.interpolate import RectBivariateSpline
import warnings
from torx.arrays import make_xarray
from torx.grid import Grid2D, Grid3D
from torx.decorators import autodoc_function
from torx.units import Normalization, Quantity
from typing import Callable, Sequence
from torx.equilibrium import EquilibriumBaseClass
from .lineouts import LineoutSet
def poloidal_integration(
grid: Grid2D, norm: Normalization, poloidal_function: xr.DataArray
):
"""
Return the 2D poloidal integral of a function.
Off-grid values are assigned a value of 0.0, and a filled grid spline
interpolator is used to evaluate the integral.
"""
assert hasattr(poloidal_function, "norm")
assert isinstance(poloidal_function.norm, Quantity)
grid_limits = {
"xa": grid.r_s.min(),
"xb": grid.r_s.max(),
"ya": grid.z_s.min(),
"yb": grid.z_s.max(),
}
poloidal_function_values = np.nan_to_num(poloidal_function.values, nan=0.0)
function_integral = RectBivariateSpline(
grid.r_s, grid.z_s, poloidal_function_values.T
).integral(**grid_limits)
norm = poloidal_function.norm * norm.R0**2
return function_integral * norm
def axisymmetric_cylindrical_integration(
grid: Grid2D, norm: Normalization, poloidal_function: xr.DataArray
):
"""
Return the 3D cylindrically-integrated value of a function.
Note
----
Automatically takes care of the "R" factor introduced by the theta
integral.
"""
r_s_grid, _ = np.meshgrid(grid.r_s, grid.z_s)
poloidal_integration_function = make_xarray(
poloidal_function * r_s_grid, norm=poloidal_function.norm
)
function_integral = (
2.0
* np.pi
* norm.R0
* poloidal_integration(grid, norm, poloidal_integration_function)
)
return function_integral
[docs]
@autodoc_function
def flux_surf_avg_3D(
grid: Grid3D,
vol: xr.DataArray,
field: xr.DataArray,
rho_func: Callable,
rho_array: Sequence,
drho: float = 1e-4,
):
"""
Return the flux-surface average of a given field on a 3D grid.
Integrates the field over an infinitesimal volume shell centered on the flux
surface [1]_.
Parameters
----------
grid : Grid3D
The 3D grid in which the field is defined.
vol : xr.DataArray
The fluxbox volume at each point of the grid
(can be obtained from the maps).
field : xr.DataArray
The field to be averaged.
Must have dimensions compatible with the grid.
rho_func : Callable
A function that computes the flux surface label (rho)
from R, Z, phi coordinates.
rho_array : Sequence
The array of rho values at which to compute the flux-surface average
drho : float, optional
The extent of rho to be used in the volume integral, by default 1e-4.
Returns
-------
xr.DataArray
The flux-surface averaged field.
References
----------
.. [1] A. Stegmeir et al., CPC 2026, Appendix C.
"""
R = grid.r_u
Z = grid.z_u
phi = grid.coords["phi"]
rho_array = xr.DataArray(rho_array, dims="rho")
# Check if the desired dimension name to be replaced exists
if 'planes' in field.dims:
field = field.rename({'planes': 'phi'})
if 'npoints_max' in field.dims:
field = field.rename({'npoints_max': 'points'})
# Compute rho field
rho_values = xr.apply_ufunc(
rho_func,
R,
Z,
phi,
input_core_dims=[["points"], ["points"], []],
output_core_dims=[["points"]],
vectorize=True,
dask='parallelized',
output_dtypes=[np.float64]
).transpose()
# Define a small function to compute the flux-surface average for one rho
def _fsa_for_rho(r, rho_values, field, vol, drho):
mask = (rho_values >= r - drho / 2) & (rho_values < r + drho / 2)
if mask.sum() == 0:
raise ValueError(f"No points found in the rho range [{r - drho / 2}, {r + drho / 2}]")
return (vol[mask]*field[mask]).sum() / vol[mask].sum()
# Use apply_ufunc to vectorize over rho_array
fsa = xr.apply_ufunc(
_fsa_for_rho,
rho_array,
rho_values,
field,
vol,
drho,
input_core_dims=[
[], ["points", "phi"],
["points", "phi"],
["points", "phi"], []
],
output_core_dims=[[]],
vectorize=True,
dask='parallelized',
output_dtypes=[np.float64]
).assign_coords(rho=rho_array)
return fsa
[docs]
@autodoc_function
def drift_flux(
grid: Grid3D,
equi: EquilibriumBaseClass,
lineout_set: LineoutSet,
norm: Normalization,
field: xr.DataArray,
coeff: xr.DataArray = None,
n_samples: int=500,
total: bool=True,
):
r"""
Return the drift flux of a given 3D field through a lineout array [1]_.
$$\int_S v * (B/B^2 \times \nabla u) dS$$,
where $S$ is the surface defined by the lineout array,
$u$ is the field to be integrated,
and $v$ is an optional coefficient, 1 by default.
The integrand can be expressed via a derivative along the polygon
$$\int v * 1/Btor du/dl J dl dphi$$
where $l$ measures the length along the polygon.
Parameters
----------
grid : Grid3D
The 3D grid in which the field is defined.
equi : EquilibriumBaseClass
The equilibrium class containing the magnetic field information.
lineout_set : LineoutSet
The array of lineouts through which to compute the drift flux.
norm : Normalization
The normalization object containing the needed constants.
field : xr.DataArray
The field to be averaged.
Must have dimensions compatible with the grid.
coeff : xr.DataArray, optional
An optional coefficient to multiply the integrand by, defaults to 1.
n_samples : int, optional
The number of samples to use for the lineout sampling, by default 500.
Only used if the lineout_set does not already have n_samples defined.
total : bool, optional
Whether to return the total drift flux or the flux at each point,
defaults to True.
Returns
-------
float
The drift flux through the lineout array.
References
----------
.. [1] K. Eder et al 2025 Plasma Phys. Control. Fusion 67 065034,
section 3.2.
"""
R = grid.r_u
Z = grid.z_u
phi = grid.coords["phi"]
dphi = phi[1] - phi[0]
nplanes = len(phi)
# Depending on how the data was loaded, the dimension names may differ.
# Rename dimensions if necessary to match the torx convention
if 'planes' in field.dims:
field = field.rename({'planes': 'phi'})
if 'npoints_max' in field.dims:
field = field.rename({'npoints_max': 'points'})
if not hasattr(lineout_set, 'n_samples'):
warnings.warn("lineout_set does not have n_samples attribute.")
lineout_set.find_points_on_grid(grid,
n_samples=n_samples)
Jacobian = xr.apply_ufunc(
equi.jacobian,
R,
Z,
phi,
input_core_dims=[["points"], ["points"], []],
output_core_dims=[["points"]],
vectorize=True,
dask='parallelized',
output_dtypes=[np.float64]
)
# kwargs={"is_unstructured": True} is needed for flare
btor = xr.apply_ufunc(
equi.magfield_component_toroidal,
R,
Z,
phi,
input_core_dims=[["points"], ["points"], []],
output_core_dims=[["points"]],
vectorize=True,
dask='parallelized',
output_dtypes=[np.float64],
kwargs={"is_unstructured": True}
)
# Define a Dataset to interpolate and roll in one go
ds_to_interp = xr.Dataset({
"Jacobian": Jacobian,
"btor": btor,
"field": field
})
if coeff is not None:
ds_to_interp["coeff"] = coeff
# Interpolate the necessary fields onto the lineout points
interp_ds = lineout_set.interpolate(ds_to_interp)
# Use roll if the lineouts are periodic, shift otherwise.
if lineout_set.is_periodic:
interp_ds_rolled = interp_ds.roll(interp_points=1)
else:
warn_msg = "Lineouts are not periodic. " + \
"One point per lineout will be a NaN."
warnings.warn(warn_msg)
interp_ds_rolled = interp_ds.shift(interp_points=1)
# If an optional field is provided, interpolate it
# and average it along the lineout, otherwise set it to 1.0
if "coeff" in interp_ds:
coeff_av = 0.5 * (interp_ds.coeff + interp_ds_rolled.coeff)
else:
coeff_av = 1.0
# Average the Jacobian and btor between adjacent points.
jac_av = 0.5 * (interp_ds.Jacobian + interp_ds_rolled.Jacobian)
btor_av = 0.5 * (interp_ds.btor + interp_ds_rolled.btor)
# Compute the drift flux
interp_field_diff = interp_ds_rolled.field - interp_ds.field
drift_flux = coeff_av * jac_av * interp_field_diff * dphi / btor_av
# Compute the normalization factor for the drift flux.
coeff_norm = coeff.attrs.get('norm', 1.0) if hasattr(coeff, 'attrs') else 1.0
drift_flux_norm = coeff_norm * norm.R0 * field.norm / norm.B0
if total:
return make_xarray(
drift_flux.sum(),
norm=drift_flux_norm,
)
else:
return make_xarray(
drift_flux,
norm=drift_flux_norm,
)
