import os
import numpy as np
from .accessor_mixin import ArrayAccessorMixin
from .bp_reader_mixin import BPReaderMixin
from .plane_data import PlaneData
[docs]
class FsourceData(BPReaderMixin, ArrayAccessorMixin):
"""
Reader for ``xgc.fsourcediag.XXXXX.bp`` source diagnostic files.
Data are stored in ``self.data`` using the common XGC-Analysis pattern::
self.data[var_name][step_index] = PlaneData | scalar
``step_index`` is a sequential internal index over loaded ``(file_index,
bp_step)`` pairs. Use ``get_step_info(step_index)`` to recover the source
file index and ADIOS step id.
Mesh-based variables are wrapped as :class:`PlaneData` on ``mesh.get_plane(0)``.
Scalar variables (for example ``time``) are stored as raw NumPy/Python values.
Notes
-----
This diagnostic is currently treated as poloidal-plane data (axisymmetric
storage on a single plane). For a 3D version, a ``MeshData`` branch can be
added following ``FieldData`` / ``FMomentData``.
"""
def __init__(
self,
mesh,
work_dir="./",
file_indices=None,
variables=None,
read_all_steps=False,
catalog=None,
steps=None,
missing="raise",
source_reader=None,
):
"""
Initialize the Fsource diagnostic reader.
Parameters
----------
mesh : Mesh
Mesh object used to provide plane geometry for :class:`PlaneData`.
work_dir : str, optional
Directory containing ``xgc.fsourcediag.XXXXX.bp`` files.
file_indices : list[int], optional
List of file indices to read. If omitted, reads ``[0]``.
variables : list[str] | str | None, optional
Optional explicit variable names to read. When provided, each
requested variable must exist in the first loaded ADIOS step.
If a requested variable is missing in later steps, zeros are
substituted using the first-seen shape/dtype. If ``None``, all
variables in each step are read.
read_all_steps : bool, optional
If ``True``, load every ADIOS step in each BP file. If ``False``
(default), load only the last ADIOS step in each file.
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.
Returns
-------
None
Reader state is initialized and requested files are loaded into
``self.data``.
"""
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]
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.data = {}
if self.catalog is None:
raise RuntimeError("FsourceData requires a catalog; direct xgc.fsourcediag filename reads are disabled.")
self._read_fsource_from_catalog()
def _read_file_steps(self, fname):
"""
Read file steps using Fsource reader variable selection rules.
Parameters
----------
fname : str
Full path to one ``xgc.fsourcediag`` 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.requested_vars)
def _catalog_variables(self):
"""
Return variables to request from the catalog.
Returns
-------
list[str]
Explicitly requested variables, or all variables advertised by
``xgc.fsourcediag.bp`` when no explicit selection was supplied.
"""
if self.requested_vars is not None:
return list(self.requested_vars)
return sorted(self.catalog.get_product("xgc.fsourcediag.bp").variables)
def _read_fsource_from_catalog(self):
"""
Read ``xgc.fsourcediag.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.fsourcediag.bp",
self._catalog_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 variable, value in self._iter_step_items(
variables,
self.step_index_info[step_index]["source_path"],
self.step_index_info[step_index]["bp_step"],
):
self._store_fsource_value(step_index, variable, value)
def _store_fsource_value(self, step_index, variable, value):
"""
Store one fsource variable using the legacy wrapping policy.
Parameters
----------
step_index : int
Reader-local step index.
variable : str
Variable name.
value : object
Raw ADIOS value.
"""
if variable == "time":
self.data.setdefault(variable, {})[step_index] = self._scalar_value(value)
return
if isinstance(value, np.ndarray):
if value.ndim == 1:
value = PlaneData(
plane=self.mesh.get_plane(0),
data_array=value,
n_components=1,
dtype=value.dtype,
)
self.data.setdefault(variable, {})[step_index] = value
return
print(
f"Warning: Variable '{variable}' at step {step_index} "
"has unexpected non-array type. Storing as raw value."
)
self.data.setdefault(variable, {})[step_index] = value
def _read_fsource(self):
"""
Read ``xgc.fsourcediag.XXXXX.bp`` files into ``self.data``.
Parameters
----------
None
Returns
-------
None
Populates ``self.data`` and ``self.step_index_info`` in place.
Each requested file index is read independently. Variables with 1D mesh
data are wrapped as :class:`PlaneData` using ``mesh.get_plane(0)``.
"""
for idx in self.file_indices:
fname = os.path.join(self.work_dir, f"xgc.fsourcediag.{idx:05d}.bp")
fsource_data = self._read_file_steps(fname)
for bp_step in sorted(fsource_data.keys()):
variables = fsource_data[bp_step]
step_index = self._register_step(idx, bp_step)
for variable, value in self._iter_step_items(variables, fname, bp_step):
self._store_fsource_value(step_index, variable, value)
# --- Typed accessors ---
[docs]
def get_source(self, var_name, step_index=0):
"""Return the raw stored source-diagnostic 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_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))