"""Execution helpers for catalog ``ReadPlan`` objects."""
from __future__ import annotations
from dataclasses import dataclass, field
import inspect
from pathlib import Path
from typing import Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
from xgc_analysis.adios_file_reader import read_adios_file_steps
from .types import ReadPlan
[docs]
@dataclass(frozen=True)
class ReadPlanRecord:
"""
One value read from a catalog plan fragment.
A record keeps the logical XGC step next to the physical ADIOS source
coordinates. Reader classes can store values under their own sequential
step indices while retaining enough metadata to map those indices back to
catalog/source coordinates.
Attributes
----------
file_index : int or None
Filename-derived source index when available. Internal-step sources
may leave this unset; readers should use ``logical_step`` and
``source_id`` for unambiguous provenance.
"""
logical_step: int
source_id: str
source_path: Path
adios_step: int
variable: str
value: object
file_index: Optional[int] = None
[docs]
@dataclass(frozen=True)
class ReadPlanExecution:
"""
Values produced by executing one ``ReadPlan``.
Attributes
----------
plan : ReadPlan
Original read plan.
records : list[ReadPlanRecord]
Values read from all plan fragments. Fragment order is preserved.
missing_steps : list[int]
Logical steps that the catalog could not resolve. The executor does
not synthesize zero values; callers decide how to handle these based on
the plan's missing-step policy and variable-specific shape information.
"""
plan: ReadPlan
records: List[ReadPlanRecord] = field(default_factory=list)
missing_steps: List[int] = field(default_factory=list)
[docs]
def records_by_logical_step(self) -> Dict[int, ReadPlanRecord]:
"""
Return read records keyed by logical step.
Later records overwrite earlier records for duplicate logical steps.
Normal plans are already deduplicated by ``SimulationCatalog`` before
execution, so overwriting only matters for hand-built test plans.
"""
return {record.logical_step: record for record in self.records}
[docs]
def values_by_logical_step(self) -> Dict[int, object]:
"""Return read values keyed by logical step."""
return {record.logical_step: record.value for record in self.records}
[docs]
def records_in_requested_order(self) -> List[ReadPlanRecord]:
"""
Return records ordered like ``plan.requested_steps``.
Missing steps are omitted. If ``requested_steps`` contains duplicate
logical steps, the same record appears more than once in the returned
list.
"""
by_step = self.records_by_logical_step()
return [by_step[step] for step in self.plan.requested_steps if step in by_step]
[docs]
@dataclass(frozen=True)
class BatchedReadPlanExecution:
"""
Values produced by executing several ``ReadPlan`` objects together.
The batch executor groups fragments by physical source and opens each source
once for the union of requested variables and ADIOS steps. Records are then
emitted in read-plan order so reader-local step indexing remains stable.
"""
plans: List[ReadPlan]
records: List[ReadPlanRecord] = field(default_factory=list)
missing_steps_by_variable: Dict[str, List[int]] = field(default_factory=dict)
SourceReader = Callable[..., Mapping[int, Mapping[str, object]]]
[docs]
def execute_read_plan(plan: ReadPlan, source_reader: SourceReader | None = None) -> ReadPlanExecution:
"""
Execute a catalog ``ReadPlan`` and return logical-step records.
Parameters
----------
plan : ReadPlan
Single-variable plan returned by ``SimulationCatalog.plan_read``.
source_reader : callable or None, optional
Backend hook with signature ``reader(source_path, variables,
adios_steps)``. If omitted, regular local BP files are read with the
FileReader-backed :func:`read_regular_bp_steps`.
Returns
-------
ReadPlanExecution
Read records and the unresolved logical steps copied from ``plan``.
Raises
------
KeyError
If a plan fragment points to an ADIOS step or variable that cannot be
read from the selected source.
ValueError
If a fragment has mismatched logical-step and ADIOS-step lists.
"""
read_source = source_reader or read_regular_bp_steps
records: List[ReadPlanRecord] = []
for fragment in plan.fragments:
if len(fragment.logical_steps) != len(fragment.adios_steps):
raise ValueError(
f"Read fragment for {fragment.source_id}:{fragment.variable} has "
"mismatched logical_steps and adios_steps lengths."
)
if not fragment.adios_steps:
continue
source_data = _call_source_reader(
read_source,
fragment.source_path,
[fragment.variable],
fragment.adios_steps,
source_id=fragment.source_id,
)
for logical_step, adios_step in zip(fragment.logical_steps, fragment.adios_steps):
if adios_step not in source_data:
raise KeyError(
f"ADIOS step {adios_step} not found while reading "
f"{fragment.source_id}:{fragment.variable}."
)
step_data = source_data[adios_step]
if fragment.variable not in step_data:
raise KeyError(
f"Variable '{fragment.variable}' not found in {fragment.source_id} "
f"at ADIOS step {adios_step}."
)
records.append(
ReadPlanRecord(
logical_step=int(logical_step),
source_id=fragment.source_id,
source_path=fragment.source_path,
adios_step=int(adios_step),
variable=fragment.variable,
value=step_data[fragment.variable],
file_index=fragment.file_index,
)
)
return ReadPlanExecution(plan=plan, records=records, missing_steps=list(plan.missing_steps))
[docs]
def execute_read_plans(
plans: Iterable[ReadPlan],
source_reader: SourceReader | None = None,
) -> BatchedReadPlanExecution:
"""
Execute several single-variable ``ReadPlan`` objects as one batched read.
Parameters
----------
plans : iterable[ReadPlan]
Plans to execute together. Plans are expected to use the same product
but may request different variables.
source_reader : callable or None, optional
Backend hook with signature ``reader(source_path, variables,
adios_steps)``. If omitted, regular local BP files are read with the
FileReader-backed :func:`read_regular_bp_steps`.
Returns
-------
BatchedReadPlanExecution
Read records and missing-step lists keyed by variable name.
Raises
------
KeyError
If a planned ADIOS step or variable cannot be read from its source.
ValueError
If any fragment has mismatched logical-step and ADIOS-step lists.
"""
plan_list = list(plans)
read_source = source_reader or read_regular_bp_steps
source_groups = _group_read_requests_by_source(plan_list)
source_cache = {}
for source_key, requests in source_groups.items():
source_id, source_path, _file_index = source_key
variables = sorted({request["variable"] for request in requests})
adios_steps = sorted({int(request["adios_step"]) for request in requests})
source_cache[source_key] = _call_source_reader(
read_source,
source_path,
variables,
adios_steps,
source_id=source_id,
)
records = []
for plan in plan_list:
for fragment in plan.fragments:
_validate_fragment_lengths(fragment)
source_key = (fragment.source_id, fragment.source_path, fragment.file_index)
source_data = source_cache.get(source_key, {})
for logical_step, adios_step in zip(fragment.logical_steps, fragment.adios_steps):
records.append(
_record_from_source_data(
source_data,
fragment.source_id,
fragment.source_path,
fragment.file_index,
fragment.variable,
logical_step,
adios_step,
)
)
return BatchedReadPlanExecution(
plans=plan_list,
records=records,
missing_steps_by_variable={plan.variable: list(plan.missing_steps) for plan in plan_list},
)
def _group_read_requests_by_source(plans: Sequence[ReadPlan]):
"""
Group plan requests by physical ADIOS-readable source.
Parameters
----------
plans : sequence[ReadPlan]
Read plans to group.
Returns
-------
dict
Keys are ``(source_id, source_path, file_index)``. Values are request
dictionaries containing ``variable`` and ``adios_step``.
"""
grouped: Dict[Tuple[str, Path, Optional[int]], List[dict]] = {}
for plan in plans:
for fragment in plan.fragments:
_validate_fragment_lengths(fragment)
source_key = (fragment.source_id, fragment.source_path, fragment.file_index)
requests = grouped.setdefault(source_key, [])
for adios_step in fragment.adios_steps:
requests.append({"variable": fragment.variable, "adios_step": int(adios_step)})
return grouped
def _call_source_reader(read_source, source_path, variables, adios_steps, *, source_id=None):
"""
Call a source-reader backend with optional source-id context.
Parameters
----------
read_source : callable
Source reader backend. Legacy/local readers may accept only
``(source_path, variables, adios_steps)``. Campaign readers can also
accept ``source_id=...`` to map unqualified reader variable names onto
campaign-qualified variable names.
source_path : pathlib.Path
Physical source path from the read fragment.
variables : sequence[str]
Reader-facing variable names.
adios_steps : sequence[int]
Physical ADIOS step ids.
source_id : str or None, optional
Source id from the read fragment.
"""
if _accepts_source_id(read_source):
return read_source(source_path, variables, adios_steps, source_id=source_id)
return read_source(source_path, variables, adios_steps)
def _accepts_source_id(read_source):
"""
Return whether ``read_source`` accepts a ``source_id`` keyword.
Callable objects and partially wrapped functions are handled through
``inspect.signature``. If inspection fails, fall back to the conservative
legacy three-argument call path.
"""
try:
parameters = inspect.signature(read_source).parameters
except (TypeError, ValueError):
return False
if "source_id" in parameters:
return True
return any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values())
def _validate_fragment_lengths(fragment):
"""
Validate that one read fragment has parallel logical and ADIOS step lists.
Parameters
----------
fragment : ReadFragment
Fragment to validate.
"""
if len(fragment.logical_steps) != len(fragment.adios_steps):
raise ValueError(
f"Read fragment for {fragment.source_id}:{fragment.variable} has "
"mismatched logical_steps and adios_steps lengths."
)
def _record_from_source_data(
source_data,
source_id,
source_path,
file_index,
variable,
logical_step,
adios_step,
):
"""
Build one ``ReadPlanRecord`` from cached source data.
Parameters
----------
source_data : mapping
Source data returned by a ``SourceReader``.
source_id, source_path, file_index
Source coordinates copied from the read fragment.
variable : str
Variable to extract.
logical_step : int
Logical XGC step.
adios_step : int
Physical ADIOS step.
"""
if adios_step not in source_data:
raise KeyError(f"ADIOS step {adios_step} not found while reading {source_id}:{variable}.")
step_data = source_data[adios_step]
if variable not in step_data:
raise KeyError(f"Variable '{variable}' not found in {source_id} at ADIOS step {adios_step}.")
return ReadPlanRecord(
logical_step=int(logical_step),
source_id=source_id,
source_path=source_path,
adios_step=int(adios_step),
variable=variable,
value=step_data[variable],
file_index=file_index,
)
[docs]
def read_regular_bp_steps(
source_path: Path,
variables: Sequence[str],
adios_steps: Sequence[int],
*,
source_id: str | None = None,
) -> Mapping[int, Mapping[str, object]]:
"""
Read selected ADIOS steps from one regular local BP source.
Parameters
----------
source_path : pathlib.Path
ADIOS-readable BP file or directory.
variables : sequence[str]
Variable names to read from each selected ADIOS step.
adios_steps : sequence[int]
Physical ADIOS step ids required by a read-plan fragment.
source_id : str or None, optional
Accepted for source-reader API compatibility. Regular BP paths do not
need it.
Returns
-------
mapping
``{adios_step: {variable_name: value}}`` for the requested step span.
Notes
-----
This is the default finite-read backend for directory-backed catalogs. It
opens the source with ``adios2.FileReader``, reads the exact requested steps
with ``step_selection``, and closes the source when the batch is done.
Campaign-backed catalogs use the same source-reader interface with a
persistent FileReader handle.
"""
return read_adios_file_steps(
source_path,
variables,
adios_steps,
source_id=source_id,
qualify_with_source_id=False,
)