"""HPC-Campaign-backed catalog discovery for XGC ADIOS outputs."""
from __future__ import annotations
from dataclasses import replace
from pathlib import Path
import re
import sqlite3
from typing import Dict, Iterable, List, Tuple
import zlib
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 .run_manifest import TextArtifact
from .types import DataProduct, SourceInfo, VariableInfo
_BP_SEQ_RE = re.compile(r"^(?P<base>.+)\.(?P<index>\d+)\.bp$")
_XGC_TIMESERIES_RE = re.compile(r"^xgc\.[A-Za-z0-9_.-]+$")
_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,
persistent_reads: 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.
persistent_reads : bool, optional
If True, reuse the catalog's open campaign ``FileReader`` for data
reads. This is the default because opening campaign archives can be
expensive, especially for remote datasets. Set False only as a
troubleshooting fallback; read-plan batching still prevents one open
per variable in that mode.
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=getattr(source_reader, "file_reader", None) or campaign_reader,
)
if persistent_reads:
def update_campaign_reader(reader):
catalog.campaign_reader = reader
source_reader = CampaignSourceReader(campaign_reader, owns_reader=True, on_reopen=update_campaign_reader)
else:
source_reader = CampaignSourceReader()
def close_campaign_reader():
close = getattr(campaign_reader, "close", None)
if close is not None:
close()
source_reader.close = close_campaign_reader
catalog = SimulationCatalog(
campaign_path.parent,
products,
refresh_func=refresh,
state_key_func=lambda: campaign_state_key(campaign_path),
source_reader=source_reader,
manifest=_campaign_text_manifest(campaign_path),
text_reader=lambda artifact_path: _read_campaign_text_with_lock(
campaign_path,
artifact_path,
source_reader,
),
)
catalog.campaign_path = campaign_path
catalog.campaign_reader = campaign_reader
return catalog
def _read_campaign_text_with_lock(campaign_path: Path, artifact_path: str, source_reader) -> str:
"""
Read campaign text through the persistent reader guarded by source lock.
Parameters
----------
campaign_path : pathlib.Path
Campaign ``.aca`` file.
artifact_path : str
Text dataset name.
source_reader : CampaignSourceReader
Source reader whose lock serializes access to its persistent campaign
FileReader.
"""
lock = getattr(source_reader, "_lock", None)
campaign_reader = getattr(source_reader, "file_reader", None)
if lock is None:
return read_campaign_text(campaign_path, artifact_path, file_reader=campaign_reader)
with lock:
campaign_reader = getattr(source_reader, "file_reader", None)
return read_campaign_text(campaign_path, artifact_path, file_reader=campaign_reader)
[docs]
def read_campaign_text(campaign_path: str | Path, artifact_path: str, *, file_reader=None) -> str:
"""
Read one embedded text artifact from an HPC-Campaign archive.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
artifact_path : str
Dataset name of the embedded text artifact, for example ``"input"`` or
``"input_dir/profile.dat"``.
Returns
-------
str
UTF-8 decoded text content.
Raises
------
KeyError
If the campaign archive does not contain the named text dataset.
RuntimeError
If the archive cannot be queried as an HPC-Campaign SQLite database.
"""
row = _campaign_text_row(campaign_path, artifact_path)
if row is not None:
compression, data = row
if data is not None:
raw = bytes(data)
if int(compression or 0) == 1:
raw = zlib.decompress(raw)
return raw.decode("utf-8")
return _read_campaign_text_variable(campaign_path, artifact_path, file_reader=file_reader)
def _campaign_text_manifest(campaign_path: Path) -> dict:
"""
Build a minimal manifest-like text-artifact section for a campaign archive.
Parameters
----------
campaign_path : pathlib.Path
Campaign ``.aca`` file to inspect.
"""
artifacts = [
TextArtifact(
path=name,
artifact_class=_campaign_text_class(name),
role=_campaign_text_role(name),
source="campaign_text_dataset",
).to_manifest()
for name in _campaign_text_names(campaign_path)
]
return {"text_artifacts": artifacts}
def _campaign_text_names(campaign_path: str | Path) -> List[str]:
"""
Return embedded text dataset names from an HPC-Campaign archive.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
"""
try:
with sqlite3.connect(Path(campaign_path)) as connection:
rows = connection.execute(
"""
select distinct dataset.name
from dataset
where dataset.fileformat = ? and ifnull(dataset.deltime, 0) = 0
order by dataset.name
""",
("TEXT",),
).fetchall()
except sqlite3.Error:
return _campaign_text_variable_names(campaign_path)
return [str(row[0]) for row in rows]
def _campaign_text_row(campaign_path: str | Path, artifact_path: str):
"""
Return the newest embedded text blob row for one campaign dataset.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
artifact_path : str
Text dataset name.
"""
try:
with sqlite3.connect(Path(campaign_path)) as connection:
return connection.execute(
"""
select file.compression, file.data
from dataset
join replica on dataset.rowid = replica.datasetid
join repfiles on replica.rowid = repfiles.replicaid
join file on repfiles.fileid = file.fileid
where dataset.name = ? and dataset.fileformat = ?
order by replica.modtime desc, file.modtime desc
limit 1
""",
(artifact_path, "TEXT"),
).fetchone()
except sqlite3.Error as exc:
raise RuntimeError(f"Could not query campaign text artifacts in '{campaign_path}'.") from exc
def _campaign_text_variable_names(campaign_path: str | Path) -> List[str]:
"""
Return ADIOS-readable text variable names from a campaign archive.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
"""
try:
with adios2.FileReader(str(campaign_path)) as reader:
vars_info = reader.available_variables() or {}
except Exception:
return []
return sorted(
name for name, info in vars_info.items()
if str(info.get("Type", "")).lower() in {"string", "char", "unsigned char"}
or name in {"input", "fort.input.used", "adios2cfg.xml", "petsc.rc"}
or name.startswith("input_dir/")
)
def _read_campaign_text_variable(campaign_path: str | Path, artifact_path: str, *, file_reader=None) -> str:
"""
Read one text dataset exposed as an ADIOS variable by HPC-Campaign.
Parameters
----------
campaign_path : str or pathlib.Path
Campaign ``.aca`` file.
artifact_path : str
ADIOS variable name for the text dataset.
"""
try:
if file_reader is not None:
vars_info = file_reader.available_variables() or {}
if artifact_path not in vars_info:
raise KeyError(f"Unknown campaign text artifact: {artifact_path}")
value = file_reader.read(artifact_path)
else:
with adios2.FileReader(str(campaign_path)) as reader:
vars_info = reader.available_variables() or {}
if artifact_path not in vars_info:
raise KeyError(f"Unknown campaign text artifact: {artifact_path}")
value = reader.read(artifact_path)
except KeyError:
raise
except Exception as exc:
raise RuntimeError(
f"Could not read campaign text artifact '{artifact_path}' from '{campaign_path}'."
) from exc
return _decode_campaign_text_value(value)
def _decode_campaign_text_value(value) -> str:
"""
Decode a text payload returned by ADIOS into a Python string.
Parameters
----------
value : object
Value returned by ``adios2.FileReader.read`` for a text dataset.
"""
if isinstance(value, str):
return value
if isinstance(value, (bytes, bytearray)):
return bytes(value).decode("utf-8")
arr = np.asarray(value)
if arr.dtype.kind in {"S", "U"}:
if arr.ndim == 0:
return str(arr.item())
return "".join(str(item) for item in arr.reshape(-1))
if arr.dtype.kind in {"i", "u"}:
return bytes(int(item) for item in arr.reshape(-1)).decode("utf-8")
raise TypeError(f"Cannot decode campaign text value with dtype={arr.dtype} shape={arr.shape}")
def _campaign_text_class(name: str) -> str:
"""
Return a coarse artifact class for one campaign text dataset name.
Parameters
----------
name : str
Campaign text dataset name.
"""
if name == "input":
return "input.parameter"
if name in {"fort.input.used", "adios2cfg.xml", "petsc.rc"}:
return "input.config"
if name.endswith((".out", ".log")):
return "log.application"
return "input.support"
def _campaign_text_role(name: str) -> str:
"""
Return an optional role for one campaign text dataset name.
Parameters
----------
name : str
Campaign text dataset name.
"""
if name == "input":
return "xgc_input"
if name == "adios2cfg.xml":
return "adios_config"
return ""
[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 = []
tindex_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,
)
tindex_values = _read_campaign_scalar_steps(
reader,
source_id,
source_vars,
("tindex",),
adios_step_count,
int,
)
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,
tindex_values=tindex_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 ``xgc.f2d/e_den``
and legacy 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, XGC time-series names such as ``xgc.2d``,
and known legacy HPC-Campaign representation aliases.
"""
return (
source_id.endswith(".bp")
or _XGC_TIMESERIES_RE.match(source_id) is not None
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"))
if source_id.endswith(".bp"):
return source_id, None
if _XGC_TIMESERIES_RE.match(source_id):
return f"{source_id}.bp", None
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 []