import re
import numpy as np
from .accessor_mixin import ArrayAccessorMixin
from .bp_reader_mixin import BPReaderMixin
from .mesh_data import MeshData
from .plane_data import PlaneData
from .velocity_grid import VelocityGrid
_DIST_VAR_RE = re.compile(r"^[A-Za-z][A-Za-z0-9]*_f$")
[docs]
class DistributionFunctionField:
"""
Wrapper for one XGC distribution-function variable on a mesh and velocity grid.
Python-facing storage order:
- canonical: ``(phi, node, vpar, vperp)``
Axisymmetric files that omit the toroidal dimension are normalized to
``phi=1`` on read, so ``self.data`` is always 4D.
"""
def __init__(self, mesh, velocity_grid, data_array, *, name=None):
"""
Store one distribution-function array with its interpretation context.
Parameters
----------
mesh : xgc_analysis.mesh.Mesh
Mesh used to validate and slice the configuration-space node and
toroidal dimensions.
velocity_grid : xgc_analysis.velocity_grid.VelocityGrid
Velocity-grid metadata used to validate the ``vpar``/``vperp``
dimensions and to compute velocity-space moments.
data_array : array-like
Distribution values in canonical Python order
``(phi, node, vpar, vperp)``.
name : str or None, optional
Variable name used for diagnostics and error context.
"""
self.mesh = mesh
self.velocity_grid = velocity_grid
self.name = name
self.data = np.asarray(data_array)
self._validate_shape()
def _validate_shape(self):
"""
Validate the canonical distribution-function array shape.
Raises
------
ValueError
If the array is not four-dimensional, if its trailing
``(node, vpar, vperp)`` dimensions do not match the mesh and
velocity grid, or if the toroidal dimension is incompatible with
the mesh.
"""
shp = self.data.shape
nmu, nvp = self.velocity_grid.shape
nnode = self.mesh.get_plane(0).n_n
if self.data.ndim == 4:
if shp[-3:] != (nnode, nvp, nmu):
raise ValueError(
f"Invalid distribution trailing shape {shp[-3:]}, "
f"expected {(nnode, nvp, nmu)} for (node, vpar, vperp)"
)
if shp[0] not in (1, self.mesh.nphi):
raise ValueError(
f"Leading toroidal dimension must be 1 or mesh.nphi={self.mesh.nphi}, "
f"got {shp[0]}"
)
self.has_toroidal_dim = True
return
raise ValueError(
"DistributionFunctionField expects a 4D array in canonical order "
"(phi, node, vpar, vperp)."
)
[docs]
def get_data(self):
"""Return the underlying NumPy array in XGC storage order."""
return self.data
@property
def nphi(self):
"""Return the stored toroidal dimension length."""
return self.data.shape[0]
@property
def is_axisymmetric_storage(self):
"""Return True when the array is stored with a single toroidal plane."""
return self.nphi == 1
[docs]
def plane_slice(self, iphi=0):
"""Return a view with shape ``(node, vpar, vperp)``."""
return self.data[iphi]
[docs]
def node_slice(self, inode, *, iphi=0):
"""Return a view with shape ``(vpar, vperp)`` for one configuration-space node."""
return self.plane_slice(iphi=iphi)[inode, :, :]
[docs]
def velocity_integral(self, *, data_includes_vperp=True, include_gyroangle=True):
"""
Integrate over velocity space and return config-space data.
Returns
-------
np.ndarray
Axisymmetric: shape ``(n_node,)``
3D: shape ``(nphi, n_node)``
"""
out = self.velocity_grid.integrate_over_velocity(
self.data,
axis_mu=3,
axis_vp=2,
data_includes_vperp=data_includes_vperp,
include_gyroangle=include_gyroangle,
)
if self.data.shape[0] == 1:
return out[0]
return out
[docs]
def velocity_integral_data(self, *, data_includes_vperp=True, include_gyroangle=True):
"""
Velocity-integrated result packaged as PlaneData or MeshData.
This is intentionally a thin wrapper around :meth:`velocity_integral`.
"""
arr = self.velocity_integral(
data_includes_vperp=data_includes_vperp,
include_gyroangle=include_gyroangle,
)
if arr.ndim == 1:
return PlaneData(
plane=self.mesh.get_plane(0),
data_array=arr,
n_components=1,
dtype=arr.dtype,
)
return MeshData(
self.mesh,
data_array=arr,
n_components=1,
dtype=arr.dtype,
mesh_is_axisym=self.mesh.is_axisymmetric,
)
[docs]
def compute_f0_moments(self, inputs, *, calculator=None):
"""
Compute electrostatic f0 moments using ``distribution_moments.F0MomentCalculator``.
Parameters
----------
inputs : F0MomentInputs
Geometry/field inputs for the requested moments.
calculator : F0MomentCalculator or None, optional
Reuse an existing calculator instance. If omitted, a new one is
created from this field's ``velocity_grid``.
"""
if calculator is None:
from .distribution_moments import F0MomentCalculator
calculator = F0MomentCalculator(self.velocity_grid)
return calculator.compute(self, inputs)
[docs]
class DistributionFunctionData(BPReaderMixin, ArrayAccessorMixin):
"""
Reader for XGC ``xgc.f0.XXXXX.bp`` distribution-function files.
Data layout matches other XGC-Analysis readers:
``self.data[var_name][file_step_index] = object``
where ``object`` is usually:
- :class:`DistributionFunctionField` for species ``*_f`` arrays
- :class:`PlaneData` / :class:`MeshData` for configuration-space arrays
- scalar / np.ndarray for metadata
"""
DEFAULT_IGNORED_VARS = {
"eden_f0",
"iden_f0_approx",
"iphi",
"nphi",
"nnode",
"nmup1",
"ndata",
"imu1m1",
"inode1m1",
"mudata",
"vpdata",
}
def __init__(
self,
mesh=None,
work_dir=".",
file_indices=None,
*,
simulation=None,
velocity_grid=None,
f0_mesh_filename="xgc.f0.mesh.bp",
ignored_vars=None,
variables=None,
catalog=None,
steps=None,
read_all_steps=False,
source_reader=None,
missing="raise",
):
"""
Initialize the distribution-function reader.
Parameters
----------
mesh : xgc_analysis.mesh.Mesh or None, optional
Mesh used to wrap configuration-space and distribution-function
arrays. If omitted and ``simulation`` is provided,
``simulation.mesh`` is used.
work_dir : str, optional
Directory containing ``xgc.f0.XXXXX.bp`` files for the legacy
filename-based read path. If ``simulation`` is provided and
``work_dir`` is left at the default, ``simulation.data_directory``
is used.
file_indices : iterable[int] or None, optional
Legacy filename indices to read when ``steps`` is omitted. With
catalog-backed reads these are interpreted as logical steps. This
preserves the common file-sequence convention where
``xgc.f0.00010.bp`` contributes logical step ``10`` when no
explicit ``gstep`` is present.
simulation : xgc_analysis.simulation.Simulation or None, optional
Preferred source for shared interpretation state. When provided,
the reader inherits ``mesh``, ``velocity_grid``, and ``catalog``
unless explicit overrides are supplied.
velocity_grid : xgc_analysis.velocity_grid.VelocityGrid or None, optional
Velocity-grid metadata. If omitted, ``simulation.velocity_grid`` is
used when available. If needed, the reader constructs a
:class:`VelocityGrid` from the catalog static product.
f0_mesh_filename : str, optional
Velocity-grid BP product used when constructing ``velocity_grid``.
ignored_vars : iterable[str] or None, optional
Variables skipped during default reads. Explicit ``variables``
requests are not filtered by this set.
variables : str or iterable[str] or None, optional
Optional explicit variable names to read. ``None`` reads all
non-ignored variables, preserving the original behavior.
catalog : xgc_analysis.catalog.SimulationCatalog or None, optional
Catalog used to discover and read ``xgc.f0.bp``. If omitted and
``simulation`` has a catalog, ``simulation.catalog`` is used.
steps : int or iterable[int] or None, optional
Explicit logical XGC steps to read. This supersedes
``file_indices`` when provided.
read_all_steps : bool, optional
If True and neither ``steps`` nor ``file_indices`` is provided, read
all logical steps advertised by the catalog product.
source_reader : callable or None, optional
Optional catalog source-reader backend. If omitted, the catalog's
installed source reader is used.
missing : {"raise", "skip", "zero"}, optional
Missing-step policy passed to catalog read planning.
"""
if simulation is not None:
if mesh is None:
mesh = getattr(simulation, "mesh", None)
if velocity_grid is None:
velocity_grid = getattr(simulation, "velocity_grid", None)
if catalog is None:
catalog = getattr(simulation, "catalog", None)
if work_dir == "." and getattr(simulation, "data_directory", None) is not None:
work_dir = simulation.data_directory
if mesh is None:
raise ValueError("DistributionFunctionData requires `mesh` or `simulation.mesh`.")
self.mesh = mesh
self.work_dir = work_dir
self.file_indices = self._normalize_steps(file_indices) if file_indices is not None else None
self.steps = self._normalize_steps(steps) if steps is not None else None
self.simulation = simulation
self.catalog = catalog
self.source_reader = source_reader
self.missing = missing
self.velocity_grid = (
velocity_grid
if velocity_grid is not None
else VelocityGrid(work_dir=work_dir, filename=f0_mesh_filename, catalog=catalog, source_reader=source_reader)
)
self.ignored_vars = set(self.DEFAULT_IGNORED_VARS if ignored_vars is None else ignored_vars)
self._init_bp_reader_state(variables=variables, read_all_steps=read_all_steps)
self.data = {}
self._read_files()
@staticmethod
def _normalize_variables(variables):
"""
Normalize an optional variable selector.
Parameters
----------
variables : str or iterable[str] or None
User-facing variable selector.
Returns
-------
list[str] or None
Deduplicated variable list, or ``None`` for default reads.
"""
if variables is None:
return None
if isinstance(variables, str):
variable_list = [variables]
else:
variable_list = list(variables)
if not variable_list:
raise ValueError("`variables` must not be empty when provided.")
normalized = []
for variable in variable_list:
if not isinstance(variable, str):
raise TypeError("Each entry in `variables` must be a string.")
if variable not in normalized:
normalized.append(variable)
return normalized
@staticmethod
def _normalize_steps(steps):
"""
Normalize an integer or iterable step selector.
Parameters
----------
steps : int or iterable[int]
Logical steps or legacy file indices supplied by the caller.
Returns
-------
list[int]
Deduplicated integer step list preserving caller order.
"""
if isinstance(steps, int):
step_list = [steps]
else:
step_list = list(steps)
normalized = []
for step in step_list:
step_value = int(step)
if step_value not in normalized:
normalized.append(step_value)
return normalized
@staticmethod
def _is_distribution_var(name):
"""
Return whether ``name`` follows the XGC distribution-function convention.
Distribution arrays are currently identified by variables ending in
``"_f"`` with an alphanumeric species or component prefix.
"""
return bool(_DIST_VAR_RE.match(name))
def _wrap_config_space_array(self, value):
"""
Wrap non-distribution arrays in mesh-aware containers when possible.
Parameters
----------
value : object
Raw value returned by catalog read-plan execution.
Returns
-------
PlaneData | MeshData | np.ndarray
One-dimensional or axisymmetric two-dimensional arrays become
``PlaneData``. Full toroidal-plane arrays become ``MeshData``.
Other shapes are returned unchanged as NumPy arrays.
"""
arr = np.asarray(value)
if arr.ndim == 1:
return PlaneData(
plane=self.mesh.get_plane(0),
data_array=arr,
n_components=1,
dtype=arr.dtype,
)
if arr.ndim == 2 and arr.shape[0] == 1:
return PlaneData(
plane=self.mesh.get_plane(0),
data_array=np.squeeze(arr, axis=0),
n_components=1,
dtype=arr.dtype,
)
if arr.ndim == 2 and arr.shape[0] == self.mesh.nphi:
return MeshData(
self.mesh,
data_array=arr,
n_components=1,
dtype=arr.dtype,
mesh_is_axisym=self.mesh.is_axisymmetric,
)
return arr
@staticmethod
def _transpose_distribution_array(value):
"""
Convert XGC storage order to canonical Python order ``(phi, node, vpar, vperp)``.
XGC orders:
- axisymmetric: (vperp, node, vpar)
- 3D: (phi, vperp, node, vpar)
"""
arr = np.asarray(value)
if arr.ndim == 3:
# (vperp, node, vpar) -> (1, node, vpar, vperp)
return np.transpose(arr, (1, 2, 0))[np.newaxis, ...]
if arr.ndim == 4:
# (phi, vperp, node, vpar) -> (phi, node, vpar, vperp)
return np.transpose(arr, (0, 2, 3, 1))
raise ValueError(
f"Unexpected distribution array rank {arr.ndim}; expected 3D or 4D."
)
def _read_files(self):
"""
Read distribution-function data through catalog read plans.
``file_indices`` are retained as a legacy user-facing selector, but the
actual source and ADIOS-step mapping comes from the catalog. This makes
file sequences, internal ADIOS steps, and future campaign sources use
the same read path.
"""
if self._read_selector_is_empty():
return
if self.catalog is None:
raise RuntimeError(
"DistributionFunctionData requires a catalog for reading xgc.f0.bp; "
"direct local BP filename reads are disabled."
)
product_key = "xgc.f0.bp"
variables = self._catalog_variables(product_key)
if not variables:
return
file_data = self._read_catalog_product(
self.catalog,
product_key,
variables,
steps=self.steps,
default_steps=self.file_indices,
read_all_steps=self.read_all_steps,
missing=self.missing,
require_all_variables=True,
source_reader=self.source_reader,
)
for step_index in sorted(file_data):
self._store_step_variables(step_index, file_data[step_index])
def _read_selector_is_empty(self):
"""
Return True when caller input explicitly requests no distribution reads.
This preserves lightweight construction patterns used by tests and
setup code that only need inherited ``mesh``/``velocity_grid`` state.
"""
if self.steps == []:
return True
return self.steps is None and self.file_indices == []
def _catalog_variables(self, product_key):
"""
Return distribution variables to request from the catalog product.
Explicit ``variables`` are honored exactly. Default reads use all
variables advertised by the product except ``DEFAULT_IGNORED_VARS``.
"""
if self.requested_vars is not None:
return list(self.requested_vars)
product = self.catalog.get_product(product_key)
return [name for name in sorted(product.variables) if name not in self.ignored_vars]
def _store_step_variables(self, step_index, variables):
"""
Wrap and store all variables read for one reader-local step.
Parameters
----------
step_index : int
Reader-local step index produced by catalog read-plan execution.
variables : dict[str, object]
Raw ADIOS values for that step.
"""
for var_name, value in variables.items():
if var_name == "time":
self.data.setdefault(var_name, {})[step_index] = float(np.asarray(value).item())
continue
if var_name in ("gstep", "tindex", "step"):
self.data.setdefault(var_name, {})[step_index] = int(np.asarray(value).item())
continue
if isinstance(value, np.ndarray) and self._is_distribution_var(var_name):
self.data.setdefault(var_name, {})[step_index] = DistributionFunctionField(
self.mesh,
self.velocity_grid,
self._transpose_distribution_array(value),
name=var_name,
)
continue
if isinstance(value, np.ndarray):
self.data.setdefault(var_name, {})[step_index] = self._wrap_config_space_array(value)
else:
self.data.setdefault(var_name, {})[step_index] = value
[docs]
def distribution_variables(self):
"""Return sorted variable names that store distribution-function data."""
return sorted(k for k, v in self.data.items() if v and hasattr(next(iter(v.values())), "velocity_grid"))
# --- Typed accessors ---
[docs]
def get_distribution(self, var_name, step_index=0):
"""Return ``var_name`` at ``step_index`` as :class:`DistributionFunctionField`."""
return self.get_as(var_name, step_index, DistributionFunctionField)
[docs]
def get_plane_data(self, var_name, step_index=0):
"""Return ``var_name`` at ``step_index`` as :class:`PlaneData`."""
return self.get_as(var_name, step_index, PlaneData)
[docs]
def get_mesh_data(self, var_name, step_index=0):
"""Return ``var_name`` at ``step_index`` as :class:`MeshData`."""
return self.get_as(var_name, step_index, MeshData)
[docs]
def get_scalar(self, var_name, step_index=0):
"""Return ``var_name`` at ``step_index`` as a scalar numeric value."""
return self.get_as(var_name, step_index, (int, float, np.integer, np.floating))