import os
import warnings
import numpy as np
from matplotlib.tri import Triangulation
from .accessor_mixin import ArrayAccessorMixin
from .bp_reader_mixin import BPReaderMixin
from .mesh_data import MeshData
from .plane_data import PlaneData
[docs]
class FMomentData(BPReaderMixin, ArrayAccessorMixin):
"""
Reader for XGC ``f2d`` / ``f3d`` moment diagnostics.
Data are stored in ``self.data`` with the common XGC-Analysis layout::
self.data[var_name][step_index] = PlaneData | MeshData | 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.
``FMomentData`` inherits :class:`ArrayAccessorMixin`, so ``get_array(var)``
is available for converting stored values to plain NumPy arrays.
"""
def __init__(
self,
mesh,
work_dir=".",
file_indices=None,
is_axisymmetric=False,
split_n0_turb=False,
variables=None,
read_all_steps=False,
catalog=None,
steps=None,
missing="raise",
source_reader=None,
):
"""
Initialize the FMoment diagnostic reader.
Parameters
----------
mesh : Mesh
Mesh object providing plane/mesh geometry for ``PlaneData`` and
``MeshData`` wrappers.
work_dir : str, optional
Directory containing ``xgc.f2d.XXXXX.bp`` or ``xgc.f3d.XXXXX.bp``.
file_indices : list[int], optional
List of file indices to read. If omitted, reads ``[0]``.
is_axisymmetric : bool, optional
If ``True``, read ``xgc.f2d`` files and store plane-based data.
If ``False``, read ``xgc.f3d`` files and store mesh-based data where
available.
split_n0_turb : bool, optional
Reserved flag for future post-processing; currently stored but not used.
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``. If omitted and
``file_indices`` is provided, ``file_indices`` are treated as
logical steps for the catalog path. If omitted and
``read_all_steps`` is True, all available catalog steps are read.
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.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.mesh = mesh
self.is_axisymmetric = is_axisymmetric
self.split_n0_turb = split_n0_turb
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 is_axisymmetric:
if self.catalog is None:
raise RuntimeError("FMomentData requires a catalog; direct xgc.f2d filename reads are disabled.")
self._read_f2d_from_catalog(mesh)
else:
if self.catalog is None:
raise RuntimeError("FMomentData requires a catalog; direct xgc.f3d filename reads are disabled.")
self._read_f3d_from_catalog(mesh)
def _read_file_steps(self, fname):
"""
Read file steps using FMoment reader variable selection rules.
Parameters
----------
fname : str
Full path to one ``xgc.f2d`` or ``xgc.f3d`` 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_product_key(self):
"""
Return the catalog product key for this moment reader.
Returns
-------
str
``"xgc.f2d.bp"`` for axisymmetric moments, otherwise
``"xgc.f3d.bp"``.
"""
return "xgc.f2d.bp" if self.is_axisymmetric else "xgc.f3d.bp"
def _catalog_variables(self):
"""
Return variables to request from the catalog.
Returns
-------
list[str]
Explicitly requested variables, or all variables advertised by the
catalog product when no explicit selection was supplied.
"""
if self.requested_vars is not None:
return list(self.requested_vars)
product = self.catalog.get_product(self._catalog_product_key())
return sorted(product.variables)
def _read_f2d_from_catalog(self, mesh):
"""
Read ``xgc.f2d.bp`` through catalog read plans.
Parameters
----------
mesh : Mesh
Mesh object used to build :class:`PlaneData` wrappers.
"""
default_steps = self.file_indices if self._file_indices_provided else None
step_variables = self._read_catalog_product(
self.catalog,
self._catalog_product_key(),
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_f2d_value(mesh, step_index, variable, value)
def _read_f3d_from_catalog(self, mesh):
"""
Read ``xgc.f3d.bp`` through catalog read plans.
Parameters
----------
mesh : Mesh
Mesh object used to build :class:`MeshData` / :class:`PlaneData`
wrappers.
"""
default_steps = self.file_indices if self._file_indices_provided else None
step_variables = self._read_catalog_product(
self.catalog,
self._catalog_product_key(),
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_f3d_value(mesh, step_index, variable, value)
def _store_f2d_value(self, mesh, step_index, variable, value):
"""
Store one f2d value using the legacy wrapping policy.
Parameters
----------
mesh : Mesh
Mesh object used to construct :class:`PlaneData`.
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 variable == "phi":
self.data.setdefault(variable, {})[step_index] = np.squeeze(value)
return
if isinstance(value, np.ndarray) and value.ndim == 1:
value = PlaneData(
plane=mesh.get_plane(0),
data_array=value,
n_components=1,
dtype=value.dtype,
)
self.data.setdefault(variable, {})[step_index] = value
def _store_f3d_value(self, mesh, step_index, variable, value):
"""
Store one f3d value using the legacy wrapping policy.
Parameters
----------
mesh : Mesh
Mesh object used to construct :class:`MeshData` or
:class:`PlaneData`.
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 variable == "phi":
self.data.setdefault(variable, {})[step_index] = np.squeeze(value)
return
if isinstance(value, np.ndarray):
if value.ndim == 2 and value.shape[0] == 1:
value = MeshData(
mesh,
data_array=value[0],
n_components=1,
dtype=value.dtype,
mesh_is_axisym=True,
)
elif value.ndim == 1:
value = PlaneData(
plane=mesh.get_plane(0),
data_array=value,
n_components=1,
dtype=value.dtype,
)
self.data.setdefault(variable, {})[step_index] = value
def _read_f2d_files(self, mesh):
"""
Read ``xgc.f2d.XXXXX.bp`` files and store values in ``self.data``.
Parameters
----------
mesh : Mesh
Mesh object used to build :class:`PlaneData` wrappers.
Returns
-------
None
Populates ``self.data`` and ``self.step_index_info`` in place.
Mesh fields are wrapped as :class:`PlaneData`. Scalars (e.g. ``time``)
and non-mesh arrays (e.g. ``phi``) are stored as raw NumPy arrays.
"""
for idx in self.file_indices:
fname = os.path.join(self.work_dir, f"xgc.f2d.{idx:05d}.bp")
file_data = self._read_file_steps(fname)
for bp_step in sorted(file_data.keys()):
variables = file_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_f2d_value(mesh, step_index, variable, value)
def _read_f3d_files(self, mesh):
"""
Read ``xgc.f3d.XXXXX.bp`` files and store values in ``self.data``.
Parameters
----------
mesh : Mesh
Mesh object used to build :class:`MeshData` / :class:`PlaneData`
wrappers.
Returns
-------
None
Populates ``self.data`` and ``self.step_index_info`` in place.
Variables shaped like ``(1, n_n)`` are treated as axisymmetric mesh
fields and wrapped as :class:`MeshData` with repeated geometry. Pure 1D
arrays are wrapped as :class:`PlaneData`. Scalars (e.g. ``time``) and
special arrays (e.g. ``phi``) are stored directly.
"""
for idx in self.file_indices:
fname = os.path.join(self.work_dir, f"xgc.f3d.{idx:05d}.bp")
file_data = self._read_file_steps(fname)
for bp_step in sorted(file_data.keys()):
variables = file_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_f3d_value(mesh, step_index, variable, value)
[docs]
def export_vtu(self, mesh):
"""
Export loaded FMoment data to VTK-compatible structures.
Parameters
----------
mesh : Mesh
Mesh object providing connectivity and coordinates.
Returns
-------
None
Creates output directory metadata and prepares VTK structures.
"""
import pyvista as pv
vtu_dir = os.path.join(self.work_dir, "vtus")
if not os.path.exists(vtu_dir):
os.makedirs(vtu_dir)
if len(self.file_indices) > 1:
warnings.warn(
"`export_vtu()` is called with multiple file_indices. "
"This will work, but may require a large amount of memory if you plan to generate many VTK files. "
"Consider instantiating separate FMomentData objects for individual VTK file indices.",
UserWarning,
)
tri_obj = Triangulation(mesh.rz[:, 0], mesh.rz[:, 1], mesh.nd_connect_list)
points_vtu = np.column_stack([mesh.rz[:, 0], mesh.rz[:, 1], np.zeros_like(mesh.rz[:, 0])])
# --- Typed accessors ---
[docs]
def get_moment(self, var_name, step_index=0):
"""Return the raw stored moment item for ``var_name`` and ``step_index``."""
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))
# read campaign file
[docs]
def read_from_campaign(self, mesh):
"""
Reads the campaign file and loads the mesh data.
"""