"""Helpers for reading static catalog products."""
from __future__ import annotations
from pathlib import Path
from time import perf_counter
from typing import Dict, Iterable, Mapping, Sequence
from .read_plan_executor import SourceReader, execute_read_plans
from .types import MissingStepPolicy
StaticProductBuffer = Dict[str, Dict[str, object]]
[docs]
def build_static_buffer(
catalog,
requests: Mapping[str, Sequence[str]],
*,
source_reader: SourceReader | None = None,
missing: str | MissingStepPolicy = MissingStepPolicy.RAISE,
logical_step: int = 0,
optional_products: Iterable[str] | None = None,
return_timing: bool = False,
) -> StaticProductBuffer | tuple[StaticProductBuffer, dict[str, object]]:
"""
Read several static products into a reusable in-memory buffer.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog containing the requested products.
requests : mapping[str, sequence[str]]
Product-to-variable mapping describing the union of static values to
prefetch.
source_reader : callable or None, optional
Optional static-product backend with the standard ``SourceReader``
signature.
missing : {"raise", "skip", "zero"} or MissingStepPolicy, optional
Missing-variable policy used when a requested product is present but a
requested variable is absent. ``"zero"`` is treated like ``"skip"``
because static readers cannot infer a product-specific fill value.
logical_step : int, optional
Logical step used for static products. Directory-backed static products
normally expose logical step ``0``.
optional_products : iterable[str] or None, optional
Product keys that may be absent from the catalog without failing the
buffer build.
return_timing : bool, optional
If True, return ``(buffer, timing)`` where ``timing`` contains a small
breakdown of plan building, campaign/local reads, and in-memory buffer
materialization time.
Returns
-------
dict[str, dict[str, object]]
Nested mapping keyed by product key and variable name.
"""
missing_policy = MissingStepPolicy(missing)
total_start = perf_counter()
plan_start = total_start
optional_product_keys = set(optional_products or ())
source_reader = source_reader or getattr(catalog, "source_reader", None)
plan_list = []
for product_key, variables in requests.items():
if product_key not in getattr(catalog, "products", {}):
if product_key in optional_product_keys:
continue
raise KeyError(f"Catalog product '{product_key}' not found.")
product_missing_policy = (
MissingStepPolicy.SKIP
if product_key in optional_product_keys
else missing_policy
)
plan_list.extend(
_build_static_plans(
catalog,
product_key,
variables,
missing_policy=product_missing_policy,
logical_step=logical_step,
)
)
plan_seconds = perf_counter() - plan_start
if not plan_list:
timing = {
"plan_s": plan_seconds,
"read_s": 0.0,
"materialize_s": 0.0,
"total_s": perf_counter() - total_start,
"planned_variable_reads": 0,
"record_count": 0,
"product_count": 0,
}
if return_timing:
return {}, timing
return {}
read_start = perf_counter()
execution = execute_read_plans(plan_list, source_reader=source_reader)
read_seconds = perf_counter() - read_start
materialize_start = perf_counter()
record_index = _index_static_records(execution.records)
buffer: StaticProductBuffer = {}
for plan in plan_list:
product_values = buffer.setdefault(plan.product_key, {})
for fragment in plan.fragments:
_validate_static_fragment(fragment, plan)
for logical_value, adios_value in zip(fragment.logical_steps, fragment.adios_steps):
key = _static_record_key(
fragment.source_id,
fragment.source_path,
fragment.file_index,
plan.variable,
logical_value,
adios_value,
)
records = record_index.get(key, [])
if not records:
continue
product_values[plan.variable] = records.pop(0).value
materialize_seconds = perf_counter() - materialize_start
buffer = {product_key: values for product_key, values in buffer.items() if values}
timing = {
"plan_s": plan_seconds,
"read_s": read_seconds,
"materialize_s": materialize_seconds,
"total_s": perf_counter() - total_start,
"planned_variable_reads": len(plan_list),
"record_count": len(execution.records),
"product_count": len(buffer),
}
if return_timing:
return buffer, timing
return buffer
[docs]
def read_static_variables(
catalog,
product_key: str,
variables: Sequence[str],
*,
source_reader: SourceReader | None = None,
missing: str | MissingStepPolicy = MissingStepPolicy.RAISE,
logical_step: int = 0,
static_buffer: Mapping[str, Mapping[str, object]] | None = None,
) -> Mapping[str, object]:
"""
Read variables from one static catalog product.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog containing ``product_key``.
product_key : str
Static product key such as ``"xgc.mesh.bp"`` or ``"xgc.f0.mesh.bp"``.
variables : sequence[str]
ADIOS variable names to read.
source_reader : callable or None, optional
Optional read backend with the standard ``SourceReader`` signature.
``None`` uses ``catalog.source_reader`` when present, otherwise the
FileReader-backed regular BP backend.
missing : {"raise", "skip", "zero"} or MissingStepPolicy, optional
Missing-variable or missing-step policy. ``"zero"`` is treated like
``"skip"`` here because static readers do not have enough context to
synthesize product-specific zero values.
logical_step : int, optional
Logical step used for static products. Directory-backed static
products normally expose logical step ``0``.
static_buffer : mapping[str, mapping[str, object]] or None, optional
Optional reusable buffer previously built with
:func:`build_static_buffer`. Variables already present in the buffer are
returned directly, while missing variables fall back to the catalog
read path.
Returns
-------
dict[str, object]
Values keyed by variable name. Missing values are omitted when
``missing`` is ``"skip"`` or ``"zero"``.
Raises
------
KeyError
If the product or a required variable/step is unavailable.
"""
missing_policy = MissingStepPolicy(missing)
buffered_values, missing_variables = _split_buffered_variables(
static_buffer,
product_key,
variables,
missing_policy,
)
if not missing_variables:
return buffered_values
source_reader = source_reader or getattr(catalog, "source_reader", None)
plans = _build_static_plans(
catalog,
product_key,
missing_variables,
missing_policy=missing_policy,
logical_step=logical_step,
)
execution = execute_read_plans(plans, source_reader=source_reader)
read_values = {record.variable: record.value for record in execution.records}
buffered_values.update(read_values)
return buffered_values
[docs]
def static_product_source_path(catalog, product_key: str) -> Path:
"""
Return the selected path for a static catalog product.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog containing ``product_key``.
product_key : str
Static product key to resolve.
Returns
-------
pathlib.Path
Path for the newest source advertised by the product.
Notes
-----
This is a compatibility helper for code paths that still require a local
BP path, such as lazy sparse-matrix wrappers. New readers should prefer
:func:`read_static_variables` so the same code can later use a
campaign-backed ``SourceReader``.
"""
product = catalog.get_product(product_key)
if not product.sources:
raise KeyError(f"Catalog product '{product_key}' has no sources.")
return max(product.sources, key=lambda source: (source.mtime, str(source.path))).path
[docs]
def has_static_product(catalog, product_key: str, variables: Iterable[str] | None = None) -> bool:
"""
Return whether ``catalog`` advertises a static product and optional variables.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog
Catalog to inspect.
product_key : str
Product key to check.
variables : iterable[str] or None, optional
Variables that must be advertised by the product. If omitted, only the
product itself is checked.
"""
if catalog is None or product_key not in getattr(catalog, "products", {}):
return False
if variables is None:
return True
product = catalog.get_product(product_key)
return all(variable in product.variables for variable in variables)
def _build_static_plans(
catalog,
product_key: str,
variables: Sequence[str],
*,
missing_policy: MissingStepPolicy,
logical_step: int,
):
"""Build single-variable static read plans for one product."""
product = catalog.get_product(product_key)
plan_missing = (
MissingStepPolicy.SKIP
if missing_policy == MissingStepPolicy.ZERO
else missing_policy
)
plans = []
for variable in dict.fromkeys(variables):
if variable not in product.variables:
if missing_policy == MissingStepPolicy.RAISE:
raise KeyError(
f"Variable '{variable}' not found in catalog product '{product_key}'."
)
continue
plans.append(
catalog.plan_read(product_key, variable, [int(logical_step)], missing=plan_missing)
)
return plans
def _split_buffered_variables(static_buffer, product_key, variables, missing_policy):
"""Split a read request into buffered and still-missing variables."""
if not static_buffer or product_key not in static_buffer:
return {}, list(dict.fromkeys(variables))
product_values = static_buffer[product_key]
buffered = {}
missing_variables = []
for variable in dict.fromkeys(variables):
if variable in product_values:
buffered[variable] = product_values[variable]
continue
if missing_policy == MissingStepPolicy.RAISE:
missing_variables.append(variable)
continue
missing_variables.append(variable)
return buffered, missing_variables
def _index_static_records(records):
"""Return records grouped by their source/variable/step coordinates."""
record_index = {}
for record in records:
key = _static_record_key(
record.source_id,
record.source_path,
record.file_index,
record.variable,
record.logical_step,
record.adios_step,
)
record_index.setdefault(key, []).append(record)
return record_index
def _static_record_key(source_id, source_path, file_index, variable, logical_step, adios_step):
"""Build a stable dictionary key for one static read record."""
return (
source_id,
source_path,
file_index,
variable,
int(logical_step),
int(adios_step),
)
def _validate_static_fragment(fragment, plan):
"""Ensure a static plan fragment has consistent logical and ADIOS steps."""
if len(fragment.logical_steps) != len(fragment.adios_steps):
raise ValueError(
f"Read fragment for {plan.product_key}:{plan.variable} has "
"mismatched logical_steps and adios_steps lengths."
)