import os
import numpy as np
from .catalog import has_static_product, read_static_variables
[docs]
class VelocityGrid:
"""
Velocity-space grid metadata for XGC distribution functions.
The XGC f0 mesh file stores a cylindrical velocity grid in
(v_perp, v_parallel) normalized to a reference thermal speed, with the
gyrophase angle removed. XGC distribution data on this grid already includes
the factor ``v_perp`` from the cylindrical Jacobian.
"""
REQUIRED_VARIABLES = (
"f0_nmu",
"f0_nvp",
"f0_smu_max",
"f0_vp_max",
"f0_dsmu",
"f0_dvp",
"f0_fg_T_ev",
"f0_T_ev",
"f0_den",
"f0_flow",
)
OPTIONAL_VARIABLES = (
"gradpsi",
"nb_curl_nb",
"v_curv",
"v_gradb",
"f0_grid_vol_vonly",
)
def __init__(self, work_dir=".", filename="xgc.f0.mesh.bp", catalog=None, source_reader=None):
"""
Initialize velocity-grid metadata.
Parameters
----------
work_dir : str, optional
Directory containing ``filename`` for the legacy local-file path.
filename : str, optional
BP product containing f0 velocity-grid metadata.
catalog : xgc_analysis.catalog.SimulationCatalog or None, optional
Optional catalog used to read ``xgc.f0.mesh.bp`` through the shared
static-product access path. Direct local-file fallback reads are
disabled.
source_reader : callable or None, optional
Optional static-product read backend. ``None`` reads regular local
BP files; future campaign-backed readers can pass an alternate
backend without changing this class.
"""
self.work_dir = work_dir
self.filename = os.path.join(work_dir, filename)
self.catalog = catalog
self.source_reader = source_reader
if has_static_product(catalog, filename, self.REQUIRED_VARIABLES):
values = read_static_variables(
catalog,
filename,
self.REQUIRED_VARIABLES + self.OPTIONAL_VARIABLES,
source_reader=source_reader,
missing="skip",
)
self._read_metadata_from_values(values)
else:
raise RuntimeError(
f"VelocityGrid requires catalog product '{filename}' with required velocity-grid variables; "
"direct local BP fallback is disabled."
)
self._build_coordinates()
def _read_metadata(self):
"""
Disabled legacy direct-reader placeholder.
Velocity-grid metadata is now read through catalog static products so
directory and campaign backends use the same access path.
"""
raise RuntimeError("VelocityGrid direct local BP reads are disabled; use a catalog.")
def _read_metadata_from_values(self, values):
"""
Populate velocity-grid metadata from catalog-read values.
Parameters
----------
values : mapping[str, object]
Values keyed by ADIOS variable name, typically returned by
:func:`xgc_analysis.catalog.read_static_variables`.
"""
self.nmu = int(np.asarray(values["f0_nmu"]).item())
self.nvp = int(np.asarray(values["f0_nvp"]).item())
self.smu_max = float(np.asarray(values["f0_smu_max"]).item())
self.vp_max = float(np.asarray(values["f0_vp_max"]).item())
self.dsmu = float(np.asarray(values["f0_dsmu"]).item())
self.dvp = float(np.asarray(values["f0_dvp"]).item())
self.fg_T_ev = np.asarray(values["f0_fg_T_ev"])
self.T_ev = np.asarray(values["f0_T_ev"])
self.den = np.asarray(values["f0_den"])
self.flow = np.asarray(values["f0_flow"])
self.gradpsi = self._optional_array(values, "gradpsi")
self.nb_curl_nb = self._optional_array(values, "nb_curl_nb")
self.v_curv = self._optional_array(values, "v_curv")
self.v_gradb = self._optional_array(values, "v_gradb")
self.grid_vol_vonly = self._optional_array(values, "f0_grid_vol_vonly")
@staticmethod
def _optional_array(values, name):
"""
Return an optional catalog-read array or ``None``.
Parameters
----------
values : mapping[str, object]
Static product values.
name : str
Optional variable name.
"""
if name not in values:
return None
return np.asarray(values[name])
def _build_coordinates(self):
# XGC stores nmu and nvp as interval counts. The actual grid vertices are:
# n_mu_points = nmu + 1
# n_vp_points = 2*nvp + 1 (symmetric around zero)
self.n_mu_points = self.nmu + 1
self.n_vp_points = 2 * self.nvp + 1
# Normalized v_perp coordinate used by XGC interpolation. The first point
# is shifted from 0 to dsmu/3 to preserve information with linear basis
# functions while the stored distribution already contains the v_perp
# Jacobian factor.
self.vperp_norm = np.arange(self.n_mu_points, dtype=float) * self.dsmu
if self.n_mu_points > 0:
self.vperp_norm[0] = self.dsmu / 3.0
self.vpara_norm = (
np.arange(self.n_vp_points, dtype=float) * self.dvp - self.vp_max
)
self.mu0_factor = self.vperp_norm[0] / self.dsmu if self.dsmu != 0 else 0.0
@property
def shape(self):
"""Velocity-grid point shape in XGC storage order: (n_mu, n_vp)."""
return (self.n_mu_points, self.n_vp_points)
[docs]
def mu_edge_factors(self):
fac = np.ones(self.n_mu_points, dtype=float)
if self.n_mu_points > 0:
fac[0] = 0.5
fac[-1] = 0.5
return fac
[docs]
def vp_edge_factors(self):
fac = np.ones(self.n_vp_points, dtype=float)
if self.n_vp_points > 0:
fac[0] = 0.5
fac[-1] = 0.5
return fac
[docs]
def mu_integration_weights(self, *, data_includes_vperp=True):
"""
1D weights for the v_perp axis (normalized units).
If ``data_includes_vperp=True`` (XGC default), only trapezoidal weights
and ``dsmu`` are applied because the stored values already carry the
cylindrical Jacobian factor ``v_perp``.
"""
w = self.mu_edge_factors() * self.dsmu
if not data_includes_vperp:
w = w * self.vperp_norm
return w
[docs]
def vp_integration_weights(self):
"""1D trapezoidal weights for the v_parallel axis (normalized units)."""
return self.vp_edge_factors() * self.dvp
[docs]
def integration_weights_2d(self, *, data_includes_vperp=True, include_gyroangle=True):
"""
2D velocity-space integration weights in normalized coordinates.
Returns weights for arrays shaped ``(..., n_mu, ..., n_vp)`` after axes
are aligned, with optional ``2*pi`` gyrophase factor.
"""
w2d = np.outer(
self.mu_integration_weights(data_includes_vperp=data_includes_vperp),
self.vp_integration_weights(),
)
if include_gyroangle:
w2d = (2.0 * np.pi) * w2d
return w2d
[docs]
def integrate_over_velocity(
self,
values,
*,
axis_mu=-3,
axis_vp=-1,
node_axis=None,
data_includes_vperp=True,
include_gyroangle=True,
apply_grid_vol_vonly=False,
species_index=None,
):
"""
Integrate an array over (v_perp, v_parallel) in normalized coordinates.
Parameters
----------
values : np.ndarray
Array containing velocity-grid data.
axis_mu, axis_vp : int
Indices of the ``v_perp`` and ``v_parallel`` axes in ``values``.
node_axis : int | None
Index of the configuration-space node axis in ``values``. Required
when ``apply_grid_vol_vonly=True``.
data_includes_vperp : bool
``True`` for XGC distribution data (`*_f`), which already includes
the ``v_perp`` Jacobian factor.
include_gyroangle : bool
Multiply by ``2*pi`` if the gyrophase angle has been integrated out.
apply_grid_vol_vonly : bool
If ``True``, multiply by species/node-dependent ``f0_grid_vol_vonly``
before integrating over velocity.
species_index : int | None
Species index into ``f0_grid_vol_vonly``. Required when
``apply_grid_vol_vonly=True`` and multiple species are present.
"""
arr = np.asarray(values)
if apply_grid_vol_vonly:
if self.grid_vol_vonly is None:
raise ValueError("f0_grid_vol_vonly is not available in this VelocityGrid.")
if node_axis is None:
raise ValueError("node_axis must be provided when apply_grid_vol_vonly=True.")
gv = np.asarray(self.grid_vol_vonly)
if gv.ndim == 2:
if species_index is None:
if gv.shape[0] == 1:
gv = gv[0]
else:
raise ValueError(
"species_index is required for multi-species f0_grid_vol_vonly."
)
else:
gv = gv[int(species_index)]
elif gv.ndim != 1:
raise ValueError(
f"Unsupported f0_grid_vol_vonly shape {gv.shape}; expected (nnode,) or (nsp,nnode)."
)
if gv.ndim != 1:
gv = np.asarray(gv).reshape(-1)
node_axis_norm = int(node_axis)
if node_axis_norm < 0:
node_axis_norm += arr.ndim
if node_axis_norm < 0 or node_axis_norm >= arr.ndim:
raise ValueError(f"node_axis={node_axis} is out of range for array with ndim={arr.ndim}.")
if arr.shape[node_axis_norm] != gv.shape[0]:
raise ValueError(
f"Node axis length {arr.shape[node_axis_norm]} does not match f0_grid_vol_vonly length {gv.shape[0]}."
)
shape = [1] * arr.ndim
shape[node_axis_norm] = gv.shape[0]
arr = arr * gv.reshape(shape)
arr = np.moveaxis(arr, (axis_mu, axis_vp), (-2, -1))
if arr.shape[-2] != self.n_mu_points or arr.shape[-1] != self.n_vp_points:
raise ValueError(
"Velocity axes do not match grid shape "
f"{self.shape}; got {arr.shape[-2:]}"
)
if apply_grid_vol_vonly:
# f0_grid_vol_vonly already carries species/node-dependent velocity
# normalization including dsmu*dvp. Keep only boundary factors here.
w2d = np.outer(self.mu_edge_factors(), self.vp_edge_factors())
else:
w2d = self.integration_weights_2d(
data_includes_vperp=data_includes_vperp,
include_gyroangle=include_gyroangle,
)
return np.sum(arr * w2d, axis=(-2, -1))