"""Read-only readers for XGC diffusion workflow BP products.
These readers intentionally do not inherit from ``BPReaderMixin`` for now.
Diffusion products have product-specific indexing rules: coefficient files use
species-suffixed variables that may eventually need a project-wide species
handling policy, and profile files contain multiple XGC sample steps inside one
ADIOS step. Keeping the BP handling explicit here avoids baking those rules
into the generic BP-reader mixin before the broader reader/catalog interface is
settled.
"""
from __future__ import annotations
import os
from typing import Dict, Iterable, List, Optional, Tuple
import numpy as np
from .accessor_mixin import ArrayAccessorMixin
from .read_bp_file import ReadBPFile
[docs]
class DiffusionCoefficientData(ArrayAccessorMixin):
"""
Read ``xgc.diffusion_coeff.bp`` as an analysis data product.
This class is intentionally read-only and independent of the existing
workflow-oriented ``DiffusionCoefficients`` class, which opens an append
stream for live coefficient updates. Each ADIOS step is stored as one
reader-local ``step_index``. Species-suffixed coefficient variables are
stacked into arrays with shape ``(n_species, npsi)`` and exposed under the
base coefficient names:
- ``ptl_diffusivity``
- ``momentum_diffusivity``
- ``heat_conductivity``
- ``ptl_pinch_velocity``
Metadata variables such as ``psi``, ``n_species``, and ``gstep`` are stored
in the same ``self.data[var_name][step_index]`` layout.
"""
COEFFICIENT_NAMES = (
"ptl_diffusivity",
"momentum_diffusivity",
"heat_conductivity",
"ptl_pinch_velocity",
)
SPECIES_SUFFIXES = ("_elec", "_ion", "_imp1", "_imp2", "_imp3", "_imp4", "_imp5")
def __init__(
self,
data_dir: str = "./",
filename: str = "xgc.diffusion_coeff.bp",
read_all_steps: bool = True,
):
"""
Initialize the read-only diffusion-coefficient reader.
Parameters
----------
data_dir : str, optional
Directory containing ``filename``.
filename : str, optional
BP product name. Defaults to ``xgc.diffusion_coeff.bp``.
read_all_steps : bool, optional
If True, read all ADIOS steps. If False, read only the last step.
"""
self.data_dir = data_dir
self.filename = filename
self.file_path = os.path.join(self.data_dir, self.filename)
self.read_all_steps = bool(read_all_steps)
self.data: Dict[str, Dict[int, object]] = {}
self.raw_data: Dict[str, Dict[int, object]] = {}
self.step_index_info: Dict[int, Dict[str, object]] = {}
self._read_file()
def _read_file(self):
"""
Read selected ADIOS steps and populate coefficient arrays.
Returns
-------
None
Missing files produce an empty reader. Malformed coefficient files
raise ``KeyError`` or ``ValueError`` with the offending variable
name and step context.
"""
if not os.path.exists(self.file_path):
return
step_range = (0, 10**9) if self.read_all_steps else None
file_data = ReadBPFile(self.file_path, step_range=step_range)
for bp_step in sorted(file_data):
variables = file_data[bp_step]
step_index = self._register_step(bp_step, variables)
self._store_raw_variables(step_index, variables)
self._store_aggregated_coefficients(step_index, bp_step, variables)
def _register_step(self, bp_step: int, variables: Dict[str, object]) -> int:
"""
Register one ADIOS step and return its reader-local step index.
Parameters
----------
bp_step : int
ADIOS step id inside ``xgc.diffusion_coeff.bp``.
variables : dict
Variables read for that ADIOS step.
"""
step_index = len(self.step_index_info)
self.step_index_info[step_index] = {
"bp_step": int(bp_step),
"gstep": _optional_scalar_int(variables.get("gstep")),
"tindex": _optional_scalar_int(variables.get("tindex")),
"time": _optional_scalar_float(variables.get("time")),
}
return step_index
def _store_raw_variables(self, step_index: int, variables: Dict[str, object]):
"""
Store raw variables after scalar/vector normalization.
Parameters
----------
step_index : int
Reader-local step index.
variables : dict
Raw variable dictionary returned by :func:`ReadBPFile`.
"""
for name, value in variables.items():
normalized = _normalize_value(value)
self.raw_data.setdefault(name, {})[step_index] = normalized
if name in {"psi", "n_species", "gstep", "tindex", "time"}:
self.data.setdefault(name, {})[step_index] = normalized
def _store_aggregated_coefficients(self, step_index: int, bp_step: int, variables: Dict[str, object]):
"""
Stack species-suffixed coefficient variables for one ADIOS step.
Parameters
----------
step_index : int
Reader-local step index.
bp_step : int
ADIOS step id used for error messages.
variables : dict
Raw variable dictionary for the ADIOS step.
"""
n_species = int(np.asarray(variables["n_species"]).item())
if n_species > len(self.SPECIES_SUFFIXES):
raise ValueError(
f"xgc.diffusion_coeff.bp step {bp_step} has n_species={n_species}, "
f"but only {len(self.SPECIES_SUFFIXES)} suffixes are known."
)
psi = np.asarray(variables["psi"]).reshape(-1)
npsi = psi.size
suffixes = self.SPECIES_SUFFIXES[:n_species]
for coeff_name in self.COEFFICIENT_NAMES:
coeff = np.zeros((n_species, npsi), dtype=np.float64)
for species_index, suffix in enumerate(suffixes):
var_name = coeff_name + suffix
if var_name not in variables:
raise KeyError(
f"Missing diffusion coefficient variable '{var_name}' in ADIOS step {bp_step}."
)
coeff[species_index, :] = np.asarray(variables[var_name], dtype=np.float64).reshape(-1)
self.data.setdefault(coeff_name, {})[step_index] = coeff
[docs]
def get_step_info(self, step_index: int) -> Dict[str, object]:
"""Return ADIOS-step metadata for one reader-local step index."""
if step_index not in self.step_index_info:
raise KeyError(f"Step index '{step_index}' not found.")
return dict(self.step_index_info[step_index])
[docs]
def get_coefficient(self, name: str, step_index: int = 0) -> np.ndarray:
"""Return one aggregated coefficient array with shape ``(n_species, npsi)``."""
return self.get_as(name, step_index, np.ndarray)
[docs]
def get_scalar(self, name: str, step_index: int = 0):
"""Return one scalar metadata value such as ``n_species`` or ``gstep``."""
return self.get_as(name, step_index, (int, float, np.integer, np.floating))
[docs]
class DiffusionProfileData(ArrayAccessorMixin):
"""
Read ``xgc.diffusion_profiles.bp`` as buffered profile snapshots.
Each ADIOS step contains multiple XGC simulation samples. The profile
variables ``density``, ``flow``, and ``temp`` therefore have shape
``(n_species, n_samples, n_surf)`` for each ADIOS step. This reader keeps
that native layout in ``self.data[var_name][step_index]`` and provides
helpers to address an inner sample by ``sample_index`` or by the XGC sample
step stored in the ``steps`` variable.
"""
PROFILE_VARIABLES = ("density", "flow", "temp")
SAMPLE_STEP_NAMES = ("steps", "sample_steps", "gsteps")
def __init__(
self,
data_dir: str = "./",
filename: str = "xgc.diffusion_profiles.bp",
read_all_steps: bool = True,
):
"""
Initialize the diffusion-profile reader.
Parameters
----------
data_dir : str, optional
Directory containing ``filename``.
filename : str, optional
BP product name. Defaults to ``xgc.diffusion_profiles.bp``.
read_all_steps : bool, optional
If True, read all ADIOS steps. If False, read only the last ADIOS
step. In either case, every buffered sample inside each selected
ADIOS step is retained.
"""
self.data_dir = data_dir
self.filename = filename
self.file_path = os.path.join(self.data_dir, self.filename)
self.read_all_steps = bool(read_all_steps)
self.data: Dict[str, Dict[int, object]] = {}
self.step_index_info: Dict[int, Dict[str, object]] = {}
self.sample_index_by_gstep: Dict[int, Tuple[int, int]] = {}
self._read_file()
def _read_file(self):
"""
Read selected ADIOS steps and index their buffered samples.
Returns
-------
None
Missing files produce an empty reader. Shape inconsistencies in
profile variables raise ``ValueError``.
"""
if not os.path.exists(self.file_path):
return
step_range = (0, 10**9) if self.read_all_steps else None
file_data = ReadBPFile(self.file_path, step_range=step_range)
for bp_step in sorted(file_data):
variables = file_data[bp_step]
step_index = self._register_step(bp_step, variables)
self._store_variables(step_index, variables)
self._index_samples(step_index)
def _register_step(self, bp_step: int, variables: Dict[str, object]) -> int:
"""
Register one ADIOS step and its inner sample coordinates.
Parameters
----------
bp_step : int
ADIOS step id inside ``xgc.diffusion_profiles.bp``.
variables : dict
Variables read for that ADIOS step.
"""
step_index = len(self.step_index_info)
sample_steps = _first_present_vector(variables, self.SAMPLE_STEP_NAMES, dtype=np.int64)
sample_times = _optional_vector(variables.get("time"), dtype=np.float64)
n_samples = _optional_scalar_int(variables.get("n_samples"))
if n_samples is None and sample_steps is not None:
n_samples = int(sample_steps.size)
self.step_index_info[step_index] = {
"bp_step": int(bp_step),
"gstep": _optional_scalar_int(variables.get("gstep")),
"tindex": _optional_scalar_int(variables.get("tindex")),
"n_species": _optional_scalar_int(variables.get("n_species")),
"n_samples": n_samples,
"n_surf": _optional_scalar_int(variables.get("n_surf")),
"sample_steps": sample_steps,
"sample_times": sample_times,
}
return step_index
def _store_variables(self, step_index: int, variables: Dict[str, object]):
"""
Store all variables from one ADIOS step after normalization.
Parameters
----------
step_index : int
Reader-local step index.
variables : dict
Raw variable dictionary returned by :func:`ReadBPFile`.
"""
for name, value in variables.items():
normalized = _normalize_value(value)
self.data.setdefault(name, {})[step_index] = normalized
for name in self.PROFILE_VARIABLES:
if name in self.data and step_index in self.data[name]:
self._validate_profile_shape(name, step_index)
def _validate_profile_shape(self, name: str, step_index: int):
"""
Validate one profile variable shape against ``n_samples`` metadata.
Parameters
----------
name : str
Profile variable name.
step_index : int
Reader-local step index.
"""
arr = np.asarray(self.data[name][step_index])
if arr.ndim != 3:
raise ValueError(
f"Diffusion profile variable '{name}' at step {step_index} has shape {arr.shape}; "
"expected (n_species, n_samples, n_surf)."
)
n_samples = self.step_index_info[step_index].get("n_samples")
if n_samples is not None and arr.shape[1] != int(n_samples):
raise ValueError(
f"Diffusion profile variable '{name}' at step {step_index} has {arr.shape[1]} samples, "
f"but n_samples={n_samples}."
)
def _index_samples(self, step_index: int):
"""
Add buffered sample steps to ``sample_index_by_gstep``.
Later ADIOS steps overwrite earlier entries for duplicate sample step
values, matching the catalog's newer-source preference at the sample
index level.
"""
sample_steps = self.step_index_info[step_index].get("sample_steps")
if sample_steps is None:
return
for sample_index, gstep in enumerate(np.asarray(sample_steps).reshape(-1)):
self.sample_index_by_gstep[int(gstep)] = (step_index, int(sample_index))
[docs]
def available_samples(self) -> List[Dict[str, object]]:
"""
Return all buffered sample coordinates in ADIOS-step order.
Returns
-------
list[dict]
Each entry contains ``step_index``, ``sample_index``, ``gstep`` and
``time``. ``gstep`` or ``time`` may be ``None`` when absent from
the file.
"""
samples = []
for step_index in sorted(self.step_index_info):
info = self.step_index_info[step_index]
sample_steps = info.get("sample_steps")
sample_times = info.get("sample_times")
n_samples = info.get("n_samples") or 0
for sample_index in range(int(n_samples)):
gstep = None
time = None
if sample_steps is not None and sample_index < len(sample_steps):
gstep = int(sample_steps[sample_index])
if sample_times is not None and sample_index < len(sample_times):
time = float(sample_times[sample_index])
samples.append(
{
"step_index": step_index,
"sample_index": sample_index,
"gstep": gstep,
"time": time,
}
)
return samples
[docs]
def get_step_info(self, step_index: int) -> Dict[str, object]:
"""Return ADIOS-step and buffered-sample metadata for one step index."""
if step_index not in self.step_index_info:
raise KeyError(f"Step index '{step_index}' not found.")
return dict(self.step_index_info[step_index])
[docs]
def get_profile(
self,
var_name: str,
step_index: int = 0,
*,
sample_index: Optional[int] = None,
gstep: Optional[int] = None,
) -> np.ndarray:
"""
Return a diffusion profile variable or one buffered sample.
Parameters
----------
var_name : str
Profile variable name, usually ``density``, ``flow``, or ``temp``.
step_index : int, optional
Reader-local ADIOS step index used when ``gstep`` is not provided.
sample_index : int or None, optional
Buffered sample index inside the ADIOS step. If omitted, the full
native array with shape ``(n_species, n_samples, n_surf)`` is
returned.
gstep : int or None, optional
XGC sample step to select from the buffered ``steps`` array. When
provided, it overrides ``step_index`` and ``sample_index``.
Returns
-------
np.ndarray
Full native profile array or one sample with shape
``(n_species, n_surf)``.
"""
if gstep is not None:
if int(gstep) not in self.sample_index_by_gstep:
raise KeyError(f"Buffered diffusion-profile sample gstep={gstep} not found.")
step_index, sample_index = self.sample_index_by_gstep[int(gstep)]
arr = self.get_as(var_name, step_index, np.ndarray)
if sample_index is None:
return arr
return arr[:, int(sample_index), :]
def _normalize_value(value):
"""
Normalize one ADIOS value to a scalar or squeezed NumPy array.
Scalars become Python/NumPy scalar values through ``item()``. Arrays keep
their full dimensionality except for length-one ADIOS step axes that have
already been selected by :func:`ReadBPFile`.
"""
arr = np.asarray(value)
if arr.ndim == 0:
return arr.item()
return np.squeeze(arr)
def _optional_scalar_int(value) -> Optional[int]:
"""Return ``value`` as ``int`` or ``None`` when absent."""
if value is None:
return None
return int(np.asarray(value).reshape(-1)[0])
def _optional_scalar_float(value) -> Optional[float]:
"""Return ``value`` as ``float`` or ``None`` when absent."""
if value is None:
return None
return float(np.asarray(value).reshape(-1)[0])
def _optional_vector(value, *, dtype) -> Optional[np.ndarray]:
"""Return ``value`` as a one-dimensional array or ``None`` when absent."""
if value is None:
return None
return np.asarray(value, dtype=dtype).reshape(-1)
def _first_present_vector(variables: Dict[str, object], names: Iterable[str], *, dtype) -> Optional[np.ndarray]:
"""
Return the first present vector from a list of candidate variable names.
Parameters
----------
variables : dict
Raw variable dictionary for one ADIOS step.
names : iterable[str]
Candidate names in priority order.
dtype
NumPy dtype used for the returned vector.
"""
for name in names:
if name in variables:
return _optional_vector(variables[name], dtype=dtype)
return None