"""HPC-Campaign-backed catalog discovery for XGC ADIOS outputs."""
from __future__ import annotations
from dataclasses import replace
from pathlib import Path
import re
from typing import Dict, Iterable, List, Tuple
import adios2
import numpy as np
from .campaign_source_reader import CampaignSourceReader
from .directory_catalog import (
SimulationCatalog,
classify_product,
_layout_for_sources,
_merge_variables,
_variables_from_adios_info,
)
from .types import DataProduct, SourceInfo, VariableInfo
_BP_SEQ_RE = re.compile(r"^(?P<base>.+)\.(?P<index>\d+)\.bp$")
_LOGICAL_STEP_VARIABLE_NAMES = ("gstep", "step", "timestep")
_CAMPAIGN_SOURCE_ALIASES = {
"2d": "xgc.2d.bp",
"3d": "xgc.3d.bp",
"bfield": "xgc.bfield.bp",
"bfieldm": "xgc.bfieldm.bp",
"cnv_from_surf": "xgc.cnv_from_surf.bp",
"cnv_to_surf": "xgc.cnv_to_surf.bp",
"equil": "xgc.equil.bp",
"f0": "xgc.f0.bp",
"f0.mesh": "xgc.f0.mesh.bp",
"f2d": "xgc.f2d.bp",
"f3d": "xgc.f3d.bp",
"fsourcediag": "xgc.fsourcediag.bp",
"grad_rz": "xgc.grad_rz.bp",
"heatdiag2": "xgc.heatdiag2.bp",
"hyp_vis_rad": "xgc.hyp_vis_rad.bp",
"loading.0": "xgc.loading.0.bp",
"loading.1": "xgc.loading.1.bp",
"loop_vol": "xgc.loop_vol.bp",
"mesh": "xgc.mesh.bp",
"neutrals": "xgc.neutrals.bp",
"oneddiag": "xgc.oneddiag.bp",
"saved_volumes": "xgc.saved_volumes.bp",
"sheathdiag": "xgc.sheathdiag.bp",
"smooth_pol": "xgc.smooth_pol.bp",
"units": "xgc.units.bp",
"volumes": "xgc.volumes.bp",
}
[docs]
def open_campaign_catalog(campaign_path: str | Path, *, collect_metadata: bool = True) -> SimulationCatalog:
"""
Open an HPC-Campaign-backed XGC catalog.
Parameters
----------
campaign_path : str or pathlib.Path
Path to an ADIOS-readable campaign ``.aca`` file.
collect_metadata : bool, optional
Whether to inspect campaign variable metadata. The initial campaign
backend requires metadata collection to build useful product records.
Returns
-------
xgc_analysis.catalog.SimulationCatalog
Catalog populated from campaign-qualified ADIOS variable names.
"""
campaign_path = Path(campaign_path)
campaign_reader = adios2.FileReader(str(campaign_path))
products = discover_campaign_products(
campaign_path,
collect_metadata=collect_metadata,
file_reader=campaign_reader,
)
def refresh():
return discover_campaign_products(
campaign_path,
collect_metadata=collect_metadata,
file_reader=campaign_reader,
)
catalog = SimulationCatalog(
campaign_path.parent,
products,
refresh_func=refresh,
state_key_func=lambda: campaign_state_key(campaign_path),
source_reader=CampaignSourceReader(),
)
catalog.campaign_path = campaign_path
catalog.campaign_reader = campaign_reader
return catalog
[docs]
def discover_campaign_products(
campaign_path: Path,
*,
collect_metadata: bool = True,
file_reader=None,
) -> List[DataProduct]:
"""
Discover XGC products advertised by a campaign ``.aca`` file.
Parameters
----------
campaign_path : pathlib.Path
Campaign file to inspect.
collect_metadata : bool, optional
If False, return no products. Campaign discovery currently depends on
the campaign's ADIOS metadata because there is no filesystem listing to
scan.
file_reader : adios2.FileReader or None, optional
Existing open campaign reader. Supplying this avoids reopening the ACA
file during catalog construction and refresh.
"""
if not collect_metadata:
return []
campaign_path = Path(campaign_path)
if file_reader is not None:
return _discover_campaign_products_from_reader(campaign_path, file_reader)
with adios2.FileReader(str(campaign_path)) as reader:
return _discover_campaign_products_from_reader(campaign_path, reader)
def _discover_campaign_products_from_reader(campaign_path: Path, reader) -> List[DataProduct]:
"""
Build campaign product records from an open FileReader.
Parameters
----------
campaign_path : pathlib.Path
Campaign file path used for source records and state metadata.
reader : adios2.FileReader
Open campaign FileReader handle.
"""
vars_info = reader.available_variables() or {}
attrs_info = reader.available_attributes() or {}
grouped = _group_campaign_metadata(vars_info, attrs_info)
sources_by_product = {}
for source_id, source_vars, source_attrs in grouped:
product_key, file_index = _campaign_product_key(source_id)
variables = _campaign_variables_from_info(source_vars, source_attrs)
adios_step_count = _campaign_step_count(source_vars)
step_values = []
time_values = []
# ADIOS 2.12.1 can crash when scalar step variables are read from
# campaign time-series aliases. Literal BP-backed campaign sources keep
# the richer logical-step metadata; aliases fall back to ADIOS ordinals.
if source_id.endswith(".bp"):
step_values = _read_campaign_scalar_steps(
reader,
source_id,
source_vars,
_LOGICAL_STEP_VARIABLE_NAMES,
adios_step_count,
int,
)
time_values = _read_campaign_scalar_steps(
reader,
source_id,
source_vars,
("time",),
adios_step_count,
float,
)
try:
mtime = campaign_path.stat().st_mtime
except OSError:
mtime = 0.0
sources_by_product.setdefault(product_key, []).append(
SourceInfo(
source_id=source_id,
path=campaign_path,
file_index=file_index,
mtime=mtime,
variables=variables,
step_values=step_values,
time_values=time_values,
adios_step_count=adios_step_count,
)
)
products = []
for product_key in sorted(sources_by_product):
sources = sorted(
sources_by_product[product_key],
key=lambda source: (-1 if source.file_index is None else source.file_index, source.source_id),
)
product_type, family = classify_product(product_key)
products.append(
DataProduct(
key=product_key,
label=product_key,
product_type=product_type,
layout=_layout_for_sources(sources),
sources=sources,
variables=_merge_variables(sources),
product_family=family,
)
)
return products
[docs]
def campaign_state_key(campaign_path: Path):
"""
Return a cheap state key for one campaign file.
Parameters
----------
campaign_path : pathlib.Path
Campaign file path.
"""
campaign_path = Path(campaign_path)
try:
stat = campaign_path.stat()
except OSError:
return tuple()
return ((campaign_path.name, "file", int(stat.st_mtime_ns), int(stat.st_size)),)
def _group_campaign_metadata(vars_info, attrs_info) -> Iterable[Tuple[str, Dict[str, dict], Dict[str, dict]]]:
"""
Group campaign-qualified metadata by source id.
Campaign variables are expected to use names like
``xgc.f2d.00010.bp/e_den`` or time-series aliases such as ``f2d/e_den``.
The returned variable and attribute mappings are stripped to reader-facing
names such as ``e_den`` and ``e_den/unit``.
"""
grouped_vars: Dict[str, Dict[str, dict]] = {}
grouped_attrs: Dict[str, Dict[str, dict]] = {}
for qualified_name, info in (vars_info or {}).items():
if "/" not in qualified_name:
continue
source_id, variable = qualified_name.split("/", 1)
if not _is_supported_campaign_source(source_id):
continue
if not variable:
continue
grouped_vars.setdefault(source_id, {})[variable] = info
for qualified_name, info in (attrs_info or {}).items():
if "/" not in qualified_name:
continue
source_id, attr_name = qualified_name.split("/", 1)
if not _is_supported_campaign_source(source_id):
continue
if not attr_name:
continue
grouped_attrs.setdefault(source_id, {})[attr_name] = info
for source_id in sorted(grouped_vars):
yield source_id, grouped_vars[source_id], grouped_attrs.get(source_id, {})
def _is_supported_campaign_source(source_id: str) -> bool:
"""
Return whether a campaign source prefix maps to an XGC data product.
Parameters
----------
source_id : str
Campaign source prefix before the first ``/`` in an ADIOS variable name.
Returns
-------
bool
True for literal BP names and known HPC-Campaign representation aliases.
"""
return source_id.endswith(".bp") or source_id in _CAMPAIGN_SOURCE_ALIASES
def _campaign_product_key(source_id: str) -> Tuple[str, int | None]:
"""
Return logical product key and optional file index for one campaign source id.
"""
if source_id in _CAMPAIGN_SOURCE_ALIASES:
return _CAMPAIGN_SOURCE_ALIASES[source_id], None
match = _BP_SEQ_RE.match(source_id)
if match:
return f"{match.group('base')}.bp", int(match.group("index"))
return source_id, None
def _campaign_variables_from_info(vars_info, attrs_info) -> Dict[str, VariableInfo]:
"""
Convert campaign variable metadata to unqualified ``VariableInfo`` records.
"""
variables = _variables_from_adios_info(vars_info, attrs_info)
for name, info in vars_info.items():
variables[name] = replace(
variables[name],
step_count=_available_steps_count(info),
)
return variables
def _campaign_step_count(vars_info) -> int:
"""
Return the maximum ``AvailableStepsCount`` advertised by a source.
"""
return max([_available_steps_count(info) for info in vars_info.values()] or [1])
def _available_steps_count(info) -> int:
"""
Parse ADIOS ``AvailableStepsCount`` metadata.
"""
try:
return int(info.get("AvailableStepsCount", 0) or 0)
except Exception:
return 0
def _read_campaign_scalar_steps(reader, source_id, vars_info, names, step_count, value_type):
"""
Read scalar coordinate values for one campaign source.
Parameters
----------
reader : adios2.FileReader
Open campaign reader.
source_id : str
Campaign source prefix.
vars_info : dict
Unqualified variable metadata for the source.
names : iterable[str]
Candidate unqualified variable names.
step_count : int
Number of ADIOS steps to inspect.
value_type : type
Conversion type for returned scalar values.
"""
for name in names:
info = vars_info.get(name)
if info is None:
continue
if str(info.get("SingleValue", "")).lower() not in {"true", "1", "yes"}:
continue
qualified_name = f"{source_id}/{name}"
values = []
for step in range(int(step_count or 0)):
try:
value = reader.read(qualified_name, step_selection=[step, 1])
except Exception:
return []
try:
values.append(value_type(np.asarray(value).reshape(-1)[0]))
except Exception:
return []
return values
return []