"""Reader for xgc.neutrals.XXXXX.bp neutral diagnostic files."""
import os
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
[docs]
class NeutralData(BPReaderMixin, ArrayAccessorMixin):
"""
Read neutral Monte-Carlo diagnostic data from ``xgc.neutrals.XXXXX.bp`` files.
The data layout follows the pattern used by ``FieldData``/``FsourceData``:
self.data[var_name][file_step_index] = PlaneData | MeshData | scalar
For axisymmetric data, mesh fields are stored as ``PlaneData``.
For non-axisymmetric data, mesh fields are stored as ``MeshData``.
Axisymmetry is inferred from the shape of the variables in the neutrals file
(not from the mesh geometry):
- axisymmetric neutrals: ``(n_n,)`` or ``(1, n_n)``
- non-axisymmetric neutrals: ``(nphi, n_n)``
The inferred mode is tracked in ``self.is_axisymmetric``. If multiple files are
loaded, all neutrals fields are required to use the same axisymmetry mode.
"""
required_vars = ["den_neut", "rel_std", "temp_neut"]
def __init__(
self,
mesh,
work_dir="./",
file_indices=None,
is_axisymmetric=None,
variables=None,
read_all_steps=False,
catalog=None,
steps=None,
missing="raise",
source_reader=None,
):
"""
Initialize the neutrals reader.
Parameters
----------
mesh : Mesh
XGC mesh object used to construct ``PlaneData``/``MeshData`` wrappers.
work_dir : str, optional
Directory containing ``xgc.neutrals.XXXXX.bp`` files.
file_indices : list[int], optional
File indices to load (e.g. ``[100, 200]`` reads
``xgc.neutrals.00100.bp`` and ``xgc.neutrals.00200.bp``).
is_axisymmetric : bool or None, optional
Expected neutrals-data symmetry. If ``None`` (default), it is inferred
from the first neutrals field encountered. If provided, a mismatch
between the inferred shape and this value raises ``ValueError``.
variables : list[str] | str | None, optional
Optional explicit variable names. If omitted, the standard neutral
field variables are read.
read_all_steps : bool, optional
If True, read all ADIOS steps from selected sources.
catalog : xgc_analysis.catalog.SimulationCatalog or None, optional
Optional catalog used to resolve logical steps into BP sources.
Direct filename fallback reads are disabled; construct a directory
or campaign catalog before constructing this reader.
steps : iterable[int] or None, optional
Logical XGC steps to read from ``catalog``.
missing : {"raise", "skip", "zero"}, optional
Missing-step policy for explicit variable requests.
source_reader : callable or None, optional
Optional read-plan backend hook.
"""
self.mesh = mesh
self.work_dir = work_dir
self._file_indices_provided = file_indices is not None
self.file_indices = file_indices if file_indices is not None else [0]
# Inferred from the data shape when not explicitly provided.
self.is_axisymmetric = is_axisymmetric
self.catalog = catalog
self.catalog_steps = None if steps is None else [int(step) for step in steps]
self.missing = missing
self.source_reader = source_reader
self._init_bp_reader_state(variables=variables, read_all_steps=read_all_steps)
self.variables = self.requested_vars if self.requested_vars is not None else list(self.required_vars)
self.data = {}
if self.catalog is None:
raise RuntimeError("NeutralData requires a catalog; direct xgc.neutrals filename reads are disabled.")
self._read_from_catalog()
def _read_file_steps(self, fname):
"""
Read file steps using neutral-reader variable selection rules.
Parameters
----------
fname : str
Full path to one ``xgc.neutrals`` BP file.
Returns
-------
dict[int, dict[str, object]]
Mapping from BP step id to variable/value dictionary.
"""
return super()._read_file_steps(fname, read_vars=self.variables)
def _read_from_catalog(self):
"""
Read ``xgc.neutrals.bp`` through catalog read plans.
Returns
-------
None
Populates ``self.data`` and ``self.step_index_info`` in place.
"""
default_steps = self.file_indices if self._file_indices_provided else None
step_variables = self._read_catalog_product(
self.catalog,
"xgc.neutrals.bp",
self.variables,
steps=self.catalog_steps,
default_steps=default_steps,
read_all_steps=self.read_all_steps,
missing=self.missing,
require_all_variables=self.requested_vars is not None,
source_reader=self.source_reader,
)
for step_index in sorted(step_variables):
variables = step_variables[step_index]
for var in self.variables:
value = variables.get(var, 0.0)
self.data.setdefault(var, {})[step_index] = self._wrap_mesh_field(value)
def _read_files(self):
"""
Read requested neutrals files and populate ``self.data``.
Each variable is stored as ``self.data[var_name][i]``, where ``i`` is the
index into ``self.file_indices``. Values are wrapped as ``PlaneData`` or
``MeshData`` depending on the neutrals-data axisymmetry inferred from the
variable shape.
"""
for idx in self.file_indices:
fname = os.path.join(self.work_dir, f"xgc.neutrals.{idx:05d}.bp")
file_data = self._read_file_steps(fname)
for bp_step, variables in file_data.items():
step_index = self._register_step(idx, bp_step)
for var in self.variables:
value = variables.get(var, 0.0)
wrapped = self._wrap_mesh_field(value)
self.data.setdefault(var, {})[step_index] = wrapped
def _wrap_mesh_field(self, value):
"""
Wrap a raw neutrals field array as ``PlaneData`` or ``MeshData``.
Parameters
----------
value : np.ndarray or scalar
Raw variable value from catalog read-plan execution.
Returns
-------
PlaneData | MeshData | scalar
Axisymmetric neutrals fields become ``PlaneData``. Non-axisymmetric
neutrals fields become ``MeshData``. Non-array values are returned
unchanged.
"""
if not isinstance(value, np.ndarray):
return value
data_is_axisymmetric = self._infer_data_axisymmetry(value)
if self.is_axisymmetric is None:
self.is_axisymmetric = data_is_axisymmetric
elif self.is_axisymmetric != data_is_axisymmetric:
raise ValueError(
"Inconsistent neutrals-data axisymmetry detected: "
f"expected is_axisymmetric={self.is_axisymmetric}, "
f"but encountered array shape {value.shape} "
f"(inferred is_axisymmetric={data_is_axisymmetric})."
)
if data_is_axisymmetric:
arr = np.squeeze(value)
if arr.ndim != 1:
raise ValueError(
f"Expected axisymmetric neutral field to be 1D after squeeze, got shape {value.shape}"
)
return PlaneData(
plane=self.mesh.get_plane(0),
data_array=arr,
n_components=1,
dtype=arr.dtype,
)
# Non-axisymmetric storage: expect native (nphi, n_n). The axisymmetry
# check above rejects 1D/(1,n_n) arrays once a non-axisymmetric mode is
# established.
if value.ndim == 2 and value.shape[0] == self.mesh.nphi:
arr = value
else:
squeezed = np.squeeze(value)
if squeezed.ndim == 2 and squeezed.shape[0] == self.mesh.nphi:
arr = squeezed
else:
raise ValueError(f"Unsupported neutral field shape {value.shape}")
return MeshData(
self.mesh,
data_array=arr,
n_components=1,
dtype=arr.dtype,
mesh_is_axisym=getattr(self.mesh, "is_axisymmetric", False),
)
def _infer_data_axisymmetry(self, value):
"""
Infer neutrals-data axisymmetry from array shape.
Parameters
----------
value : np.ndarray
Raw neutrals field array.
Returns
-------
bool
``True`` for axisymmetric neutrals data (``(n_n,)`` or ``(1, n_n)``),
``False`` for non-axisymmetric neutrals data (``(nphi, n_n)``).
Raises
------
ValueError
If the shape does not match any supported neutrals layout.
"""
arr = np.asarray(value)
squeezed = np.squeeze(arr)
if squeezed.ndim == 1:
return True
if squeezed.ndim == 2 and squeezed.shape[0] == self.mesh.nphi:
return False
# Allow unsqueezed axisymmetric shape (1, n_n)
if arr.ndim == 2 and arr.shape[0] == 1:
return True
raise ValueError(f"Cannot infer neutral-data axisymmetry from shape {arr.shape}")
# --- Typed accessors ---
[docs]
def get_neutral_field(self, var_name, step_index=0):
"""Return the raw stored neutral field item for ``var_name``."""
return self.get_item(var_name, step_index)
[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))