"""
Directory-backed catalog for XGC ADIOS output discovery.
This module scans one dataset root, groups XGC ``*.bp`` outputs into logical
products, collects ADIOS metadata when requested, and builds read plans that map
user-facing logical steps to concrete ADIOS source/step pairs. It does not read
payload arrays; readers execute the returned ``ReadPlan`` objects.
"""
from __future__ import annotations
from collections import defaultdict
from dataclasses import replace
from pathlib import Path
import re
from typing import Callable, Dict, Iterable, List, Optional, Tuple
import numpy as np
from .types import (
CatalogArrayRead,
DataProduct,
MissingStepPolicy,
ProductLayout,
ProductType,
ReadFragment,
ReadPlan,
SourceInfo,
StepFragment,
StepInfo,
VariableInfo,
VariableMetadata,
)
from xgc_analysis.adios_file_reader import AdiosFileSourceReader
_BP_SEQ_RE = re.compile(r"^(?P<base>.+)\.(?P<index>\d+)\.bp$")
_ALL_STEPS_RANGE = range(0, 10**12)
_LOGICAL_STEP_VARIABLE_NAMES = ("gstep", "step", "timestep")
MetadataReader = Callable[[Path], Tuple[Dict[str, VariableInfo], List[int], List[float], int]]
CatalogRefresh = Callable[[], Iterable[DataProduct]]
CatalogStateKey = Tuple[Tuple[str, str, int, int], ...]
CatalogStateReader = Callable[[], CatalogStateKey]
[docs]
class SimulationCatalog:
"""
Catalog of ADIOS-readable XGC products in one dataset root.
The catalog is intentionally independent of ``Simulation`` so GUI code and
automation scripts can inspect available products without constructing mesh
or field objects. A ``Simulation`` may hold a catalog instance when heavier
analysis context is needed.
"""
def __init__(
self,
root_dir: Path,
products: Iterable[DataProduct],
*,
refresh_func: Optional[CatalogRefresh] = None,
state_key_func: Optional[CatalogStateReader] = None,
state_key: Optional[CatalogStateKey] = None,
source_reader=None,
):
"""
Store discovered products for one root directory.
Parameters
----------
root_dir : pathlib.Path
Dataset root used for discovery.
products : iterable of DataProduct
Product entries already discovered by a backend.
refresh_func : callable or None, optional
Backend hook that returns a freshly discovered product iterable.
Directory-backed catalogs install this hook automatically through
:func:`open_catalog`; hand-built catalogs may omit it.
state_key_func : callable or None, optional
Backend hook that returns a cheap signature for the current dataset
state. The directory backend uses BP path names, file types,
modification times, and sizes.
state_key : tuple or None, optional
Initial state signature. If omitted and ``state_key_func`` is
provided, the constructor computes the initial signature.
source_reader : callable or None, optional
Backend hook used by readers to execute catalog read plans. The
directory backend installs a FileReader source reader by default.
"""
self.root_dir = Path(root_dir)
self.products: Dict[str, DataProduct] = {product.key: product for product in products}
self._refresh_func = refresh_func
self._state_key_func = state_key_func
self.state_key = state_key if state_key is not None else self._read_state_key()
self.source_reader = source_reader
[docs]
def close(self) -> None:
"""
Release backend resources held by this catalog.
Directory-backed catalogs normally do not hold open files. Campaign
catalogs may keep an ACA ``FileReader`` open so repeated read-plan
execution does not reopen an expensive archive. If the installed
``source_reader`` exposes ``close()``, this method delegates to it.
"""
close = getattr(self.source_reader, "close", None)
if close is not None:
close()
def __enter__(self):
"""Return this catalog for ``with`` statement use."""
return self
def __exit__(self, exc_type, exc, traceback) -> None:
"""Close backend resources when leaving a ``with`` block."""
self.close()
[docs]
def refresh(self) -> "SimulationCatalog":
"""
Rebuild this catalog from its backend and update it in place.
Directory-backed catalogs created by :func:`open_catalog` rescan the
dataset root, regroup products, reread ADIOS metadata when configured to
do so, and replace ``self.products`` with the refreshed product map.
Existing :class:`ReadPlan` objects are not mutated; callers should
request new read plans after refreshing.
Returns
-------
SimulationCatalog
``self``, for convenient chaining.
Raises
------
RuntimeError
If the catalog was constructed without a refresh backend.
"""
if self._refresh_func is None:
raise RuntimeError("This catalog was constructed without a refresh backend.")
products = self._refresh_func()
self.products = {product.key: product for product in products}
self.state_key = self._read_state_key()
return self
[docs]
def refresh_if_changed(self) -> bool:
"""
Refresh this catalog only if the backend state key has changed.
This is intended for polling running simulations from scripts or the
GUI. The directory backend keeps the state-key computation cheap by
stat-ing BP paths and their contents rather than opening ADIOS files.
When a change is detected, this method calls :meth:`refresh`.
Returns
-------
bool
True if the catalog was refreshed, False if the current state key
matches ``self.state_key``.
Raises
------
RuntimeError
If the catalog was constructed without a state-key backend.
"""
if self._state_key_func is None:
raise RuntimeError("This catalog was constructed without a state-key backend.")
current_state_key = self._read_state_key()
if current_state_key == self.state_key:
return False
self.refresh()
return True
def _read_state_key(self) -> Optional[CatalogStateKey]:
"""Return the backend state key, or ``None`` when no backend provides one."""
if self._state_key_func is None:
return None
return self._state_key_func()
[docs]
def list_products(self) -> List[DataProduct]:
"""
Return all discovered products ordered by product key.
Returns
-------
list[DataProduct]
Stable, alphabetically ordered product list for UI display or
scripted inspection.
"""
return [self.products[key] for key in sorted(self.products)]
[docs]
def get_product(self, product_key: str) -> DataProduct:
"""
Return one product entry by key.
Parameters
----------
product_key : str
Catalog key such as ``xgc.3d.bp`` or ``xgc.oneddiag.bp``.
Raises
------
KeyError
If the product is not present in the catalog.
"""
if product_key not in self.products:
raise KeyError(f"Unknown XGC product: {product_key}")
return self.products[product_key]
[docs]
def list_variables(self, product_key: str) -> List[VariableInfo]:
"""
Return variables available for one product.
Parameters
----------
product_key : str
Catalog key for the product to inspect.
Returns
-------
list[VariableInfo]
Variable metadata ordered by native variable name.
"""
product = self.get_product(product_key)
return [product.variables[key] for key in sorted(product.variables)]
[docs]
def available_steps(self, product_key: str, variable: Optional[str] = None) -> List[StepInfo]:
"""
Return deduplicated logical steps for one product and optional variable.
Parameters
----------
product_key : str
Catalog product key.
variable : str or None, optional
If provided, only sources that advertise the variable contribute
steps. If omitted, all sources for the product contribute.
Returns
-------
list[StepInfo]
Logical steps in ascending order. When several sources advertise
the same logical step, the selected fragment follows the
``newer_source_wins`` policy and other candidates are retained as
duplicates.
"""
product = self.get_product(product_key)
candidates: Dict[int, List[StepFragment]] = defaultdict(list)
for source in product.sources:
if variable is not None and variable not in source.variables:
continue
for logical_step, adios_step, time_value in _iter_source_steps(source):
candidates[int(logical_step)].append(
StepFragment(
source_id=source.source_id,
source_path=source.path,
logical_step=int(logical_step),
adios_step=int(adios_step),
file_index=source.file_index,
time=time_value,
)
)
out = []
for logical_step in sorted(candidates):
fragments = candidates[logical_step]
selected = max(fragments, key=_fragment_newness_key)
duplicates = [
StepFragment(
source_id=fragment.source_id,
source_path=fragment.source_path,
logical_step=fragment.logical_step,
adios_step=fragment.adios_step,
file_index=fragment.file_index,
time=fragment.time,
selected=False,
)
for fragment in fragments
if fragment != selected
]
out.append(
StepInfo(
logical_step=logical_step,
selected_fragment=selected,
duplicate_fragments=duplicates,
)
)
return out
[docs]
def plan_read(
self,
product_key: str,
variable: str,
steps: Iterable[int],
missing: str | MissingStepPolicy = MissingStepPolicy.RAISE,
) -> ReadPlan:
"""
Resolve a variable read over logical steps into ADIOS source fragments.
Parameters
----------
product_key : str
Catalog product key.
variable : str
Native XGC/ADIOS variable name to read.
steps : iterable[int]
Requested logical step values, typically XGC ``gstep`` values or
filename-derived indices when no explicit step variable exists.
missing : {"raise", "skip", "zero"} or MissingStepPolicy, optional
Policy for requested logical steps that are not available.
Returns
-------
ReadPlan
Source fragments grouped by ADIOS-readable source.
Raises
------
KeyError
If a requested step is unavailable and ``missing`` is ``"raise"``.
ValueError
If ``missing`` is not a supported missing-step policy.
"""
missing_policy = MissingStepPolicy(missing)
requested_steps = [int(step) for step in steps]
step_map = {info.logical_step: info for info in self.available_steps(product_key, variable)}
missing_steps = [step for step in requested_steps if step not in step_map]
if missing_steps and missing_policy == MissingStepPolicy.RAISE:
missing_str = ", ".join(str(step) for step in missing_steps)
raise KeyError(f"Missing logical step(s) for {product_key}:{variable}: {missing_str}")
grouped: Dict[Tuple[str, Path, Optional[int]], List[Tuple[int, int]]] = defaultdict(list)
for step in requested_steps:
if step not in step_map:
continue
selected = step_map[step].selected_fragment
grouped[(selected.source_id, selected.source_path, selected.file_index)].append((step, selected.adios_step))
fragments = [
ReadFragment(
source_id=source_id,
source_path=source_path,
variable=variable,
logical_steps=[item[0] for item in items],
adios_steps=[item[1] for item in items],
file_index=file_index,
)
for (source_id, source_path, file_index), items in grouped.items()
]
return ReadPlan(
product_key=product_key,
variable=variable,
requested_steps=requested_steps,
fragments=fragments,
missing_steps=missing_steps,
missing_policy=missing_policy,
)
[docs]
def read_arrays(
self,
product_key: str,
variables: Iterable[str] | str | None = None,
steps: Iterable[int] | int | None = None,
*,
missing: str | MissingStepPolicy = MissingStepPolicy.RAISE,
source_reader=None,
) -> CatalogArrayRead:
"""
Read catalog variables into plain NumPy arrays.
Parameters
----------
product_key : str
Catalog product key such as ``"xgc.2d.bp"``.
variables : iterable[str], str, or None, optional
Variables to read. If omitted, all variables advertised by the
product are read.
steps : iterable[int], int, or None, optional
Logical XGC steps to read. If omitted, each variable is read over
the logical steps available for that variable.
missing : {"raise", "skip", "zero"} or MissingStepPolicy, optional
Missing-step policy for explicit ``steps``. ``"zero"`` records
missing steps but does not synthesize zero arrays.
source_reader : callable or None, optional
Optional read-plan source backend. If omitted,
``self.source_reader`` is used when present, otherwise the regular
FileReader-backed local BP backend is used.
Returns
-------
CatalogArrayRead
Variable-first array mapping plus read provenance.
"""
product = self.get_product(product_key)
variable_list = _normalize_read_array_variables(product, variables)
source_reader = source_reader or self.source_reader
plans = []
for variable in variable_list:
if steps is None:
variable_steps = [info.logical_step for info in self.available_steps(product_key, variable)]
else:
variable_steps = _normalize_step_selector(steps)
plans.append(self.plan_read(product_key, variable, variable_steps, missing=missing))
from .read_plan_executor import execute_read_plans # pylint: disable=import-outside-toplevel
execution = execute_read_plans(plans, source_reader=source_reader)
arrays: Dict[str, Dict[int, np.ndarray]] = defaultdict(dict)
for record in execution.records:
arrays[record.variable][record.logical_step] = np.asarray(record.value)
return CatalogArrayRead(
arrays={variable: dict(step_values) for variable, step_values in arrays.items()},
records=list(execution.records),
missing_steps_by_variable=dict(execution.missing_steps_by_variable),
)
[docs]
def open_catalog(
root_dir: str | Path,
*,
collect_metadata: bool = True,
metadata_reader: Optional[MetadataReader] = None,
) -> SimulationCatalog:
"""
Discover XGC products in a directory and return a catalog.
Parameters
----------
root_dir : str or pathlib.Path
Directory containing XGC output files.
collect_metadata : bool, optional
If True, open each ADIOS source to collect variable and step metadata.
If False, only filenames, product grouping, and basic source metadata
are collected.
metadata_reader : callable or None, optional
Test or backend hook with signature ``reader(path) -> (variables,
step_values, time_values, adios_step_count)``.
Returns
-------
SimulationCatalog
Directory-backed catalog for ``root_dir``. Missing or non-directory
roots produce an empty catalog.
"""
root = Path(root_dir).expanduser().resolve()
products = discover_directory_products(
root,
collect_metadata=collect_metadata,
metadata_reader=metadata_reader,
)
return SimulationCatalog(
root,
products,
refresh_func=lambda: discover_directory_products(
root,
collect_metadata=collect_metadata,
metadata_reader=metadata_reader,
),
state_key_func=lambda: directory_state_key(root),
state_key=directory_state_key(root),
source_reader=AdiosFileSourceReader(),
)
def _normalize_read_array_variables(product: DataProduct, variables) -> List[str]:
"""
Normalize a catalog ``read_arrays`` variable selector.
Parameters
----------
product : DataProduct
Product whose variable namespace should be used for validation.
variables : iterable[str], str, or None
User variable selector.
"""
if variables is None:
return sorted(product.variables)
if isinstance(variables, str):
variable_list = [variables]
else:
variable_list = list(variables)
missing = [variable for variable in variable_list if variable not in product.variables]
if missing:
raise KeyError(f"Variable(s) not found in catalog product '{product.key}': {', '.join(missing)}")
return variable_list
def _normalize_step_selector(steps) -> List[int]:
"""
Normalize an integer or iterable step selector to a list of integers.
Parameters
----------
steps : iterable[int] or int
Logical step selector.
"""
if isinstance(steps, int):
return [int(steps)]
return [int(step) for step in steps]
[docs]
def directory_state_key(root_dir: str | Path) -> CatalogStateKey:
"""
Return a cheap filesystem signature for directory-backed catalog contents.
The signature covers top-level ``*.bp`` files/directories and all entries
beneath BP directories. It intentionally uses only relative path, path
type, modification time, and size, so it can be used for frequent GUI
polling without opening ADIOS metadata. A changed key means callers should
call :meth:`SimulationCatalog.refresh`; an unchanged key means the catalog
can usually be reused.
Parameters
----------
root_dir : str or pathlib.Path
Dataset root scanned by the directory backend.
Returns
-------
tuple
Stable, sorted records ``(relative_path, kind, mtime_ns, size)``. The
tuple is empty when the root does not exist or is not a directory.
"""
root = Path(root_dir).expanduser().resolve()
if not root.exists() or not root.is_dir():
return ()
records = []
for path in sorted(root.iterdir(), key=lambda item: item.name):
if path.name.endswith(".bp"):
records.extend(_path_state_records(root, path))
return tuple(records)
def _path_state_records(root: Path, path: Path) -> List[Tuple[str, str, int, int]]:
"""
Return state-key records for one BP file or directory tree.
Parameters
----------
root : pathlib.Path
Dataset root used to make returned paths relative and stable.
path : pathlib.Path
Top-level BP file or BP directory.
Returns
-------
list[tuple[str, str, int, int]]
Records containing relative POSIX path, path kind (``"dir"`` or
``"file"``), nanosecond modification time, and size. Paths that
disappear during polling are skipped so concurrently updated outputs do
not make the catalog refresh check fail.
"""
paths = [path]
if path.is_dir():
try:
paths.extend(sorted(path.rglob("*"), key=lambda item: item.relative_to(root).as_posix()))
except OSError:
pass
records = []
for item in paths:
try:
stat = item.stat()
except OSError:
continue
kind = "dir" if item.is_dir() else "file"
records.append((item.relative_to(root).as_posix(), kind, int(stat.st_mtime_ns), int(stat.st_size)))
return records
[docs]
def discover_directory_products(
root_dir: Path,
*,
collect_metadata: bool = True,
metadata_reader: Optional[MetadataReader] = None,
) -> List[DataProduct]:
"""
Discover BP products directly contained in ``root_dir``.
Files are grouped using XGC naming conventions. A sequence such as
``xgc.3d.00010.bp`` and ``xgc.3d.00012.bp`` becomes one product keyed as
``xgc.3d.bp`` with two sources. Non-sequence files such as
``xgc.oneddiag.bp`` become single-source products.
Parameters
----------
root_dir : pathlib.Path
Directory to scan. The scan is non-recursive.
collect_metadata : bool, optional
Whether to open ADIOS sources and collect variables/steps.
metadata_reader : callable or None, optional
Optional metadata reader override, mainly for tests.
Returns
-------
list[DataProduct]
Discovered products sorted by product key.
"""
if not root_dir.exists() or not root_dir.is_dir():
return []
entries = [path for path in root_dir.iterdir() if path.name.endswith(".bp")]
grouped: Dict[str, List[Tuple[Optional[int], Path]]] = defaultdict(list)
for path in entries:
match = _BP_SEQ_RE.match(path.name)
if match:
grouped[f"{match.group('base')}.bp"].append((int(match.group("index")), path))
else:
grouped[path.name].append((None, path))
read_metadata = metadata_reader or _read_adios_metadata
products = []
for key in sorted(grouped):
source_items = sorted(grouped[key], key=lambda item: (-1 if item[0] is None else item[0], item[1].name))
sources = [
_build_source_info(
path,
file_index=file_index,
collect_metadata=collect_metadata,
metadata_reader=read_metadata,
)
for file_index, path in source_items
]
product_type, family = classify_product(key)
layout = _layout_for_sources(sources)
variables = _merge_variables(sources)
products.append(
DataProduct(
key=key,
label=key,
product_type=product_type,
layout=layout,
sources=sources,
variables=variables,
product_family=family,
)
)
return products
[docs]
def classify_product(product_key: str) -> Tuple[ProductType, str]:
"""
Classify an XGC product key by established filename conventions.
Parameters
----------
product_key : str
Logical product key after sequence grouping, for example
``xgc.3d.bp`` rather than ``xgc.3d.00010.bp``.
Returns
-------
tuple[ProductType, str]
Product type and coarse product family. Unknown products return
``(ProductType.UNKNOWN, "")`` so callers can still inspect them.
"""
if product_key == "xgc.mesh.bp":
return ProductType.MESH_GEOMETRY, "mesh"
if product_key == "xgc.equil.bp":
return ProductType.EQUILIBRIUM, "mesh"
if product_key in {"xgc.bfield.bp", "xgc.bfieldm.bp", "xgc.current_drive.bp"}:
return ProductType.MAGNETIC_FIELD, "mesh"
if product_key in {
"xgc.f0.mesh.bp",
"xgc.volumes.bp",
"xgc.saved_volumes.bp",
"xgc.grad_rz.bp",
"xgc.smooth_pol.bp",
"xgc.hyp_vis_rad.bp",
"xgc.cnv_to_surf.bp",
"xgc.cnv_from_surf.bp",
"xgc.ff_1dp_fwd.bp",
"xgc.ff_1dp_rev.bp",
"xgc.ff_hdp_fwd.bp",
"xgc.ff_hdp_rev.bp",
}:
return ProductType.MESH_GEOMETRY, "mesh"
if product_key in {"xgc.2d.bp", "xgc.loop_vol.bp"}:
return ProductType.FIELD_2D, "field"
if product_key == "xgc.3d.bp":
return ProductType.FIELD_3D, "field"
if product_key == "xgc.f0.bp":
return ProductType.DISTRIBUTION_FUNCTION, "distribution"
if product_key == "xgc.f2d.bp":
return ProductType.FMOMENT_2D, "moment"
if product_key == "xgc.f3d.bp":
return ProductType.FMOMENT_3D, "moment"
if product_key == "xgc.heatdiag2.bp":
return ProductType.HEAT_DIAG, "diagnostic"
if product_key == "xgc.oneddiag.bp":
return ProductType.ONE_D_DIAG, "diagnostic"
if product_key == "xgc.neutrals.bp":
return ProductType.NEUTRAL_DIAG, "diagnostic"
if product_key == "xgc.fsourcediag.bp":
return ProductType.FSOURCE_DIAG, "diagnostic"
if product_key == "xgc.sheathdiag.bp":
return ProductType.SHEATH_DIAG, "diagnostic"
if product_key == "xgc.diffusion_coeff.bp":
return ProductType.DIFFUSION_COEFFICIENTS, "workflow"
if product_key == "xgc.diffusion_profiles.bp":
return ProductType.DIFFUSION_PROFILES, "workflow"
if product_key == "xgc.units.bp":
return ProductType.ANALYSIS, "analysis"
return ProductType.UNKNOWN, ""
def _build_source_info(
path: Path,
*,
file_index: Optional[int],
collect_metadata: bool,
metadata_reader: MetadataReader,
) -> SourceInfo:
"""
Build source metadata for one physical ADIOS-readable file.
Parameters
----------
path : pathlib.Path
Source path.
file_index : int or None
Sequence index parsed from the filename, if present.
collect_metadata : bool
Whether to call ``metadata_reader``.
metadata_reader : callable
Function used to collect variables, step values, time values, and ADIOS
step count.
Returns
-------
SourceInfo
Source record. Metadata-read failures are intentionally downgraded to
empty metadata so filename discovery remains usable on incomplete or
temporarily unreadable output.
"""
try:
mtime = path.stat().st_mtime
except OSError:
mtime = 0.0
variables: Dict[str, VariableInfo] = {}
step_values: List[int] = []
time_values: List[float] = []
adios_step_count = 1
if collect_metadata:
try:
variables, step_values, time_values, adios_step_count = metadata_reader(path)
except Exception:
variables = {}
step_values = []
time_values = []
adios_step_count = 1
return SourceInfo(
source_id=path.name,
path=path,
file_index=file_index,
mtime=mtime,
variables=variables,
step_values=step_values,
time_values=time_values,
adios_step_count=max(1, int(adios_step_count or 1)),
)
def _read_adios_metadata(path: Path) -> Tuple[Dict[str, VariableInfo], List[int], List[float], int]:
"""
Read lightweight ADIOS metadata from one BP or ACA source.
Parameters
----------
path : pathlib.Path
ADIOS-readable BP source or HPC-Campaign ACA file.
Returns
-------
tuple
``(variables, step_values, time_values, adios_step_count)`` where
``variables`` maps native variable names to ``VariableInfo`` objects,
``step_values`` comes from ``gstep``, ``step``, or ``timestep`` when readable,
``time_values`` comes from ``time`` when readable, and
``adios_step_count`` is observed or inferred from ADIOS metadata.
Notes
-----
Metadata discovery is a finite snapshot operation. ``FileReader`` exposes
the steps available when the source is opened, which matches catalog refresh
semantics for offline and periodically monitored datasets.
"""
import adios2 # pylint: disable=import-outside-toplevel
vars_info = {}
attrs_info = {}
step_values = []
time_values = []
step_count = 0
with adios2.FileReader(str(path)) as reader:
vars_info = reader.available_variables() or {}
attrs_info = reader.available_attributes() or {}
step_count = _max_step_count_from_vars(vars_info)
for step_id in range(int(step_count or 0)):
_append_scalar_step_value(
reader,
vars_info,
_LOGICAL_STEP_VARIABLE_NAMES,
step_values,
int,
step_id,
)
_append_scalar_step_value(reader, vars_info, ("time",), time_values, float, step_id)
return _variables_from_adios_info(vars_info, attrs_info), step_values, time_values, step_count
def _append_scalar_step_value(reader, vars_info, names, out, value_type, step_id):
"""
Append the first readable scalar value from a set of candidate names.
Parameters
----------
reader
Open ADIOS FileReader.
vars_info : dict
Available variable metadata for the current source.
names : iterable[str]
Candidate variable names in priority order.
out : list
Destination list mutated in place.
value_type : type
Conversion type, usually ``int`` for step values or ``float`` for time.
step_id : int
ADIOS step id to inspect.
"""
for name in names:
if name not in vars_info:
continue
try:
value = reader.read(name, step_selection=[int(step_id), 1])
except Exception:
return
try:
import numpy as np # pylint: disable=import-outside-toplevel
out.append(value_type(np.asarray(value).reshape(-1)[0]))
except Exception:
return
return
def _variables_from_adios_info(vars_info, attrs_info) -> Dict[str, VariableInfo]:
"""
Convert raw ADIOS metadata dictionaries into catalog variable records.
Parameters
----------
vars_info : dict
Mapping returned by ``available_variables``.
attrs_info : dict
Mapping returned by ``available_attributes``.
Returns
-------
dict[str, VariableInfo]
Parsed variable metadata keyed by native ADIOS variable name.
"""
attrs_by_var = _attrs_by_variable(attrs_info)
variables = {}
for name, info in (vars_info or {}).items():
shape = str(info.get("Shape", "") or "")
dtype = str(info.get("Type", "") or "")
shape_dims, step_count, is_scalar = _parse_shape_spec(shape)
step_count = _available_steps_count(info) or step_count
variables[name] = VariableInfo(
name=name,
dtype=dtype,
shape=shape,
shape_dims=shape_dims,
is_scalar=is_scalar,
step_count=step_count,
metadata=_semantic_metadata(attrs_by_var.get(name, {})),
)
return variables
def _attrs_by_variable(attrs_info) -> Dict[str, Dict[str, object]]:
"""
Group ADIOS attributes by variable name.
ADIOS commonly exposes variable attributes as keys like
``variable/description``. Attributes without a variable prefix are ignored
by the initial catalog model.
"""
attrs = defaultdict(dict)
for key, info in (attrs_info or {}).items():
if "/" not in key:
continue
var_name, attr_name = key.split("/", 1)
attrs[var_name][attr_name] = info
return dict(attrs)
def _semantic_metadata(attrs: Dict[str, object]) -> VariableMetadata:
"""
Extract XGC semantic metadata fields from one variable's attributes.
Missing fields are returned as empty strings or empty lists so old outputs
without rich metadata remain valid catalog inputs. Both ``unit`` and
``units`` are accepted because existing XGC outputs use the singular form
while the metadata guidance uses the plural field name.
"""
return VariableMetadata(
description=_attr_value(attrs.get("description")),
units=_attr_value(attrs.get("units")) or _attr_value(attrs.get("unit")),
axes=_split_attr_list(attrs.get("axes")),
mesh_context=_attr_value(attrs.get("mesh_context")),
coordinate_context=_attr_value(attrs.get("coordinate_context")),
species_context=_attr_value(attrs.get("species_context")),
time_context=_attr_value(attrs.get("time_context")),
centering=_attr_value(attrs.get("centering")),
normalization=_attr_value(attrs.get("normalization")),
)
def _attr_value(info) -> str:
"""
Return the string value from one ADIOS attribute info dictionary.
Parameters
----------
info : object
Attribute metadata entry. Non-dictionary entries return ``""``.
"""
if not isinstance(info, dict):
return ""
value = info.get("Value", "")
if isinstance(value, (list, tuple)):
return ", ".join(_strip_attr_quotes(str(item)) for item in value)
return _strip_attr_quotes(str(value))
def _strip_attr_quotes(value: str) -> str:
"""
Remove one pair of ADIOS string-attribute quotes if present.
Some ADIOS builds return string attribute values with the literal quote
characters included, for example ``"m^-3"``. Numeric-looking or unquoted
values are returned unchanged.
"""
text = str(value)
if len(text) >= 2 and text[0] == '"' and text[-1] == '"':
return text[1:-1]
return text
def _split_attr_list(info) -> List[str]:
"""
Parse a comma-separated ADIOS attribute value into a list of strings.
Empty or missing attributes return an empty list.
"""
value = _attr_value(info)
if not value:
return []
return [item.strip() for item in value.split(",") if item.strip()]
def _parse_shape_spec(shape_text: str) -> Tuple[List[int], int, bool]:
"""
Parse an ADIOS shape string into payload dimensions and step count.
Parameters
----------
shape_text : str
ADIOS shape text such as ``10*{150}``, ``10*scalar``, or ``{1000, 8}``.
Returns
-------
tuple[list[int], int, bool]
Payload dimensions, leading ADIOS step count if encoded in the shape,
and whether the payload is scalar.
"""
text = str(shape_text or "").strip()
step_count = 0
payload = text
match = re.match(r"^\s*(\d+)\s*\*\s*(.+)\s*$", text)
if match:
step_count = int(match.group(1))
payload = match.group(2).strip()
nums = re.findall(r"\d+", payload or "")
payload_dims = [int(num) for num in nums]
payload_is_scalar = payload == "" or "scalar" in payload.lower()
return payload_dims, step_count, payload_is_scalar
def _max_step_count_from_vars(vars_info) -> int:
"""
Infer the maximum ADIOS step count advertised by variable metadata.
Parameters
----------
vars_info : dict
Raw ADIOS variable metadata.
"""
step_counts = []
for info in vars_info.values():
_dims, step_count, _is_scalar = _parse_shape_spec(str(info.get("Shape", "") or ""))
step_counts.append(_available_steps_count(info) or step_count)
return max(step_counts or [0])
def _available_steps_count(info) -> int:
"""
Return ``AvailableStepsCount`` from an 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
def _layout_for_sources(sources: List[SourceInfo]) -> ProductLayout:
"""
Infer the physical product layout from source and step counts.
Multiple one-step sources are treated as a file sequence. A single
multi-step source is treated as an internal-step product. Multiple sources
with at least one multi-step source are marked mixed.
"""
if len(sources) > 1:
if any(source.adios_step_count > 1 for source in sources):
return ProductLayout.MIXED
return ProductLayout.FILE_SEQUENCE
if sources and sources[0].adios_step_count > 1:
return ProductLayout.INTERNAL_STEPS
return ProductLayout.STATIC
def _merge_variables(sources: List[SourceInfo]) -> Dict[str, VariableInfo]:
"""
Merge source variable metadata using the newer-source-wins policy.
Later sources in the sorted newness order overwrite earlier metadata for
variables with the same name. The returned product-level ``step_count`` is
the number of deduplicated logical steps available for that variable across
all sources, not just the ADIOS-internal step prefix from one source shape.
"""
variable_sources: Dict[str, List[SourceInfo]] = defaultdict(list)
for source in sources:
for name in source.variables:
variable_sources[name].append(source)
variables = {}
for name, sources_for_variable in variable_sources.items():
newest_source = max(sources_for_variable, key=_source_newness_key)
variables[name] = replace(
newest_source.variables[name],
step_count=_deduplicated_step_count(sources_for_variable),
)
return variables
def _deduplicated_step_count(sources: List[SourceInfo]) -> int:
"""
Count unique logical steps advertised by a set of sources.
Parameters
----------
sources : list[SourceInfo]
Sources that contain the variable being merged.
Returns
-------
int
Number of unique logical steps after applying the same logical-step
mapping used by read planning.
"""
logical_steps = set()
for source in sources:
for logical_step, _adios_step, _time_value in _iter_source_steps(source):
logical_steps.add(int(logical_step))
return len(logical_steps)
def _iter_source_steps(source: SourceInfo):
"""
Yield logical/adios/time step triples for one source.
Logical steps prefer explicit ``gstep``, ``step``, or ``timestep`` values in
that order. If no explicit step values are available, file-sequence sources
use the filename index plus ADIOS step offset, and non-sequence sources fall
back to ADIOS step ordinal.
"""
step_count = max(1, int(source.adios_step_count or 1))
for adios_step in range(step_count):
if adios_step < len(source.step_values):
logical_step = source.step_values[adios_step]
elif source.file_index is not None:
logical_step = source.file_index + adios_step
else:
logical_step = adios_step
time_value = source.time_values[adios_step] if adios_step < len(source.time_values) else None
yield int(logical_step), adios_step, time_value
def _source_newness_key(source: SourceInfo):
"""
Return the ordering key used to prefer newer source-level metadata.
Modification time is primary. The filename index and path provide stable
tie-breakers.
"""
file_index = source.file_index if source.file_index is not None else -1
return source.mtime, file_index, str(source.path)
def _fragment_newness_key(fragment: StepFragment):
"""
Return the ordering key used to select duplicate logical-step fragments.
The policy intentionally prefers newer source modification times over older
file-sequence indices so regenerated append-only products can supersede old
one-step files.
"""
file_index = fragment.file_index if fragment.file_index is not None else -1
try:
mtime = fragment.source_path.stat().st_mtime
except OSError:
mtime = 0.0
return mtime, file_index, str(fragment.source_path)