"""Shared helpers for BP-based diagnostic readers."""
import warnings
import numpy as np
from .catalog import execute_read_plans
[docs]
class BPReaderMixin:
"""
Mixin for readers that load ADIOS BP files into ``self.data``.
This mixin handles:
- optional explicit variable selection,
- optional all-step vs last-step file reads,
- first-step variable validation for explicit requests,
- step-index bookkeeping and
- zero-filling of later-step missing requested variables.
"""
_ALL_STEPS_RANGE = (0, 10**9)
def _init_bp_reader_state(self, variables=None, read_all_steps=False):
"""
Initialize shared BP-reader attributes.
Parameters
----------
variables : list[str] | str | None, default=None
Optional explicit variable names requested by the caller. ``None``
means use reader-specific defaults or "all available variables"
depending on the concrete reader.
read_all_steps : bool, default=False
If ``True``, readers load all ADIOS steps from each file.
If ``False``, readers load only the last ADIOS step.
Returns
-------
None
This method mutates ``self`` by setting:
``requested_vars``, ``read_all_steps``, ``step_index_info``,
and ``_requested_var_templates``.
"""
self.requested_vars = self._normalize_variables(variables)
self.read_all_steps = bool(read_all_steps)
self.step_index_info = {}
self._requested_var_templates = {}
@staticmethod
def _normalize_variables(variables):
"""
Normalize ``variables`` input to a deduplicated list of strings.
Parameters
----------
variables : list[str] | str | None
Raw variable selector provided by user-facing reader constructors.
Returns
-------
list[str] | None
``None`` when ``variables`` is ``None``; otherwise a deduplicated
list preserving original order.
Raises
------
ValueError
If ``variables`` is provided but empty.
TypeError
If any entry is not a string.
"""
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 var in variable_list:
if not isinstance(var, str):
raise TypeError("Each entry in `variables` must be a string.")
if var not in normalized:
normalized.append(var)
return normalized
def _register_step(self, file_index, bp_step):
"""
Register one loaded ``(file_index, bp_step)`` pair.
Parameters
----------
file_index : int
File index from the filename pattern (for example ``42`` for
``xgc.3d.00042.bp``).
bp_step : int
ADIOS step id inside that file.
Returns
-------
int
Sequential internal ``step_index`` used for storing values in
``self.data[var_name][step_index]``.
"""
step_index = len(self.step_index_info)
self.step_index_info[step_index] = {
"file_index": int(file_index),
"bp_step": int(bp_step),
}
return step_index
@staticmethod
def _read_plan_record_key(record):
"""
Return the identity key used to combine variables into reader steps.
Parameters
----------
record : xgc_analysis.catalog.ReadPlanRecord
One value read from a catalog read plan.
Returns
-------
tuple
Stable key for one logical/physical source step.
"""
return (
int(record.logical_step),
record.source_id,
str(record.source_path),
int(record.adios_step),
)
def _register_read_plan_record(self, record):
"""
Register the source step represented by one read-plan record.
Parameters
----------
record : xgc_analysis.catalog.ReadPlanRecord
Read record whose source coordinates should become one reader-local
``step_index``.
Returns
-------
int
Sequential reader-local step index.
Notes
-----
``file_index`` is retained for compatibility with existing reader code.
Internal-step sources may not have a filename-derived index; in that
case the logical step is used as a compatibility fallback while the
exact source coordinates remain available through ``logical_step``,
``source_id``, ``source_path``, and ``bp_step``.
"""
step_index = len(self.step_index_info)
file_index = record.file_index if record.file_index is not None else record.logical_step
self.step_index_info[step_index] = {
"file_index": int(file_index),
"bp_step": int(record.adios_step),
"logical_step": int(record.logical_step),
"source_id": record.source_id,
"source_path": str(record.source_path),
}
return step_index
def _all_available_catalog_steps(self, catalog, product_key, variables):
"""
Return all logical steps available for the first advertised variable.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog containing the product.
product_key : str
Catalog product key.
variables : iterable[str]
Candidate variables in reader-preferred order.
Returns
-------
list[int]
Logical steps for the first variable that has available steps.
"""
product = catalog.get_product(product_key)
for variable in variables:
if variable not in product.variables:
continue
steps = catalog.available_steps(product_key, variable)
if steps:
return [step.logical_step for step in steps]
return []
def _resolve_catalog_steps(
self,
catalog,
product_key,
variables,
steps=None,
default_steps=None,
read_all_steps=False,
):
"""
Resolve caller step options into logical catalog steps.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog containing the product.
product_key : str
Catalog product key.
variables : iterable[str]
Variables used to discover all available steps when requested.
steps : iterable[int] or None, optional
Explicit logical steps requested by the caller.
default_steps : iterable[int] or None, optional
Reader-specific fallback steps, often legacy file indices.
read_all_steps : bool, optional
If True and neither ``steps`` nor ``default_steps`` is provided,
return all available logical steps for the product.
Returns
-------
list[int]
Logical steps to use for catalog read planning.
"""
if steps is not None:
return [int(step) for step in steps]
if default_steps is not None:
return [int(step) for step in default_steps]
if read_all_steps:
return self._all_available_catalog_steps(catalog, product_key, variables)
return [0]
def _catalog_read_plans(
self,
catalog,
product_key,
variables,
logical_steps,
missing="raise",
require_all_variables=True,
missing_for_optional_variables="skip",
):
"""
Build single-variable read plans for a catalog product.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog containing the product.
product_key : str
Product to read.
variables : iterable[str]
Variables requested by the concrete reader.
logical_steps : iterable[int]
Logical XGC steps to read.
missing : {"raise", "skip", "zero"}, optional
Missing-step policy used when ``require_all_variables`` is True.
require_all_variables : bool, optional
If True, every variable must be advertised by the product. If
False, unadvertised variables are skipped so concrete readers can
apply default-fill behavior.
missing_for_optional_variables : {"raise", "skip", "zero"}, optional
Missing-step policy used when ``require_all_variables`` is False.
Returns
-------
list[xgc_analysis.catalog.ReadPlan]
Single-variable read plans.
"""
product = catalog.get_product(product_key)
plans = []
for variable in variables:
if variable not in product.variables:
if require_all_variables:
raise KeyError(f"Requested variable '{variable}' not found in catalog product '{product_key}'.")
continue
plan_missing = missing if require_all_variables else missing_for_optional_variables
plans.append(catalog.plan_read(product_key, variable, logical_steps, missing=plan_missing))
return plans
def _read_catalog_product(
self,
catalog,
product_key,
variables,
steps=None,
default_steps=None,
read_all_steps=False,
missing="raise",
require_all_variables=True,
source_reader=None,
):
"""
Plan and read raw values for one catalog product.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog containing the product.
product_key : str
Product key to read.
variables : iterable[str]
Variables requested by the concrete reader.
steps : iterable[int] or None, optional
Explicit logical steps requested by the caller.
default_steps : iterable[int] or None, optional
Reader-specific fallback steps.
read_all_steps : bool, optional
If True and no explicit/default steps are provided, read all
available logical steps for the product.
missing : {"raise", "skip", "zero"}, optional
Missing-step policy for required variables.
require_all_variables : bool, optional
If False, variables absent from product metadata are skipped and can
be filled by the concrete reader's ordered-variable policy.
source_reader : callable or None, optional
Optional backend passed to the batched read-plan executor.
Returns
-------
dict[int, dict[str, object]]
Mapping from reader-local ``step_index`` to raw variable values.
"""
source_reader = source_reader or getattr(catalog, "source_reader", None)
logical_steps = self._resolve_catalog_steps(
catalog,
product_key,
variables,
steps=steps,
default_steps=default_steps,
read_all_steps=read_all_steps,
)
read_plans = self._catalog_read_plans(
catalog,
product_key,
variables,
logical_steps,
missing=missing,
require_all_variables=require_all_variables,
)
return self._read_from_read_plans(read_plans, source_reader=source_reader)
def _read_from_read_plans(self, read_plans, source_reader=None):
"""
Execute catalog read plans as a batch and group values by step index.
Parameters
----------
read_plans : iterable[xgc_analysis.catalog.ReadPlan]
Single-variable read plans. Plans for different variables are
batched by physical source so each source is opened once for the
union of requested variables and ADIOS steps.
source_reader : callable or None, optional
Optional backend hook passed to
:func:`xgc_analysis.catalog.execute_read_plans`. ``None`` uses the
catalog source reader when present, otherwise the FileReader-backed
regular BP backend.
Returns
-------
dict[int, dict[str, object]]
Mapping from reader-local ``step_index`` to raw variable values.
Concrete readers are responsible for product-specific wrapping and
storage in ``self.data[var_name][step_index]``.
"""
step_index_by_record = {}
variables_by_step_index = {}
execution = execute_read_plans(read_plans, source_reader=source_reader)
for record in execution.records:
record_key = self._read_plan_record_key(record)
if record_key not in step_index_by_record:
step_index_by_record[record_key] = self._register_read_plan_record(record)
step_index = step_index_by_record[record_key]
variables_by_step_index.setdefault(step_index, {})[record.variable] = record.value
return variables_by_step_index
def _validate_requested_variables(self, file_data, fname):
"""
Validate explicit variable requests against the first loaded step.
Parameters
----------
file_data : dict[int, dict[str, object]]
Step-major mapping returned by a raw source reader.
fname : str
File path used for error context.
Returns
-------
None
Raises
------
KeyError
If ``self.requested_vars`` is set and any requested variable is
missing in the first loaded ADIOS step.
"""
if self.requested_vars is None:
return
if not file_data:
return
first_step = min(file_data.keys())
first_variables = file_data[first_step]
missing = [var for var in self.requested_vars if var not in first_variables]
if missing:
missing_str = ", ".join(missing)
raise KeyError(
f"Requested variable(s) missing in first step {first_step} of '{fname}': "
f"{missing_str}"
)
@staticmethod
def _zero_like(value):
"""
Build a zero value preserving type/shape information.
Parameters
----------
value : object
Template value used to infer output type and, for arrays, shape.
Returns
-------
object
Zero-like value with matching dtype/shape when possible.
"""
if isinstance(value, np.ndarray):
return np.zeros_like(value)
if isinstance(value, np.generic):
return value.dtype.type(0)
if isinstance(value, bool):
return False
if isinstance(value, (int, float, complex)):
return type(value)(0)
return 0.0
@staticmethod
def _scalar_value(value):
"""
Return a Python scalar for scalar-like ADIOS values.
Parameters
----------
value : object
Value returned by ADIOS or supplied by default-fill policy.
Returns
-------
object
``value.item()`` for scalar or length-one NumPy arrays; otherwise
the original value.
"""
arr = np.asarray(value)
if arr.size == 1:
return arr.reshape(-1)[0].item()
return value
def _read_file_steps(self, fname, read_vars=None):
"""
Disabled legacy direct-file read helper.
Parameters
----------
fname : str
Full file path that would have been read by the legacy path.
read_vars : list[str] | None, default=None
Ignored. Retained only so old private helper call sites fail with
a clear message.
Raises
------
RuntimeError
Always raised. Converted readers must use catalog read plans so the
same implementation works for directory and campaign backends.
"""
raise RuntimeError(
f"Direct BP file read helper is disabled for '{fname}'. "
"Use a SimulationCatalog and catalog read plans instead."
)
def _iter_step_items(self, variables, fname, bp_step, ordered_vars=None, default_value=0.0):
"""
Iterate ``(var, value)`` pairs for one BP step under reader policy.
Parameters
----------
variables : dict[str, object]
Variable dictionary for one ADIOS step.
fname : str
File path used for warning/error context.
bp_step : int
ADIOS step id used for warning/error context.
ordered_vars : list[str] | None, default=None
Ordered variable names to iterate when ``self.requested_vars`` is
``None``. If ``None``, iterate all keys in ``variables``.
default_value : object, default=0.0
Fallback value used when ``ordered_vars`` is provided and a variable
is missing from ``variables``.
Returns
-------
collections.abc.Iterable[tuple[str, object]]
Iterable of ``(variable_name, value)`` pairs.
Raises
------
KeyError
If ``self.requested_vars`` is set, a requested variable is missing,
and no template is available for zero-fill.
Notes
-----
If explicit ``variables`` were requested, missing values in later steps
are filled with zeros using first-seen variable templates and a warning.
"""
if self.requested_vars is None:
if ordered_vars is None:
return variables.items()
return ((var, variables.get(var, default_value)) for var in ordered_vars)
variable_items = []
for var in self.requested_vars:
if var in variables:
value = variables[var]
self._requested_var_templates.setdefault(var, value)
else:
if var not in self._requested_var_templates:
raise KeyError(
f"Variable '{var}' missing in '{fname}' at step {bp_step} and no template is available."
)
warnings.warn(
f"Variable '{var}' missing in '{fname}' at step {bp_step}; filling with zeros.",
UserWarning,
)
value = self._zero_like(self._requested_var_templates[var])
variable_items.append((var, value))
return variable_items
[docs]
def get_step_info(self, step_index):
"""
Return file/ADIOS-step metadata for one internal step index.
Parameters
----------
step_index : int
Internal sequential index used as second key in ``self.data``.
Returns
-------
dict[str, int]
Dictionary with keys:
- ``file_index``: integer file index from filename
- ``bp_step``: integer ADIOS step id inside that file
Raises
------
KeyError
If ``step_index`` is not present in ``self.step_index_info``.
"""
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])