"""Reusable ADIOS FileReader helpers for finite BP and ACA reads."""
from __future__ import annotations
from pathlib import Path
from typing import Mapping, Sequence
import adios2
[docs]
def read_adios_file_steps(
source_path: Path,
variables: Sequence[str],
adios_steps: Sequence[int],
*,
source_id: str | None = None,
file_reader=None,
qualify_with_source_id: bool = False,
) -> Mapping[int, Mapping[str, object]]:
"""
Read selected variables and ADIOS steps through ``adios2.FileReader``.
Parameters
----------
source_path : pathlib.Path
ADIOS-readable BP source. This path is opened when ``file_reader`` is
omitted.
variables : sequence[str]
Reader-facing variable names to read at each selected step.
adios_steps : sequence[int]
Physical ADIOS step ids to read.
source_id : str or None, optional
Source id used to qualify campaign variables, for example
``xgc.f3d.00010.bp/e_den``.
file_reader : adios2.FileReader or None, optional
Existing open FileReader handle. Supplying this avoids reopening
expensive campaign files for every read batch.
qualify_with_source_id : bool, optional
If True, qualify unqualified variable names as ``source_id/variable``.
Returns
-------
dict[int, dict[str, object]]
Values keyed by ADIOS step and original reader-facing variable name.
"""
requested_steps = sorted({int(step) for step in adios_steps})
if not requested_steps:
return {}
if file_reader is not None:
return _read_steps_from_open_file_reader(
file_reader,
variables,
requested_steps,
source_id=source_id,
qualify_with_source_id=qualify_with_source_id,
)
with adios2.FileReader(str(Path(source_path))) as reader:
return _read_steps_from_open_file_reader(
reader,
variables,
requested_steps,
source_id=source_id,
qualify_with_source_id=qualify_with_source_id,
)
[docs]
def available_step_count(reader) -> int:
"""
Return the maximum ``AvailableStepsCount`` advertised by an open FileReader.
Parameters
----------
reader : adios2.FileReader
Open FileReader handle.
"""
vars_info = reader.available_variables() or {}
return max([available_steps_count_for_variable(info) for info in vars_info.values()] or [0])
[docs]
def available_steps_count_for_variable(info) -> int:
"""
Return ``AvailableStepsCount`` from one ADIOS variable metadata record.
Parameters
----------
info : dict
Raw metadata for one ADIOS variable.
"""
try:
return int(info.get("AvailableStepsCount", 0) or 0)
except Exception:
return 0
[docs]
def available_variable_names(reader) -> list[str]:
"""
Return sorted variable names advertised by an open FileReader.
Parameters
----------
reader : adios2.FileReader
Open FileReader handle.
"""
return sorted((reader.available_variables() or {}).keys())
[docs]
class AdiosFileSourceReader:
"""
Callable source-reader backend for finite catalog reads.
The same callable shape is used for directory-backed local BP products and
campaign-backed ACA products. Directory catalogs construct this class
without an open reader, so each physical BP source is opened once per
batched read request. Campaign catalogs pass a persistent FileReader handle
and enable source-id qualification, avoiding repeated ACA opens.
"""
def __init__(self, file_reader=None, *, qualify_with_source_id: bool = False, owns_reader: bool = False):
"""
Store FileReader backend configuration.
Parameters
----------
file_reader : adios2.FileReader or None, optional
Persistent reader handle. If omitted, ``source_path`` is opened and
closed on each call.
qualify_with_source_id : bool, optional
If True, unqualified variable names are read as
``source_id/variable``.
owns_reader : bool, optional
Whether :meth:`close` should close ``file_reader``.
"""
self.file_reader = file_reader
self.qualify_with_source_id = bool(qualify_with_source_id)
self.owns_reader = bool(owns_reader)
def __call__(
self,
source_path: Path,
variables: Sequence[str],
adios_steps: Sequence[int],
*,
source_id: str | None = None,
) -> Mapping[int, Mapping[str, object]]:
"""
Read one catalog source batch.
Parameters
----------
source_path : pathlib.Path
BP source path, or the campaign path when ``file_reader`` is
persistent.
variables : sequence[str]
Reader-facing variable names.
adios_steps : sequence[int]
Physical ADIOS step ids.
source_id : str or None, optional
Catalog source id used for campaign qualification.
"""
return read_adios_file_steps(
source_path,
variables,
adios_steps,
source_id=source_id,
file_reader=self.file_reader,
qualify_with_source_id=self.qualify_with_source_id,
)
[docs]
def close(self) -> None:
"""Close the persistent FileReader handle when this reader owns it."""
if self.file_reader is None or not self.owns_reader:
return
close = getattr(self.file_reader, "close", None)
if close is not None:
close()
self.file_reader = None
[docs]
@staticmethod
def qualified_variable_name(source_id: str | None, variable: str) -> str:
"""
Return ``source_id/variable`` unless the name is already qualified.
Parameters
----------
source_id : str or None
Catalog source id.
variable : str
Reader-facing variable name.
"""
if source_id is None or "/" in variable:
return variable
return f"{source_id}/{variable}"
def _read_steps_from_open_file_reader(
reader,
variables: Sequence[str],
requested_steps: Sequence[int],
*,
source_id: str | None,
qualify_with_source_id: bool,
) -> Mapping[int, Mapping[str, object]]:
"""
Read selected steps from an already open FileReader.
Parameters
----------
reader : adios2.FileReader
Open FileReader handle.
variables : sequence[str]
Reader-facing variable names.
requested_steps : sequence[int]
ADIOS step ids to read.
source_id : str or None
Catalog source id used when qualification is enabled.
qualify_with_source_id : bool
Whether to qualify unqualified variables with ``source_id``.
"""
available_info = reader.available_variables() or {}
fallback_step_count = available_step_count(reader)
out = {int(step): {} for step in requested_steps}
for variable in variables:
read_name = (
AdiosFileSourceReader.qualified_variable_name(source_id, variable)
if qualify_with_source_id
else variable
)
_validate_read_request(
available_info,
read_name,
variable,
requested_steps,
fallback_step_count,
)
for adios_step in requested_steps:
value = reader.read(
read_name,
step_selection=[int(adios_step), 1],
)
out[int(adios_step)][variable] = _normalize_scalar_value(value, available_info[read_name])
return out
def _validate_read_request(
available_info,
read_name: str,
requested_name: str,
requested_steps: Sequence[int],
fallback_step_count: int,
) -> None:
"""
Validate one variable read before calling ADIOS ``read``.
Parameters
----------
available_info : dict
Variable metadata returned by ``FileReader.available_variables()``.
read_name : str
Actual ADIOS variable name, possibly source-qualified for campaigns.
requested_name : str
Reader-facing variable name used in returned dictionaries and errors.
requested_steps : sequence[int]
ADIOS step ids requested from the source.
fallback_step_count : int
Source-level step count used only when the per-variable metadata does
not advertise ``AvailableStepsCount``.
"""
if read_name not in available_info:
raise KeyError(f"Variable '{requested_name}' not found in ADIOS source as '{read_name}'.")
step_count = available_steps_count_for_variable(available_info[read_name]) or int(fallback_step_count or 0)
missing_steps = [int(step) for step in requested_steps if int(step) < 0 or int(step) >= step_count]
if missing_steps:
raise KeyError(
f"ADIOS step(s) unavailable for variable '{requested_name}' "
f"as '{read_name}': {', '.join(str(step) for step in missing_steps)}"
)
def _normalize_scalar_value(value, info):
"""
Convert ADIOS scalar reads to scalar values while preserving arrays.
Parameters
----------
value : object
Value returned by ``FileReader.read``.
info : dict
ADIOS metadata for the variable.
"""
if str(info.get("SingleValue", "")).lower() != "true":
return value
try:
return value.item()
except AttributeError:
return value