diff --git a/pyproject.toml b/pyproject.toml index a591e8a1..346c6875 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "xarray", "tables", "netCDF4", + "h5netcdf", "h5py", "cad-to-dagmc", "pydagmc @ git+https://github.com/svalinn/pydagmc.git", diff --git a/results_database/convert_legacy_results.py b/results_database/convert_legacy_results.py new file mode 100644 index 00000000..831d793e --- /dev/null +++ b/results_database/convert_legacy_results.py @@ -0,0 +1,343 @@ +#!/usr/bin/env python3 +"""Convert legacy OFB HDF5 results to the current tally group schema. + +Legacy format (per group): +- dataset named like the group, shape (realization, row, column) +- coordinate datasets: realization, row, column +- typical columns: [energy low, energy high, mean, std. dev.] + +New format (per group): +- variables: mean, mc_std +- dims: (surface, energy, nuclide, score) +- attrs: filter_axes, nuclides, scores, observed_tally, optional spec_* fields +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +from pathlib import Path + +import h5py +import numpy as np +import xarray as xr +import yaml + + +def _normalize_filter_type(type_name: str) -> str: + name = str(type_name).strip().lower() + if name.endswith("filter"): + name = name[:-6] + return name + + +def _filter_bins_match(spec_filter: dict, observed_axis: dict) -> bool: + expected = spec_filter.get("values") + observed = observed_axis.get("bins") + if expected is None or observed is None: + return True + + ftype = _normalize_filter_type(spec_filter.get("type", "")) + if ftype == "energy": + try: + return np.allclose(np.asarray(expected, dtype=float), np.asarray(observed, dtype=float)) + except Exception: + return False + + try: + return list(expected) == list(observed) + except Exception: + return False + + +def _validate_tally_consistency(spec_tally: dict, observed_tally: dict) -> tuple[bool, list[str]]: + issues: list[str] = [] + + observed_scores = [str(s) for s in observed_tally.get("scores", [])] + observed_nuclides = [str(n) for n in observed_tally.get("nuclides", [])] + observed_filters = list(observed_tally.get("filters", [])) + + spec_scores = [str(s) for s in spec_tally.get("scores", [])] + if spec_scores and spec_scores != observed_scores: + issues.append(f"scores mismatch: expected {spec_scores}, observed {observed_scores}") + + spec_nuclides = [str(n) for n in spec_tally.get("nuclides", [])] + if spec_nuclides and spec_nuclides != observed_nuclides: + issues.append(f"nuclides mismatch: expected {spec_nuclides}, observed {observed_nuclides}") + + expected_particle = spec_tally.get("particle") + if expected_particle is not None: + particle_filters = [a for a in observed_filters if _normalize_filter_type(a.get("name", "")) == "particle"] + if not particle_filters: + issues.append("missing ParticleFilter in observed tally") + else: + bins = particle_filters[0].get("bins", []) + observed_particle = str(bins[0]) if bins else None + if str(expected_particle) != str(observed_particle): + issues.append( + f"particle mismatch: expected {expected_particle}, observed {observed_particle}" + ) + + spec_filters = spec_tally.get("filters", []) + observed_non_particle = [ + a for a in observed_filters if _normalize_filter_type(a.get("name", "")) != "particle" + ] + + expected_types = [_normalize_filter_type(f.get("type", "")) for f in spec_filters] + observed_types = [_normalize_filter_type(a.get("name", "")) for a in observed_non_particle] + if expected_types != observed_types: + issues.append(f"filter type/order mismatch: expected {expected_types}, observed {observed_types}") + + if len(spec_filters) == len(observed_non_particle): + for i, (spec_filter, observed_axis) in enumerate(zip(spec_filters, observed_non_particle)): + ftype = _normalize_filter_type(spec_filter.get("type", "")) + if ftype == "energy": + expected_closure = spec_filter.get("closure", "[low, high)") + observed_closure = observed_axis.get("closure", "[low, high)") + if expected_closure != observed_closure: + issues.append( + "energy closure mismatch at index " + f"{i}: expected {expected_closure}, observed {observed_closure}" + ) + + if not _filter_bins_match(spec_filter, observed_axis): + issues.append( + f"filter bins mismatch at index {i} ({spec_filter.get('type')}): " + f"expected {spec_filter.get('values')}, observed {observed_axis.get('bins')}" + ) + + return len(issues) == 0, issues + + +def _load_spec_lookup(repo_root: Path, benchmark: str | None) -> dict[str, dict]: + if not benchmark: + return {} + + spec_path = repo_root / "src" / "openmc_fusion_benchmarks" / "benchmarks" / benchmark / "specifications.yaml" + if not spec_path.exists(): + raise FileNotFoundError(f"Could not find specifications file: {spec_path}") + + with spec_path.open("r", encoding="utf-8") as f: + spec = yaml.safe_load(f) + + lookup: dict[str, dict] = {} + for entry in spec.get("tallies", []): + if isinstance(entry, dict) and entry.get("name"): + lookup[str(entry["name"])] = entry + return lookup + + +def _decode_columns(group: h5py.Group) -> list[str]: + cols_raw = group["column"][()] + cols: list[str] = [] + for c in cols_raw: + if isinstance(c, bytes): + cols.append(c.decode("utf-8")) + else: + cols.append(str(c)) + return cols + + +def _col_index(columns: list[str], candidates: list[str]) -> int: + normalized = [c.strip().lower().replace("_", " ") for c in columns] + for cand in candidates: + c = cand.strip().lower().replace("_", " ") + for i, col in enumerate(normalized): + if c == col: + return i + for cand in candidates: + c = cand.strip().lower().replace("_", " ") + for i, col in enumerate(normalized): + if c in col: + return i + raise KeyError(f"Could not find column from candidates {candidates}. Found columns: {columns}") + + +def _legacy_group_to_dataset(group_name: str, arr: np.ndarray, columns: list[str], tally_id: int, spec_tally: dict | None) -> xr.Dataset: + if arr.ndim != 3: + raise ValueError(f"Expected legacy data shape (realization, row, column), got {arr.shape}") + if arr.shape[0] < 1: + raise ValueError(f"Legacy dataset for '{group_name}' has no realization axis entries") + + low_idx = _col_index(columns, ["energy low [ev]", "energy low", "energy_low [ev]", "energy_low"]) + high_idx = _col_index(columns, ["energy high [ev]", "energy high", "energy_high [ev]", "energy_high"]) + mean_idx = _col_index(columns, ["mean"]) + std_idx = _col_index(columns, ["std. dev.", "std dev", "std_dev", "mc_std", "std"]) + + first = arr[0, :, :] + low = np.asarray(first[:, low_idx], dtype=float) + high = np.asarray(first[:, high_idx], dtype=float) + mean = np.asarray(first[:, mean_idx], dtype=float) + mc_std = np.asarray(first[:, std_idx], dtype=float) + + if len(low) == 0: + energy_edges = np.asarray([], dtype=float) + else: + energy_edges = np.concatenate([low[:1], high]) + + particle = None + if isinstance(spec_tally, dict): + particle = spec_tally.get("particle") + if not particle: + low_name = group_name.lower() + if "neutron" in low_name: + particle = "neutron" + elif "photon" in low_name or "gamma" in low_name: + particle = "photon" + + mean_5d = mean.reshape(1, 1, mean.shape[0], 1, 1) + std_5d = mc_std.reshape(1, 1, mc_std.shape[0], 1, 1) + + dims = ("particle", "surface", "energy", "nuclide", "score") + coords = { + "particle": ("particle", np.asarray([0], dtype=int)), + "surface": ("surface", np.asarray([0], dtype=int)), + "energy": ("energy", np.arange(mean.shape[0], dtype=int)), + "nuclide": ("nuclide", np.asarray([0], dtype=int)), + "score": ("score", np.asarray([0], dtype=int)), + } + + ds = xr.Dataset( + { + "mean": xr.DataArray(mean_5d, dims=dims, coords=coords), + "mc_std": xr.DataArray(std_5d, dims=dims, coords=coords), + } + ) + + filter_axes = [ + { + "name": "ParticleFilter", + "axis": "particle", + "num_bins": 1, + "bins": [particle] if particle is not None else [], + }, + { + "name": "SurfaceFilter", + "axis": "surface", + "num_bins": 1, + "bins": [7], + }, + { + "name": "EnergyFilter", + "axis": "energy", + "num_bins": int(mean.shape[0]), + "bins": energy_edges.tolist(), + "kind": "edges", + "units": "eV", + "closure": "[low, high)", + }, + ] + + scores = ["current"] + nuclides = ["total"] + + ds.attrs["filter_axes"] = json.dumps(filter_axes) + ds.attrs["scores"] = json.dumps(scores) + ds.attrs["nuclides"] = json.dumps(nuclides) + + ds.attrs["group"] = group_name + ds.attrs["tally_name"] = group_name + + ds["mean"].attrs["tally_id"] = int(tally_id) + ds["mean"].attrs["tally_name"] = group_name + ds["mean"].attrs["tally_group"] = group_name + + ds["mc_std"].attrs["tally_id"] = int(tally_id) + ds["mc_std"].attrs["tally_name"] = group_name + ds["mc_std"].attrs["tally_group"] = group_name + + observed_tally = { + "name": group_name, + "id": int(tally_id), + "filters": filter_axes, + "scores": scores, + "nuclides": nuclides, + } + ds.attrs["observed_tally"] = json.dumps(observed_tally) + + if spec_tally is not None: + ds.attrs["spec_tally"] = json.dumps(spec_tally) + consistent, issues = _validate_tally_consistency(spec_tally, observed_tally) + ds.attrs["spec_consistent"] = int(bool(consistent)) + ds.attrs["spec_consistency_issues"] = json.dumps(issues) + + return ds + + +def convert_file(input_path: Path, output_path: Path, benchmark: str | None, engine: str) -> Path: + repo_root = Path(__file__).resolve().parents[1] + spec_lookup = _load_spec_lookup(repo_root, benchmark) + + if output_path.exists(): + output_path.unlink() + + with h5py.File(input_path, "r") as src: + for idx, group_name in enumerate(src.keys(), start=1): + group = src[group_name] + if group_name not in group: + raise KeyError(f"Legacy group '{group_name}' missing data dataset '{group_name}'") + if "column" not in group: + raise KeyError(f"Legacy group '{group_name}' missing 'column' dataset") + + arr = np.asarray(group[group_name][()]) + columns = _decode_columns(group) + spec_tally = spec_lookup.get(group_name) + + ds = _legacy_group_to_dataset( + group_name=group_name, + arr=arr, + columns=columns, + tally_id=idx, + spec_tally=spec_tally, + ) + mode = "w" if idx == 1 else "a" + ds.to_netcdf(output_path, mode=mode, group=group_name, engine=engine) + + return output_path.resolve() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Convert legacy OFB result HDF5 format to the current tally schema.") + parser.add_argument("input", type=Path, help="Legacy input .h5 file") + parser.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="Output .h5 file (default: _converted.h5)", + ) + parser.add_argument( + "--benchmark", + type=str, + default=None, + help="Optional benchmark name to attach spec_tally and spec_consistency metadata.", + ) + parser.add_argument( + "--engine", + choices=["auto", "h5netcdf"], + default="auto", + help="NetCDF engine to use (default: auto, which resolves to h5netcdf).", + ) + args = parser.parse_args() + + input_path = args.input.resolve() + if not input_path.exists(): + raise FileNotFoundError(f"Input file not found: {input_path}") + + output_path = args.output.resolve() if args.output else input_path.with_name(f"{input_path.stem}_converted.h5") + engine = _select_engine() if args.engine == "auto" else args.engine + out = convert_file(input_path=input_path, output_path=output_path, benchmark=args.benchmark, engine=engine) + print(f"Converted: {input_path} -> {out}") + + +def _select_engine() -> str: + if importlib.util.find_spec("h5netcdf") is not None: + return "h5netcdf" + raise RuntimeError( + "h5netcdf is required for this converter. Install it with: pip install h5netcdf" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_all_benchmark.py b/scripts/validate_all_benchmark.py index a53ddce1..5717bf41 100644 --- a/scripts/validate_all_benchmark.py +++ b/scripts/validate_all_benchmark.py @@ -1,4 +1,4 @@ -from openmc_fusion_benchmarks import validate_benchmark +from openmc_fusion_benchmarks.validate_spec import validate_benchmark from pathlib import Path benchmarks_dir = Path("src/openmc_fusion_benchmarks/benchmarks") diff --git a/src/openmc_fusion_benchmarks/__init__.py b/src/openmc_fusion_benchmarks/__init__.py index 71c6026b..df130ddf 100644 --- a/src/openmc_fusion_benchmarks/__init__.py +++ b/src/openmc_fusion_benchmarks/__init__.py @@ -1,6 +1,7 @@ from openmc_fusion_benchmarks.benchmark import * from openmc_fusion_benchmarks.benchmark_results import * -from openmc_fusion_benchmarks.validate import * +from openmc_fusion_benchmarks.validate_spec import * +from openmc_fusion_benchmarks.validate_results import * from openmc_fusion_benchmarks.database import * import openmc_fusion_benchmarks.uq diff --git a/src/openmc_fusion_benchmarks/backends/__init__.py b/src/openmc_fusion_benchmarks/backends/__init__.py new file mode 100644 index 00000000..9726a8f6 --- /dev/null +++ b/src/openmc_fusion_benchmarks/backends/__init__.py @@ -0,0 +1 @@ +"""Backend-specific interfaces (OpenMC, Serpent, MCNP, ...).""" diff --git a/src/openmc_fusion_benchmarks/backends/openmc/__init__.py b/src/openmc_fusion_benchmarks/backends/openmc/__init__.py new file mode 100644 index 00000000..4424c95b --- /dev/null +++ b/src/openmc_fusion_benchmarks/backends/openmc/__init__.py @@ -0,0 +1,3 @@ +"""OpenMC backend adapters for openmc_fusion_benchmarks.""" + +from .tallies import * diff --git a/src/openmc_fusion_benchmarks/backends/openmc/tallies.py b/src/openmc_fusion_benchmarks/backends/openmc/tallies.py new file mode 100644 index 00000000..81723aae --- /dev/null +++ b/src/openmc_fusion_benchmarks/backends/openmc/tallies.py @@ -0,0 +1,341 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Callable, Iterable +import json + +import h5py +import numpy as np +import openmc +import xarray as xr + +from ...validate_results import normalize_filter_type, validate_tally_consistency + + +def make_default_openmc_normalizer(mesh: str | Path | object): + """ + Build a standard OpenMC tally normalizer based on geometry measures. + + The returned callback normalizes each tally over its filter axes using: + - cell filter bins -> cell volumes from a DAGMC mesh model + - surface filter bins -> surface areas from a DAGMC mesh model + + Parameters + ---------- + mesh: + Either a mesh path accepted by ``pydagmc.Model`` or an existing + object exposing ``volumes_by_id`` and ``surfaces_by_id`` mappings. + """ + if hasattr(mesh, "volumes_by_id") and hasattr(mesh, "surfaces_by_id"): + msh = mesh + else: + # Local import keeps this backend module importable even if pydagmc + # is not installed in non-OpenMC environments. + import pydagmc + + msh = pydagmc.Model(str(mesh)) + + def normalizer(tally: openmc.Tally, mean_nd: np.ndarray, std_nd: np.ndarray): + filters = list(tally.filters) + + # Filter dimensions are the leading axes in mean_nd/std_nd. + for axis, flt in enumerate(filters): + if isinstance(flt, openmc.CellFilter): + ids = np.asarray(flt.bins, dtype=int).reshape(-1) + factors = np.asarray([msh.volumes_by_id[i].volume for i in ids], dtype=float) + elif isinstance(flt, openmc.SurfaceFilter): + ids = np.asarray(flt.bins, dtype=int).reshape(-1) + factors = np.asarray([msh.surfaces_by_id[i].area for i in ids], dtype=float) + elif isinstance(flt, openmc.MaterialFilter): + raise NotImplementedError( + "Material filter normalization is not implemented yet." + ) + else: + continue + + if np.any(factors == 0.0): + raise ValueError("Normalization factor contains zero values.") + + reshape = (1,) * axis + (factors.shape[0],) + (1,) * (mean_nd.ndim - axis - 1) + norm = factors.reshape(reshape) + mean_nd = mean_nd / norm + std_nd = std_nd / norm + + return mean_nd, std_nd + + return normalizer + + +def _unique_filter_dims(filters: list[openmc.Filter]) -> list[str]: + """Build stable, unique filter dimension names.""" + counts: dict[str, int] = {} + dims: list[str] = [] + for flt in filters: + base = type(flt).__name__.replace("Filter", "").lower() or "filter" + idx = counts.get(base, 0) + counts[base] = idx + 1 + dims.append(base if idx == 0 else f"{base}_{idx}") + return dims + + +def _serialize_filter_bins(flt: openmc.Filter): + """Serialize OpenMC filter bins into JSON-compatible Python objects.""" + bins = flt.bins + arr = np.asarray(bins) + + # Store energy filters as edge lists (E0..En), consistent with specifications.yaml. + if isinstance(flt, openmc.EnergyFilter): + if arr.ndim == 1: + return arr.tolist() + if arr.ndim == 2 and arr.shape[1] == 2: + low = arr[:, 0] + high = arr[:, 1] + if low.size == 0: + return [] + return np.concatenate([low[:1], high]).tolist() + return arr.reshape(-1).tolist() + + # Scalar bins -> simple list of values. + if arr.ndim <= 1: + return arr.tolist() + + # Structured bins (for example mesh-like tuples) -> nested lists. + return [list(np.asarray(b).tolist()) for b in bins] + + +def _build_filter_axis_metadata( + filters: list[openmc.Filter], + filter_dims: list[str], + spec_tally: dict | None = None, +): + """Create rich per-filter metadata with axis names and bin definitions.""" + spec_filters = [] + if isinstance(spec_tally, dict): + spec_filters = list(spec_tally.get("filters", [])) + + spec_idx = 0 + axes = [] + for flt, dim in zip(filters, filter_dims): + ftype = normalize_filter_type(type(flt).__name__) + spec_filter = None + if ftype != "particle" and spec_idx < len(spec_filters): + spec_filter = spec_filters[spec_idx] + spec_idx += 1 + + axis_meta = { + "name": type(flt).__name__, + "axis": dim, + "num_bins": int(flt.num_bins), + "bins": _serialize_filter_bins(flt), + } + + if isinstance(spec_filter, dict) and "units" in spec_filter: + axis_meta["units"] = spec_filter.get("units") + + if ftype == "energy": + # Energy filters are represented as edge lists in OFB results. + axis_meta["kind"] = "edges" + axis_meta["units"] = ( + spec_filter.get("units", "eV") + if isinstance(spec_filter, dict) + else "eV" + ) + axis_meta["closure"] = ( + spec_filter.get("closure", "[low, high)") + if isinstance(spec_filter, dict) + else "[low, high)" + ) + + axes.append(axis_meta) + return axes + + +def _build_spec_lookup(spec_tallies) -> dict[str, dict]: + """Build a lookup map from tally name to tally specification entry.""" + if spec_tallies is None: + return {} + if isinstance(spec_tallies, dict): + return spec_tallies + + lookup = {} + for entry in spec_tallies: + if not isinstance(entry, dict): + continue + name = entry.get("name") + if name: + lookup[str(name)] = entry + return lookup + + +def _to_1d_coord(values) -> np.ndarray: + """Normalize scalar/list-like coordinate input into a 1D numpy array.""" + arr = np.asarray(values) + if arr.ndim == 0: + arr = arr.reshape(1) + return arr + + +def openmc_tally_to_dataset( + tally: openmc.Tally, + tmc_coords: dict[str, Iterable] | None = None, + normalizer: Callable[[openmc.Tally, np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray]] | None = None, + spec_tally: dict | None = None, +) -> xr.Dataset: + """ + Convert one OpenMC tally into an xarray Dataset with variables `mean` and `mc_std`. + + The output schema mirrors the one used in `uq/tmc_manager.py`: + dimensions are `(tmc dims..., filter dims..., nuclide, score)` and metadata + is stored in dataset attrs (`filter_axes`, `nuclides`, `scores`). + """ + tmc_coords = tmc_coords or {} + + filters = list(tally.filters) + filter_bins = [f.num_bins for f in filters] + n_nuclides = max(len(tally.nuclides), 1) + n_scores = len(tally.scores) + nd_shape = tuple(filter_bins) + (n_nuclides, n_scores) + + mean_nd = tally.mean.reshape(nd_shape) + std_nd = tally.std_dev.reshape(nd_shape) + + if normalizer is not None: + mean_nd, std_nd = normalizer(tally, mean_nd, std_nd) + + tmc_dims = tuple(tmc_coords.keys()) + tmc_coord_arrays = {k: _to_1d_coord(v) for k, v in tmc_coords.items()} + tmc_shape = tuple(len(v) for v in tmc_coord_arrays.values()) + + full_shape = tmc_shape + nd_shape + mean_full = mean_nd.reshape((1,) * len(tmc_shape) + nd_shape) + std_full = std_nd.reshape((1,) * len(tmc_shape) + nd_shape) + + # Broadcast to provided TMC coordinate lengths (usually all 1 for a single write). + mean_full = np.broadcast_to(mean_full, full_shape) + std_full = np.broadcast_to(std_full, full_shape) + + filter_dims = _unique_filter_dims(filters) + dims = tmc_dims + tuple(filter_dims) + ("nuclide", "score") + + coords: dict[str, tuple[str, np.ndarray] | np.ndarray] = {} + for d, c in tmc_coord_arrays.items(): + coords[d] = (d, c) + for f, d in zip(filters, filter_dims): + coords[d] = (d, np.arange(f.num_bins)) + + nuclides = [str(n) for n in tally.nuclides] if tally.nuclides else ["total"] + scores = [str(s) for s in tally.scores] + coords["nuclide"] = ("nuclide", np.asarray(nuclides, dtype="U")) + coords["score"] = ("score", np.asarray(scores, dtype="U")) + + ds = xr.Dataset( + { + "mean": xr.DataArray(mean_full, dims=dims, coords=coords), + "mc_std": xr.DataArray(std_full, dims=dims, coords=coords), + } + ) + + tally_name = tally.name or f"tally_{tally.id}" + ds["mean"].attrs["tally_id"] = int(tally.id) + ds["mean"].attrs["tally_name"] = tally_name + ds["mc_std"].attrs["tally_id"] = int(tally.id) + ds["mc_std"].attrs["tally_name"] = tally_name + + filter_axes = _build_filter_axis_metadata(filters, filter_dims, spec_tally=spec_tally) + + ds.attrs["filter_axes"] = json.dumps(filter_axes) + ds.attrs["nuclides"] = json.dumps(nuclides) + ds.attrs["scores"] = json.dumps(scores) + + observed_tally = { + "name": tally_name, + "id": int(tally.id), + "filters": filter_axes, + "scores": scores, + "nuclides": nuclides, + } + ds.attrs["observed_tally"] = json.dumps(observed_tally) + + if spec_tally is not None: + ds.attrs["spec_tally"] = json.dumps(spec_tally) + consistent, issues = validate_tally_consistency(spec_tally, observed_tally) + # h5netcdf/netCDF attrs do not support boolean dtype reliably. + ds.attrs["spec_consistent"] = int(bool(consistent)) + ds.attrs["spec_consistency_issues"] = json.dumps(issues) + + return ds + + +def save_openmc_statepoint_tallies( + statepoint: openmc.StatePoint, + filename: str | Path, + tally_names: Iterable[str] | None = None, + spec_tallies: Iterable[dict] | dict[str, dict] | None = None, + tmc_coords: dict[str, Iterable] | None = None, + append_dim: str | None = None, + normalizer: Callable[[openmc.Tally, np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray]] | None = None, + group_by: str = "name", + engine: str = "h5netcdf", +) -> Path: + """ + Write OpenMC statepoint tallies to grouped HDF5/NetCDF datasets. + + Parameters + ---------- + group_by: + Group naming strategy, either: + - ``"name"``: use tally names (recommended for benchmark workflows) + - ``"id"``: use ``tally_`` groups + """ + if group_by not in {"name", "id"}: + raise ValueError("group_by must be either 'name' or 'id'") + + filename = Path(filename) + spec_lookup = _build_spec_lookup(spec_tallies) + if tally_names is None: + selected = list(statepoint.tallies.values()) + else: + selected = [statepoint.get_tally(name=name) for name in tally_names] + + for tally in selected: + ds_new = openmc_tally_to_dataset( + tally=tally, + tmc_coords=tmc_coords, + normalizer=normalizer, + spec_tally=spec_lookup.get(str(tally.name)), + ) + if group_by == "name" and tally.name: + group = str(tally.name) + else: + group = f"tally_{int(tally.id)}" + + # Keep explicit group metadata alongside tally_id / tally_name metadata. + ds_new.attrs["group"] = group + ds_new["mean"].attrs["tally_group"] = group + ds_new["mc_std"].attrs["tally_group"] = group + + if not filename.exists(): + ds_new.to_netcdf(filename, mode="w", group=group, engine=engine) + continue + + try: + ds_old = xr.open_dataset(filename, group=group, engine=engine) + if append_dim is not None and append_dim in ds_old.dims and append_dim in ds_new.dims: + ds_combined = xr.concat([ds_old, ds_new], dim=append_dim) + else: + ds_combined = ds_new + ds_old.close() + except (OSError, ValueError, KeyError): + ds_combined = ds_new + + with h5py.File(filename, "a") as h5f: + if group in h5f: + del h5f[group] + ds_combined.to_netcdf(filename, mode="a", group=group, engine=engine) + + return filename.resolve() + + +# Backward-compatible aliases while call sites migrate. +tally_to_dataset = openmc_tally_to_dataset +save_statepoint_tallies = save_openmc_statepoint_tallies diff --git a/src/openmc_fusion_benchmarks/benchmark.py b/src/openmc_fusion_benchmarks/benchmark.py index e4725f6a..2294fc76 100644 --- a/src/openmc_fusion_benchmarks/benchmark.py +++ b/src/openmc_fusion_benchmarks/benchmark.py @@ -3,10 +3,11 @@ import warnings from abc import ABC, abstractmethod import numpy as np -import xarray as xr -import h5py -from .validate import validate_benchmark -from .utils import _openmc_to_ofb, _save_result +from .validate_spec import validate_benchmark +from .backends.openmc.tallies import ( + make_default_openmc_normalizer, + save_openmc_statepoint_tallies, +) from .uq.tmc_engine import tmc_engine import openmc @@ -18,6 +19,31 @@ LFS_DIR = Path(__file__).parents[2] / "lfs" +def _openmc_to_ofb( + spec_tallies, + statepoint: openmc.StatePoint, + mesh: str = "mesh.h5m", + realization_label: str = "baseline", +): + """Backward-compatible postprocess entry point. + + Historically this symbol was patched in tests and used as the benchmark + postprocessing hook. Keep it as a thin adapter over the new backend-aware + tally serialization path. + """ + tally_names = [t["name"] for t in spec_tallies] + normalizer = make_default_openmc_normalizer(mesh) + save_openmc_statepoint_tallies( + statepoint=statepoint, + filename="benchmark_results.h5", + tally_names=tally_names, + spec_tallies=spec_tallies, + tmc_coords={"realization": [realization_label]}, + append_dim="realization", + normalizer=normalizer, + ) + + class Benchmark(ABC): def __init__(self, name: str): self.name = name @@ -401,16 +427,27 @@ def _build_model(self): ) return model - def _postprocess(self, statepoint: openmc.StatePoint, mesh: str = 'mesh.h5m'): + def _postprocess(self, statepoint: openmc.StatePoint | str | Path, mesh: str = 'mesh.h5m'): """Post-process the model after running.""" # Retrieve tallies data from specifications tallies_data = self._benchmark_spec['tallies'] - _openmc_to_ofb( - spec_tallies=tallies_data, - statepoint=statepoint, - mesh=mesh - ) + # Accept both already-open/duck-typed StatePoint objects and file paths. + if isinstance(statepoint, openmc.StatePoint) or hasattr(statepoint, "get_tally"): + _openmc_to_ofb( + statepoint=statepoint, + spec_tallies=tallies_data, + mesh=mesh, + realization_label="baseline", + ) + else: + with openmc.StatePoint(str(statepoint)) as sp: + _openmc_to_ofb( + statepoint=sp, + spec_tallies=tallies_data, + mesh=mesh, + realization_label="baseline", + ) return diff --git a/src/openmc_fusion_benchmarks/benchmark_results.py b/src/openmc_fusion_benchmarks/benchmark_results.py index aa73c073..981e9680 100644 --- a/src/openmc_fusion_benchmarks/benchmark_results.py +++ b/src/openmc_fusion_benchmarks/benchmark_results.py @@ -1,18 +1,46 @@ from __future__ import annotations +import json from pathlib import Path from typing import Union + import h5py import xarray as xr from .database import _resolve_database_path +from .tallies import Tally + + +def _open_dataset_with_fallback(filepath: Path, group: str) -> xr.Dataset: + """Open a group using available xarray backends. + + Prefer h5netcdf but gracefully fall back when optional dependencies are + not installed in the execution environment. + """ + last_exc = None + for engine in ("h5netcdf", "netcdf4", None): + try: + if engine is None: + return xr.open_dataset(filepath, group=group) + return xr.open_dataset(filepath, group=group, engine=engine) + except (ModuleNotFoundError, ImportError, OSError, ValueError) as exc: + last_exc = exc + continue + if last_exc is not None: + raise last_exc + raise RuntimeError(f"Could not open dataset group '{group}' from {filepath}") + + +class Results: + """Class to handle OFB results stored in HDF5 files. + + This is the generic entry point and can be used for benchmark and non-benchmark + workflows as long as they follow the OFB tally group schema. -class BenchmarkResults: - """Class to handle benchmark results stored in HDF5 files. usage: - - Load from arbitrary file path: `BenchmarkResults.from_file(filepath)` - - Load from run directory: `BenchmarkResults.from_run_dir(run_dir, filename)` - - Load from package database: `BenchmarkResults.from_database(benchmark, filename)` + - Load from arbitrary file path: `Results.from_file(filepath)` + - Load from run directory: `Results.from_run_dir(run_dir, filename)` + - Load from package database: `Results.from_database(benchmark, filename)` """ def __init__(self, filepath: Union[str, Path]): @@ -21,94 +49,184 @@ def __init__(self, filepath: Union[str, Path]): raise FileNotFoundError(f"Results file not found: {self.filepath}") @classmethod - def from_file(cls, filepath: Union[str, Path]) -> "BenchmarkResults": + def from_file(cls, filepath: Union[str, Path]) -> "Results": """Load results from an arbitrary file path.""" return cls(filepath) @classmethod - def from_run_dir(cls, run_dir: Union[str, Path] = ".", filename: str = "results.h5") -> "BenchmarkResults": + def from_run_dir( + cls, + run_dir: Union[str, Path] = ".", + filename: str = "benchmark_results.h5", + ) -> "Results": """Load results from a run directory (default: current directory).""" path = Path(run_dir) / filename return cls(path) @classmethod - def from_database(cls, benchmark: str, filename: str = "reference_results.h5") -> "BenchmarkResults": + def from_database(cls, benchmark: str, filename: str = "reference_results.h5") -> "Results": """Load reference results from the package database.""" db_path = _resolve_database_path(benchmark, filename) return cls(db_path) @property def tallies(self): - with h5py.File(self.filepath, 'r') as f: + with h5py.File(self.filepath, "r") as f: tallies = list(f.keys()) return tallies - def get_tally(self, name: str) -> xr.DataArray: - return xr.load_dataarray(self.filepath, group=name) - - -# # Some utilities for data analysis -# def get_means(tally: xr.DataArray) -> xr.DataArray: -# """Extract tally means from the tally DataArray.""" -# return tally.sel(column='mean').squeeze() - -# def get_stds(tally: xr.DataArray) -> xr.DataArray: -# """Extract tally standard deviations from the tally DataArray.""" -# return tally.sel(column='std. dev.').squeeze() - -# def get_rstds(tally: xr.DataArray) -> xr.DataArray: -# """Compute tally relative standard deviations from the tally DataArray.""" -# mean_vals = get_means(tally) -# std_vals = get_stds(tally) -# return std_vals / mean_vals - - -# # UQ-TMC base analysis functions - Move in uq/ ? -# def mean_of_means(tally: xr.DataArray) -> xr.DataArray: -# """Compute the mean of the means across realizations.""" -# means = get_means(tally) -# return means.mean(dim='realization') - -# def std_of_means(tally: xr.DataArray) -> xr.DataArray: -# """Compute the standard deviation of the means across realizations.""" -# means = get_means(tally) -# return means.std(dim='realization') - -# def rstd_of_means(tally: xr.DataArray) -> xr.DataArray: -# """Compute the relative standard deviation of the means across realizations.""" -# mean_vals = mean_of_means(tally) -# std_vals = std_of_means(tally) -# return std_vals / mean_vals - -# # UQ-TMC dynamic realization analysis functions - Move in uq/ ? -# def dynamic_mean_of_means(tally: xr.DataArray) -> np.ndarray: -# """Compute the dynamic mean of the means across realizations.""" -# means = get_means(tally) -# return np.array([means[:i].mean(dim='realization') for i in range(2, len(means.realization) + 1)]) - -# def dynamic_std_of_means(tally: xr.DataArray) -> np.ndarray: -# """Compute the dynamic standard deviation of the means across realizations.""" -# means = get_means(tally) -# return np.array([ -# means[:i].std(dim='realization') for i in range(2, len(means.realization) + 1) -# ]) - -# def dynamic_rstd_of_means(tally: xr.DataArray) -> np.ndarray: -# """Compute the dynamic relative standard deviation of the means across realizations.""" -# means = get_means(tally) -# return np.array([ -# means[:i].std(dim='realization') / means[:i].mean(dim='realization') -# for i in range(2, len(means.realization) + 1) -# ]) - -# def dynamic_rstd_of_rstds(tally: xr.DataArray) -> np.ndarray: -# """Compute the dynamic relative standard deviation of the relative standard deviations across realizations.""" -# rstds = dynamic_rstd_of_means(tally) -# return np.array([ -# rstds[:i].std() / rstds[:i].mean() for i in range(2, len(rstds) + 1) -# ]) - -# def derivative_of_dynamic_rstds(tally: xr.DataArray) -> np.ndarray: -# """Compute the derivative of the dynamic relative standard deviations.""" -# dynamic_rstds = dynamic_rstd_of_means(tally) -# return np.gradient(dynamic_rstds, axis=0) \ No newline at end of file + def get_tally(self, name: str) -> Tally: + """ + Get a tally wrapper. + + `name` can be either: + - an HDF5 group name (for example `tally_1`), or + - a tally logical name stored in attrs (for example `nuclear_heating`). + """ + group = None + + with h5py.File(self.filepath, "r") as f: + if name in f: + group = name + else: + # Fallback lookup by tally_name attribute. + for candidate in f.keys(): + try: + with _open_dataset_with_fallback(self.filepath, candidate) as ds: + target_var = "mean" + if "mean" not in ds: + if len(ds.data_vars) != 1: + continue + target_var = next(iter(ds.data_vars)) + if ds[target_var].attrs.get("tally_name") == name: + group = candidate + break + except (OSError, ValueError, KeyError): + continue + + if group is None: + raise ValueError(f"No tally with name or group '{name}' found") + + ds = _open_dataset_with_fallback(self.filepath, group) + if "mean" not in ds: + # Legacy layout support: return the only data variable as DataArray. + if len(ds.data_vars) == 1: + var_name = next(iter(ds.data_vars)) + da = ds[var_name].load() + ds.close() + return da + ds.close() + raise ValueError(f"Group '{group}' does not contain a 'mean' dataset") + + da_mean = ds["mean"] + da_mc_std = ds["mc_std"] if "mc_std" in ds else None + return Tally(da_mean, da_mc_std, parent_ds=ds) + + def get_spec_consistency_report(self, only_mismatches: bool = False): + """ + Return spec-vs-observed consistency metadata for each tally group. + + Parameters + ---------- + only_mismatches: + If True, return only entries with `spec_consistent == 0`. + + Returns + ------- + list[dict] + One entry per tally group with fields: + - `group` + - `tally_name` + - `spec_consistent` (0/1/None) + - `issues` (list[str]) + - `spec_tally` (dict) + - `observed_tally` (dict) + """ + + def _loads_attr(attrs, key, default): + value = attrs.get(key) + if value is None: + return default + if isinstance(value, bytes): + value = value.decode("utf-8") + try: + return json.loads(value) + except Exception: + return default + + report = [] + with h5py.File(self.filepath, "r") as f: + for group in f.keys(): + attrs = f[group].attrs + spec_consistent = attrs.get("spec_consistent") + if spec_consistent is not None: + if hasattr(spec_consistent, "shape"): + spec_consistent = int(spec_consistent[0]) + else: + spec_consistent = int(spec_consistent) + + entry = { + "group": group, + "tally_name": attrs.get("tally_name"), + "spec_consistent": spec_consistent, + "issues": _loads_attr(attrs, "spec_consistency_issues", []), + "spec_tally": _loads_attr(attrs, "spec_tally", {}), + "observed_tally": _loads_attr(attrs, "observed_tally", {}), + } + + # Fallback tally_name from nested metadata when variable-level attrs were used. + if entry["tally_name"] in (None, ""): + entry["tally_name"] = entry["observed_tally"].get("name") + + if only_mismatches and entry["spec_consistent"] != 0: + continue + + report.append(entry) + + return report + + def format_spec_consistency_report(self, only_mismatches: bool = False) -> str: + """Return a human-readable spec consistency report string.""" + report = self.get_spec_consistency_report(only_mismatches=only_mismatches) + + if not report: + if only_mismatches: + return "Spec Consistency Report\nNo mismatches found." + return "Spec Consistency Report\nNo tallies found." + + label_w = 18 + lines: list[str] = ["Spec Consistency Report"] + + for entry in report: + status = entry.get("spec_consistent") + if status == 1: + status_text = "OK" + elif status == 0: + status_text = "MISMATCH" + else: + status_text = "N/A" + + group = entry.get("group") + tally_name = entry.get("tally_name") + issues = entry.get("issues") or [] + + lines.append("") + lines.append(f"Tally {tally_name} (group={group})") + lines.append(f"{'Status':<{label_w}}: {status_text}") + + if issues: + lines.append(f"{'Issues':<{label_w}}: {len(issues)}") + for issue in issues: + lines.append(f" - {issue}") + else: + lines.append(f"{'Issues':<{label_w}}: none") + + return "\n".join(lines) + + +class BenchmarkResults(Results): + """Backward-compatible alias class for benchmark-centric naming.""" + + +# Optional explicit alias for OFB naming style. +OFBResults = Results diff --git a/src/openmc_fusion_benchmarks/benchmarks/oktavian_al/specifications.yaml b/src/openmc_fusion_benchmarks/benchmarks/oktavian_al/specifications.yaml index e4b8e06f..036982ac 100644 --- a/src/openmc_fusion_benchmarks/benchmarks/oktavian_al/specifications.yaml +++ b/src/openmc_fusion_benchmarks/benchmarks/oktavian_al/specifications.yaml @@ -74,6 +74,7 @@ geometry: - id: 1 material: void mesh_size: 5 + mesh_size: 5 - id: 2 material: void mesh_size: 5 @@ -138,6 +139,7 @@ sources: settings: run_mode: fixed_source batches: 100 + particles_per_batch: 10000 particles_per_batch: 100000 photon_transport: true @@ -150,6 +152,7 @@ tallies: - type: surface values: [7] - type: energy + closure: "[low, high)" values: [97122., 101090., 105210., 109500., 113970., 118620., 123470., 128500., 133750., 139210., 144890., 150800., 156960., 163360., 170030., 176970., 184190., 191710., @@ -183,6 +186,7 @@ tallies: - type: surface values: [7] - type: energy + closure: "[low, high)" values: [500000., 600000., 700000., 800000., 900000., 1000000., 1100000., 1200000., 1300000., 1400000., 1500000., 1600000., 1700000., 1800000., 1900000., 2000000., 2100000., 2200000., diff --git a/src/openmc_fusion_benchmarks/tallies.py b/src/openmc_fusion_benchmarks/tallies.py new file mode 100644 index 00000000..1dd08634 --- /dev/null +++ b/src/openmc_fusion_benchmarks/tallies.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +import json + +import numpy as np +import xarray as xr + + +class BaseTally: + """Common wrapper for structured tally datasets.""" + + def __init__(self, mean_da: xr.DataArray, mc_std_da: xr.DataArray | None = None, parent_ds: xr.Dataset | None = None): + self._da = mean_da + self._da_mc_std = mc_std_da + self._parent_ds = parent_ds + + @property + def id(self): + """Tally ID.""" + return self._da.attrs.get("tally_id") + + @property + def name(self): + """Tally name.""" + return self._da.attrs.get("tally_name") + + @property + def data(self): + """Underlying mean data array values.""" + return self._da.values + + @property + def scores(self): + """List of score names.""" + if "score" in self._da.coords: + return [str(s) for s in self._da.coords["score"].values] + + if self._parent_ds is not None: + scores_json = self._parent_ds.attrs.get("scores") + if scores_json: + return json.loads(scores_json) + + scores_json = self._da.attrs.get("scores") + if scores_json: + return json.loads(scores_json) + return [] + + @property + def nuclides(self): + """List of nuclide names.""" + if "nuclide" in self._da.coords: + return [str(n) for n in self._da.coords["nuclide"].values] + + if self._parent_ds is not None: + nuclides_json = self._parent_ds.attrs.get("nuclides") + if nuclides_json: + return json.loads(nuclides_json) + + nuclides_json = self._da.attrs.get("nuclides") + if nuclides_json: + return json.loads(nuclides_json) + return [] + + @property + def filters(self): + """List of filter metadata entries.""" + if self._parent_ds is not None: + filt_json = self._parent_ds.attrs.get("filter_axes") + if filt_json: + return json.loads(filt_json) + + filt_json = self._da.attrs.get("filter_axes") + if filt_json: + return json.loads(filt_json) + return [] + + @property + def shape(self): + """Shape of the underlying mean data array.""" + return self._da.shape + + @property + def dims(self): + """Dimension names of the underlying mean data array.""" + return self._da.dims + + @property + def mean(self): + """Mean array for non-TMC tallies (already the final mean).""" + return self._da.values + + @property + def std_dev(self): + """MC standard deviation array for non-TMC tallies.""" + if self._da_mc_std is not None: + return self._da_mc_std.values + return np.zeros_like(self._da.values) + + def dimension_report(self): + """Return a compact mapping report for OFB and OpenMC-equivalent shapes.""" + dims = list(self._da.dims) + ofb_shape = tuple(int(self._da.sizes[d]) for d in dims) + + filter_axes = self.filters if isinstance(self.filters, list) else [] + filter_dims = [ + str(a.get("axis")) + for a in filter_axes + if isinstance(a, dict) and a.get("axis") in self._da.dims + ] + + if not filter_dims: + known_tmc_dims = {"realization", "sample", "replica", "batch", "iteration", "case"} + filter_dims = [d for d in dims if d not in known_tmc_dims and d not in {"nuclide", "score"}] + + filter_dim_sizes = {d: int(self._da.sizes[d]) for d in filter_dims} + flat_filter_bins = int(np.prod(list(filter_dim_sizes.values()))) if filter_dim_sizes else 1 + + nuclide_size = int(self._da.sizes["nuclide"]) if "nuclide" in self._da.sizes else max(len(self.nuclides), 1) + score_size = int(self._da.sizes["score"]) if "score" in self._da.sizes else max(len(self.scores), 1) + + tmc_dims = [d for d in dims if d not in set(filter_dims) and d not in {"nuclide", "score"}] + tmc_sizes = {d: int(self._da.sizes[d]) for d in tmc_dims} + + return { + "tally_name": self.name, + "tally_id": self.id, + "ofb_dims": dims, + "ofb_shape": ofb_shape, + "tmc_dims": tmc_dims, + "tmc_sizes": tmc_sizes, + "filter_dims": filter_dims, + "filter_dim_sizes": filter_dim_sizes, + "nuclide_size": nuclide_size, + "score_size": score_size, + "openmc_equivalent_raw_shape": (flat_filter_bins, nuclide_size, score_size), + } + + def get_dimension_report(self): + """Backward/UX alias to ``dimension_report``.""" + return self.dimension_report() + + def format_dimension_report(self): + """Return a readable text report from ``dimension_report``.""" + report = self.dimension_report() + + def _fmt_axis_sizes(axis_names, axis_sizes): + if not axis_names: + return "none" + return ", ".join(f"{name}={axis_sizes.get(name, '?')}" for name in axis_names) + + tally_name = report.get("tally_name") or "" + tally_id = report.get("tally_id") + ofb_dims = report.get("ofb_dims", []) + ofb_shape = report.get("ofb_shape", ()) + tmc_dims = report.get("tmc_dims", []) + tmc_sizes = report.get("tmc_sizes", {}) + filter_dims = report.get("filter_dims", []) + filter_sizes = report.get("filter_dim_sizes", {}) + raw_shape = report.get("openmc_equivalent_raw_shape", ()) + + label_w = 20 + lines = [ + f"Tally {tally_name} (id={tally_id})", + f"{'OFB dims':<{label_w}}: {ofb_dims}", + f"{'OFB shape':<{label_w}}: {ofb_shape}", + f"{'TMC axes':<{label_w}}: {_fmt_axis_sizes(tmc_dims, tmc_sizes)}", + f"{'Filter axes':<{label_w}}: {_fmt_axis_sizes(filter_dims, filter_sizes)}", + f"{'Nuclide/score sizes':<{label_w}}: nuclide={report.get('nuclide_size')}, score={report.get('score_size')}", + f"{'OpenMC raw equivalent':<{label_w}}: {raw_shape} (flat_filters, nuclide, score)", + ] + return "\n".join(lines) + + def get_slice(self, scores=None, nuclides=None, **filter_kwargs): + """Get a filtered xarray view of the mean tally data.""" + da = self._da + + if filter_kwargs: + da = da.sel(**filter_kwargs) + + if scores is not None and "score" in da.dims: + all_scores = self.scores + score_indices = [all_scores.index(s) for s in scores] + da = da.isel(score=score_indices) + + if nuclides is not None and "nuclide" in da.dims: + all_nuclides = self.nuclides + nuclide_indices = [all_nuclides.index(n) for n in nuclides] + da = da.isel(nuclide=nuclide_indices) + + return da + + +class Tally(BaseTally): + """Generic OFB tally wrapper for non-TMC result groups.""" + + def __repr__(self): + return f"" + +def tally_to_dataset(*args, **kwargs): + """Backward-compatible alias to the OpenMC backend serializer.""" + from .backends.openmc.tallies import openmc_tally_to_dataset + + return openmc_tally_to_dataset(*args, **kwargs) + + +def save_statepoint_tallies(*args, **kwargs): + """Backward-compatible alias to the OpenMC backend statepoint writer.""" + from .backends.openmc.tallies import save_openmc_statepoint_tallies + + return save_openmc_statepoint_tallies(*args, **kwargs) + diff --git a/src/openmc_fusion_benchmarks/uq/__init__.py b/src/openmc_fusion_benchmarks/uq/__init__.py index 17660be2..6e71c695 100644 --- a/src/openmc_fusion_benchmarks/uq/__init__.py +++ b/src/openmc_fusion_benchmarks/uq/__init__.py @@ -1,3 +1,4 @@ from .uq_utils import * from .tmc_engine import * from .tmc_manager import * +from .tmc_statepoint import * diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index 6414c668..f7744308 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -5,9 +5,9 @@ import json import copy import xarray as xr -import inspect import itertools -import h5py +from .tmc_statepoint import TMCStatePoint +from ..validate_results import validate_tally_consistency class TMCManager: @@ -31,7 +31,7 @@ def __init__(self, base_model: openmc.Model, perturbations: List[Callable], # Wrap user perturbations into indexed perturb(model, idx) self.perturbations = self._build_indexed_perturbations(perturbations) - def run(self, mode="matrix", cwd='.', *args, **kwargs): + def run(self, mode="matrix", cwd='.', benchmark_export=None, *args, **kwargs): cwd = Path(cwd).resolve() @@ -73,7 +73,7 @@ def run(self, mode="matrix", cwd='.', *args, **kwargs): f_manifest.write(json.dumps(rec) + "\n") f_manifest.flush() - self._process_tmc(manifest_path=manifest) + self._process_tmc(manifest_path=manifest, benchmark_export=benchmark_export) return # --- Matrix / diagonal modes share the same structure and manifest keys --- @@ -115,7 +115,7 @@ def run(self, mode="matrix", cwd='.', *args, **kwargs): f_manifest.flush() # Postprocess the whole TMC set - self._process_tmc(manifest_path=manifest) + self._process_tmc(manifest_path=manifest, benchmark_export=benchmark_export) def _build_indexed_perturbations(self, user_factories): """ @@ -145,7 +145,7 @@ def perturb(model, idx): return indexed - def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): + def _process_tmc(self, manifest_path="tmc_manifest.jsonl", benchmark_export=None): manifest_path = Path(manifest_path).resolve() tmc_dir = manifest_path.parent # directory containing the manifest @@ -156,7 +156,105 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): if tmc_statepoint.exists(): tmc_statepoint.unlink() + benchmark_file = None + benchmark_normalizer = None + spec_lookup = {} + if benchmark_export: + benchmark_file = Path(benchmark_export.get("filename", "benchmark_results.h5")) + if not benchmark_file.is_absolute(): + benchmark_file = (tmc_dir / benchmark_file).resolve() + if benchmark_file.exists(): + benchmark_file.unlink() + + benchmark_normalizer = benchmark_export.get("normalizer") + + spec_tallies = benchmark_export.get("spec_tallies") + if isinstance(spec_tallies, dict): + spec_lookup = spec_tallies + elif spec_tallies is not None: + for entry in spec_tallies: + if isinstance(entry, dict) and entry.get("name"): + spec_lookup[str(entry["name"])] = entry + # ---- helper: resolve statepoint path robustly ---- + def _normalize_filter_type(type_name: str) -> str: + name = str(type_name).strip().lower() + if name.endswith("filter"): + name = name[:-6] + return name + + def _unique_filter_dims(filters): + counts = {} + dims = [] + for flt in filters: + base = type(flt).__name__.replace("Filter", "").lower() or "filter" + idx = counts.get(base, 0) + counts[base] = idx + 1 + dims.append(base if idx == 0 else f"{base}_{idx}") + return dims + + def _serialize_filter_bins(flt): + bins = flt.bins + arr = np.asarray(bins) + + if isinstance(flt, openmc.EnergyFilter): + if arr.ndim == 1: + return arr.tolist() + if arr.ndim == 2 and arr.shape[1] == 2: + low = arr[:, 0] + high = arr[:, 1] + if low.size == 0: + return [] + return np.concatenate([low[:1], high]).tolist() + return arr.reshape(-1).tolist() + + if arr.ndim <= 1: + return arr.tolist() + + return [list(np.asarray(b).tolist()) for b in bins] + + def _build_filter_axis_metadata(filters, filter_dims): + axes = [] + for flt, dim in zip(filters, filter_dims): + ftype = _normalize_filter_type(type(flt).__name__) + axis_meta = { + "name": type(flt).__name__, + "axis": dim, + "num_bins": int(flt.num_bins), + "bins": _serialize_filter_bins(flt), + } + if ftype == "energy": + axis_meta["kind"] = "edges" + axis_meta["units"] = "eV" + axis_meta["closure"] = "[low, high)" + axes.append(axis_meta) + return axes + + def _safe_group_name(name: str, tid: int, used: set[str]) -> str: + base = (name or "").strip() + if not base: + candidate = f"tally_{int(tid)}" + else: + # Avoid creating nested paths in HDF groups. + candidate = base.replace("/", "_") + + if candidate not in used: + used.add(candidate) + return candidate + + with_id = f"{candidate}__{int(tid)}" + if with_id not in used: + used.add(with_id) + return with_id + + i = 1 + while True: + alt = f"{with_id}_{i}" + if alt not in used: + used.add(alt) + return alt + i += 1 + def resolve_statepoint_path(sp_str: str) -> Path: sp_path = Path(sp_str) if sp_path.is_absolute(): @@ -280,13 +378,16 @@ def resolve_statepoint_path(sp_str: str) -> Path: tally_shapes[tid] = nd_shape tally_filters[tid] = filters + filter_dims = _unique_filter_dims(filters) + filter_axes = _build_filter_axis_metadata(filters, filter_dims) + scores = [str(s) for s in tally.scores] + nuclides = [str(n) for n in tally.nuclides] if tally.nuclides else ["total"] + axis_info = { - "filter_axes": [ - {"name": type(f).__name__, "num_bins": f.num_bins} - for f in filters - ], - "nuclides": [str(n) for n in tally.nuclides] if tally.nuclides else ["total"], - "scores": list(tally.scores), + "filter_dims": filter_dims, + "filter_axes": filter_axes, + "nuclides": nuclides, + "scores": scores, } tally_axisinfo[tid] = axis_info @@ -294,11 +395,16 @@ def resolve_statepoint_path(sp_str: str) -> Path: # ---- 4. Allocate arrays: one per tally ---- tmc_data = {} # tid -> ndarray (extra_shape + nd_shape) tmc_mc_std = {} # tid -> ndarray (extra_shape + nd_shape) + benchmark_data = {} + benchmark_mc_std = {} for tid, nd_shape in tally_shapes.items(): full_shape = extra_shape + nd_shape tmc_data[tid] = np.empty(full_shape, dtype=float) tmc_mc_std[tid] = np.empty(full_shape, dtype=float) + if benchmark_file is not None: + benchmark_data[tid] = np.empty(full_shape, dtype=float) + benchmark_mc_std[tid] = np.empty(full_shape, dtype=float) # ---- 5. Fill arrays by looping over statepoints ---- if mode == "sequential": @@ -317,6 +423,13 @@ def resolve_statepoint_path(sp_str: str) -> Path: r_idx = rec["realization"] tmc_data[tid][p_idx, r_idx, ...] = mean_nd tmc_mc_std[tid][p_idx, r_idx, ...] = std_nd + if benchmark_file is not None: + if benchmark_normalizer is not None: + b_mean, b_std = benchmark_normalizer(tally, mean_nd.copy(), std_nd.copy()) + else: + b_mean, b_std = mean_nd, std_nd + benchmark_data[tid][p_idx, r_idx, ...] = b_mean + benchmark_mc_std[tid][p_idx, r_idx, ...] = b_std elif mode == "diagonal": # 1D structure: one "realization" dim @@ -332,16 +445,28 @@ def resolve_statepoint_path(sp_str: str) -> Path: std_nd = std_flat.reshape(nd_shape) tmc_data[tid][i, ...] = mean_nd tmc_mc_std[tid][i, ...] = std_nd + if benchmark_file is not None: + if benchmark_normalizer is not None: + b_mean, b_std = benchmark_normalizer(tally, mean_nd.copy(), std_nd.copy()) + else: + b_mean, b_std = mean_nd, std_nd + benchmark_data[tid][i, ...] = b_mean + benchmark_mc_std[tid][i, ...] = b_std else: # mode == "matrix" # Fill as flat (n_combos, ...) then reshape first axis into extra_shape n_combos = len(records) flat_data = {} flat_mc_std = {} + flat_benchmark_data = {} + flat_benchmark_mc_std = {} for tid, nd_shape in tally_shapes.items(): flat_shape = (n_combos,) + nd_shape flat_data[tid] = np.empty(flat_shape, dtype=float) flat_mc_std[tid] = np.empty(flat_shape, dtype=float) + if benchmark_file is not None: + flat_benchmark_data[tid] = np.empty(flat_shape, dtype=float) + flat_benchmark_mc_std[tid] = np.empty(flat_shape, dtype=float) for i, rec in enumerate(records): sp_path = resolve_statepoint_path(rec["statepoint"]) @@ -355,6 +480,13 @@ def resolve_statepoint_path(sp_str: str) -> Path: std_nd = std_flat.reshape(nd_shape) flat_data[tid][i, ...] = mean_nd flat_mc_std[tid][i, ...] = std_nd + if benchmark_file is not None: + if benchmark_normalizer is not None: + b_mean, b_std = benchmark_normalizer(tally, mean_nd.copy(), std_nd.copy()) + else: + b_mean, b_std = mean_nd, std_nd + flat_benchmark_data[tid][i, ...] = b_mean + flat_benchmark_mc_std[tid][i, ...] = b_std # reshape into extra_shape + nd_shape for tid, arr in flat_data.items(): @@ -363,19 +495,25 @@ def resolve_statepoint_path(sp_str: str) -> Path: for tid, arr in flat_mc_std.items(): arr.shape = extra_shape + tally_shapes[tid] tmc_mc_std[tid][...] = arr + if benchmark_file is not None: + for tid, arr in flat_benchmark_data.items(): + arr.shape = extra_shape + tally_shapes[tid] + benchmark_data[tid][...] = arr + for tid, arr in flat_benchmark_mc_std.items(): + arr.shape = extra_shape + tally_shapes[tid] + benchmark_mc_std[tid][...] = arr # ---- 6. Build per-tally Datasets and write each into its own group ---- + used_group_names = set() + for tid, arr in tmc_data.items(): nd_shape = tally_shapes[tid] filters = tally_filters[tid] axisinfo = tally_axisinfo[tid] - # filter dims based on filter types - filter_dims = [] - for f in filters: - filter_type = type(f).__name__.replace("Filter", "").lower() - filter_dims.append(filter_type) + # filter dims based on filter types (stable + unique) + filter_dims = list(axisinfo["filter_dims"]) # within each tally group, we can use generic "nuclide" and "score" dims = extra_dims + tuple(filter_dims) + ("nuclide", "score") @@ -389,7 +527,7 @@ def resolve_statepoint_path(sp_str: str) -> Path: # coords: nuclide / score for this tally nuclides = axisinfo["nuclides"] - scores = axisinfo["scores"] + scores = axisinfo["scores"] coords["nuclide"] = ("nuclide", np.array(nuclides, dtype="U")) coords["score"] = ("score", np.array(scores, dtype="U")) @@ -417,8 +555,24 @@ def resolve_statepoint_path(sp_str: str) -> Path: target_da.attrs["tally_id"] = tid target_da.attrs["tally_name"] = tally_name + observed_tally = { + "name": tally_name, + "id": int(tid), + "filters": axisinfo["filter_axes"], + "scores": scores, + "nuclides": nuclides, + } + + ds_tid.attrs["filter_axes"] = json.dumps(axisinfo["filter_axes"]) + ds_tid.attrs["nuclides"] = json.dumps(nuclides) + ds_tid.attrs["scores"] = json.dumps(scores) + ds_tid.attrs["observed_tally"] = json.dumps(observed_tally) + # serialize complex axisinfo at dataset level for k, v in axisinfo.items(): + if k in {"filter_axes", "nuclides", "scores"}: + # Canonical copies already written above. + continue if isinstance(v, (int, float, bool, str, np.number)): ds_tid.attrs[k] = v else: @@ -431,7 +585,10 @@ def resolve_statepoint_path(sp_str: str) -> Path: ds_tid["mean"] = da_mean ds_tid["mc_std"] = da_mc_std - group_name = f"tally_{tid}" + group_name = _safe_group_name(tally_name, tid, used_group_names) + ds_tid.attrs["group"] = group_name + ds_tid["mean"].attrs["tally_group"] = group_name + ds_tid["mc_std"].attrs["tally_group"] = group_name # append this tally-dataset as a group in the same file ds_tid.to_netcdf( @@ -441,6 +598,52 @@ def resolve_statepoint_path(sp_str: str) -> Path: engine="h5netcdf", # or "netcdf4" ) + if benchmark_file is not None: + da_mean_b = xr.DataArray( + benchmark_data[tid], + dims=dims, + coords=coords, + name="mean", + ) + da_mc_std_b = xr.DataArray( + benchmark_mc_std[tid], + dims=dims, + coords=coords, + name="mc_std", + ) + + for target_da in (da_mean_b, da_mc_std_b): + target_da.attrs["tally_id"] = tid + target_da.attrs["tally_name"] = tally_name + target_da.attrs["tally_group"] = group_name + + ds_b = xr.Dataset( + { + "mean": da_mean_b, + "mc_std": da_mc_std_b, + } + ) + ds_b.attrs["filter_axes"] = json.dumps(axisinfo["filter_axes"]) + ds_b.attrs["nuclides"] = json.dumps(nuclides) + ds_b.attrs["scores"] = json.dumps(scores) + ds_b.attrs["observed_tally"] = json.dumps(observed_tally) + ds_b.attrs["group"] = group_name + + spec_tally = spec_lookup.get(str(tally_name)) + if spec_tally is not None: + ds_b.attrs["spec_tally"] = json.dumps(spec_tally) + consistent, issues = validate_tally_consistency(spec_tally, observed_tally) + ds_b.attrs["spec_consistent"] = int(bool(consistent)) + ds_b.attrs["spec_consistency_issues"] = json.dumps(issues) + + mode_write = "a" if benchmark_file.exists() else "w" + ds_b.to_netcdf( + benchmark_file, + mode=mode_write, + group=group_name, + engine="h5netcdf", + ) + self.tmc_statepoint_path = tmc_statepoint def get_tmc_statepoint(self, path=None): @@ -465,382 +668,4 @@ def get_tmc_statepoint(self, path=None): else: path = Path(path).resolve() - return TMCStatePoint(path) - - -class TMCStatePoint: - """ - Wrapper for TMC statepoint providing an OpenMC StatePoint-like interface. - - Parameters - ---------- - path : str or Path - Path to the TMC statepoint NetCDF/HDF5 file. - """ - - def __init__(self, path): - self.path = Path(path).resolve() - # We won't keep one global ds; we'll open per-tally groups as needed. - self._tallies = None - - @property - def tallies(self): - """Dictionary of tallies, indexed by tally ID (mimics openmc.StatePoint.tallies).""" - if self._tallies is None: - self._tallies = {} - # Discover tally groups via h5py - with h5py.File(self.path, "r") as f: - for group_name in f.keys(): - if not group_name.startswith("tally_"): - continue - # Open this group as an xarray Dataset - ds = xr.open_dataset( - self.path, - group=group_name, - engine="h5netcdf", - ) - if "mean" not in ds: - continue - da_mean = ds["mean"] - da_mc_std = ds["mc_std"] if "mc_std" in ds else None - - tally_id = da_mean.attrs.get("tally_id") - if tally_id is None: - # Fallback: parse id from group name - try: - tally_id = int(group_name.split("_", 1)[1]) - except Exception: - continue - - self._tallies[tally_id] = TMCTally(da_mean, da_mc_std, parent_ds=ds) - return self._tallies - - def get_tally(self, tally_id=None, name=None): - """ - Get a tally by ID or name (mimics openmc.StatePoint.get_tally). - - Parameters - ---------- - tally_id : int, optional - Tally ID - name : str, optional - Tally name - - Returns - ------- - TMCTally - The requested tally - """ - if tally_id is not None: - try: - return self.tallies[tally_id] - except KeyError: - raise ValueError(f"No tally with id '{tally_id}' found") - elif name is not None: - for tally in self.tallies.values(): - if tally.name == name: - return tally - raise ValueError(f"No tally with name '{name}' found") - else: - raise ValueError("Must specify either 'tally_id' or 'name'") - - def close(self): - """No persistent open Dataset to close, but keep for API symmetry.""" - # If you decide to cache per-group ds objects, close them here. - pass - - def __enter__(self): - return self - - def __exit__(self, *args): - self.close() - - def __repr__(self): - n_tallies = len(self.tallies) - # Try to infer "realizations" (or more generally TMC size along first dim) - n_realizations = 0 - if self.tallies: - any_tally = next(iter(self.tallies.values())) - tmc_dims = any_tally._tmc_dims - if tmc_dims: - # product of TMC dims sizes - sz = 1 - for d in tmc_dims: - sz *= any_tally._da.sizes[d] - n_realizations = sz - return f"" - - -class TMCTally: - """ - Wrapper for a single TMC tally providing an OpenMC Tally-like interface. - - Parameters - ---------- - mean_da : xarray.DataArray - The DataArray containing the TMC mean values for this tally. - mc_std_da : xarray.DataArray, optional - The DataArray containing the MC std dev per run/combo. - parent_ds : xarray.Dataset, optional - Parent dataset (group) containing metadata attributes. - """ - - def __init__(self, mean_da, mc_std_da=None, parent_ds=None): - self._da = mean_da - self._da_mc_std = mc_std_da - self._parent_ds = parent_ds - - # Identify TMC dimensions: "perturbation" and "realization" for sequential, "perturbation_*" for matrix - self._tmc_dims = [ - d for d in self._da.dims - if d in ("perturbation", "realization") or d.startswith("perturbation_") - ] - - @property - def id(self): - """Tally ID.""" - return self._da.attrs.get("tally_id") - - @property - def name(self): - """Tally name.""" - return self._da.attrs.get("tally_name") - - # --- Metadata & helpers --- - - @property - def data(self): - """Full TMC mean data array (all TMC entries).""" - return self._da.values - - @property - def scores(self): - """List of score names.""" - # Preferred: from 'score' coordinate, if present - if "score" in self._da.coords: - return [str(s) for s in self._da.coords["score"].values] - - # Fallback: from parent_ds attrs "scores" (JSON) - if self._parent_ds is not None: - scores_json = self._parent_ds.attrs.get("scores") - if scores_json: - return json.loads(scores_json) - # Last resort: from variable attrs (if you changed writer) - scores_json = self._da.attrs.get("scores") - if scores_json: - return json.loads(scores_json) - return [] - - @property - def nuclides(self): - """List of nuclide names.""" - if "nuclide" in self._da.coords: - return [str(n) for n in self._da.coords["nuclide"].values] - - if self._parent_ds is not None: - nuclides_json = self._parent_ds.attrs.get("nuclides") - if nuclides_json: - return json.loads(nuclides_json) - nuclides_json = self._da.attrs.get("nuclides") - if nuclides_json: - return json.loads(nuclides_json) - return [] - - @property - def filters(self): - """List of filter information (type and number of bins).""" - if self._parent_ds is not None: - filt_json = self._parent_ds.attrs.get("filter_axes") - if filt_json: - return json.loads(filt_json) - filt_json = self._da.attrs.get("filter_axes") - if filt_json: - return json.loads(filt_json) - return [] - - @property - def tmc_dims(self): - """Names of TMC dimensions (realization / perturbation_*).""" - return tuple(self._tmc_dims) - - @property - def shape(self): - """Shape of the underlying mean data array.""" - return self._da.shape - - @property - def dims(self): - """Dimension names of the underlying mean data array.""" - return self._da.dims - - # --- TMC statistics --- - - @property - def mean(self): - """ - TMC mean across all TMC dimensions (realization / perturbation_*). - """ - if not self._tmc_dims: - return self._da.values - return self._da.mean(dim=self._tmc_dims).values - - @property - def std_dev(self): - """ - TMC standard deviation across all TMC dimensions. - - This is the propagated parametric uncertainty from the ensemble, - not the MC sampling error within each run. - """ - if not self._tmc_dims: - return np.zeros_like(self._da.values) - return self._da.std(dim=self._tmc_dims).values - - @property - def per_realization_mean(self): - """ - Raw mean value for each TMC point (all TMC dims retained). - Shape: (TMC dims..., filters..., nuclide, score). - """ - return self._da.values - - @property - def per_realization_std_dev(self): - """ - Monte Carlo standard deviation for each run/combo. - - This is the statistical uncertainty from particle sampling within each - individual OpenMC run, shaped like self._da. - """ - if self._da_mc_std is not None: - return self._da_mc_std.values - return np.zeros_like(self._da.values) - - @property - def per_perturbation_mean(self): - """ - Mean value for each perturbation type (averaging over all realizations). - - For sequential mode: shape (n_perturbations, filters..., nuclide, score) - - Averages over realizations, keeping separate perturbation results - - For matrix mode: shape (n_perturbations, filters..., nuclide, score) - - Averages over all perturbation_i dimensions (all realization grids) - - Returns one value per perturbation type - - For diagonal mode: returns overall mean (single realization dimension) - """ - dims = tuple(self._da.dims) - has_pert = "perturbation" in dims - has_real = "realization" in dims - - if has_pert and has_real: - # Sequential mode: average over realizations only - result = self._da.mean(dim="realization") - return result.values - elif has_pert: - # Edge case: perturbation dim exists but no realization dim - return self._da.values - else: - # Matrix mode: average over all perturbation_i dimensions - pert_dims = [d for d in self._tmc_dims if d.startswith("perturbation_")] - if pert_dims: - # Average over all perturbation dimensions (all realization grids) - result = self._da.mean(dim=pert_dims) - # Result shape: (filters..., nuclide, score) - # Expand to add perturbation axis: (n_perturbations, filters..., nuclide, score) - n_perturbations = len(pert_dims) - # Repeat the result for each perturbation - result_expanded = np.tile(result.values, (n_perturbations,) + (1,) * (result.ndim)) - return result_expanded - else: - # Diagonal or other: collapse all TMC dims - return self.mean - @property - def per_perturbation_std_dev(self): - """ - Standard deviation for each perturbation type (across all realizations). - - For sequential mode: shape (n_perturbations, filters..., nuclide, score) - - Std deviation across realizations, keeping separate perturbation results - - For matrix mode: shape (n_perturbations, filters..., nuclide, score) - - Std deviation over all perturbation_i dimensions (all realization grids) - - Returns one value per perturbation type - - For diagonal mode: returns overall std_dev (single realization dimension) - """ - dims = tuple(self._da.dims) - has_pert = "perturbation" in dims - has_real = "realization" in dims - - if has_pert and has_real: - # Sequential mode: std over realizations only - result = self._da.std(dim="realization") - return result.values - elif has_pert: - # Edge case: perturbation dim exists but no realization dim - return np.zeros_like(self._da.values) - else: - # Matrix mode: std over all perturbation_i dimensions - pert_dims = [d for d in self._tmc_dims if d.startswith("perturbation_")] - if pert_dims: - # Std over all perturbation dimensions (all realization grids) - result = self._da.std(dim=pert_dims) - # Result shape: (filters..., nuclide, score) - # Expand to add perturbation axis: (n_perturbations, filters..., nuclide, score) - n_perturbations = len(pert_dims) - # Repeat the result for each perturbation - result_expanded = np.tile(result.values, (n_perturbations,) + (1,) * (result.ndim)) - return result_expanded - else: - # Diagonal or other: collapse all TMC dims - return self.std_dev - - - @property - def perturbation_dims(self): - """Names of perturbation dimensions in matrix mode (e.g. 'perturbation_0', ...).""" - return tuple(d for d in self._tmc_dims if d.startswith("perturbation_")) - - def get_slice(self, scores=None, nuclides=None, **filter_kwargs): - """ - Get a slice of the TMC data with optional filtering. - - Parameters - ---------- - scores : list of str, optional - Score names to select - nuclides : list of str, optional - Nuclide names to select - **filter_kwargs : optional - Additional dimension filters (e.g., energy=slice(0, 10)) - - Returns - ------- - xarray.DataArray - Filtered TMC mean data - """ - da = self._da - - # Apply filter dimension selections (energy, cell, mesh, perturbation_x, etc.) - if filter_kwargs: - da = da.sel(**filter_kwargs) - - # Apply score selection - if scores is not None and "score" in da.dims: - all_scores = self.scores - score_indices = [all_scores.index(s) for s in scores] - da = da.isel(score=score_indices) - - # Apply nuclide selection - if nuclides is not None and "nuclide" in da.dims: - all_nuclides = self.nuclides - nuclide_indices = [all_nuclides.index(n) for n in nuclides] - da = da.isel(nuclide=nuclide_indices) - - return da - - - def __repr__(self): - return f"" \ No newline at end of file + return TMCStatePoint(path) \ No newline at end of file diff --git a/src/openmc_fusion_benchmarks/uq/tmc_statepoint.py b/src/openmc_fusion_benchmarks/uq/tmc_statepoint.py new file mode 100644 index 00000000..912ad58e --- /dev/null +++ b/src/openmc_fusion_benchmarks/uq/tmc_statepoint.py @@ -0,0 +1,322 @@ +from pathlib import Path + +import h5py +import numpy as np +import xarray as xr + +from ..tallies import BaseTally + + +def _open_dataset_with_fallback(path: Path, group_name: str) -> xr.Dataset: + """Open an HDF5/NetCDF group using whichever backend is available.""" + last_exc = None + for engine in ("h5netcdf", "netcdf4", None): + try: + if engine is None: + return xr.open_dataset(path, group=group_name) + return xr.open_dataset(path, group=group_name, engine=engine) + except (ModuleNotFoundError, ImportError, OSError, ValueError) as exc: + last_exc = exc + continue + if last_exc is not None: + raise last_exc + raise RuntimeError(f"Could not open dataset group '{group_name}' from {path}") + + +class TMCStatePoint: + """ + Wrapper for TMC statepoint providing an OpenMC StatePoint-like interface. + + Parameters + ---------- + path : str or Path + Path to the TMC statepoint NetCDF/HDF5 file. + """ + + def __init__(self, path): + self.path = Path(path).resolve() + # We won't keep one global ds; we'll open per-tally groups as needed. + self._tallies = None + self._tally_groups = None + self._tallies_by_name = None + + @property + def tallies_by_id(self): + """Dictionary of tallies indexed by tally ID (OpenMC-like access).""" + if self._tallies is None: + self._tallies = {} + self._tally_groups = [] + self._tallies_by_name = {} + # Discover tally groups via h5py + with h5py.File(self.path, "r") as f: + for group_name in f.keys(): + if not isinstance(f[group_name], h5py.Group): + continue + # Open this group as an xarray Dataset + ds = _open_dataset_with_fallback(self.path, group_name) + if "mean" in ds: + mean_var_name = "mean" + elif len(ds.data_vars) == 1: + # Backward-compatibility for legacy single-var groups. + mean_var_name = next(iter(ds.data_vars)) + else: + ds.close() + continue + da_mean = ds[mean_var_name] + da_mc_std = ds["mc_std"] if "mc_std" in ds else None + + tally_id = da_mean.attrs.get("tally_id") + if tally_id is None: + # Fallback: parse id from group name + try: + tally_id = int(group_name.split("_", 1)[1]) + except Exception: + ds.close() + continue + + try: + tally_id = int(tally_id) + except Exception: + ds.close() + continue + + self._tally_groups.append(group_name) + + tally_obj = TMCTally(da_mean, da_mc_std, parent_ds=ds) + self._tallies[tally_id] = tally_obj + + tname = tally_obj.name + if tname: + self._tallies_by_name[str(tname)] = tally_obj + return self._tallies + + @property + def tallies(self): + """List available tally group names (same style as Results.tallies).""" + _ = self.tallies_by_id + return list(self._tally_groups) + + def get_tally(self, tally_id=None, name=None): + """ + Get a tally by ID or name (mimics openmc.StatePoint.get_tally). + + Parameters + ---------- + tally_id : int, optional + Tally ID + name : str, optional + Tally name + + Returns + ------- + TMCTally + The requested tally + """ + # Convenience: allow get_tally("name") while preserving + # existing get_tally() behavior. + if name is None and isinstance(tally_id, str): + name = tally_id + tally_id = None + + # Ensure tallies are loaded (and name index is built). + _ = self.tallies_by_id + + if name is not None: + tally = self._tallies_by_name.get(str(name)) + if tally is not None: + return tally + raise ValueError(f"No tally with name '{name}' found") + + if tally_id is not None: + try: + return self.tallies_by_id[tally_id] + except KeyError: + raise ValueError(f"No tally with id '{tally_id}' found") + + raise ValueError("Must specify either 'name' or 'tally_id'") + + def close(self): + """No persistent open Dataset to close, but keep for API symmetry.""" + # If you decide to cache per-group ds objects, close them here. + pass + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + def __repr__(self): + n_tallies = len(self.tallies_by_id) + # Try to infer "realizations" (or more generally TMC size along first dim) + n_realizations = 0 + if self.tallies_by_id: + any_tally = next(iter(self.tallies_by_id.values())) + tmc_dims = any_tally._tmc_dims + if tmc_dims: + # product of TMC dims sizes + sz = 1 + for d in tmc_dims: + sz *= any_tally._da.sizes[d] + n_realizations = sz + return f"" + + +class TMCTally(BaseTally): + """ + Wrapper for a single TMC tally providing an OpenMC Tally-like interface. + + Parameters + ---------- + mean_da : xarray.DataArray + The DataArray containing the TMC mean values for this tally. + mc_std_da : xarray.DataArray, optional + The DataArray containing the MC std dev per run/combo. + parent_ds : xarray.Dataset, optional + Parent dataset (group) containing metadata attributes. + """ + + def __init__(self, mean_da, mc_std_da=None, parent_ds=None): + super().__init__(mean_da=mean_da, mc_std_da=mc_std_da, parent_ds=parent_ds) + + # Identify TMC dimensions: "perturbation" and "realization" for sequential, "perturbation_*" for matrix + self._tmc_dims = [ + d for d in self._da.dims + if d in ("perturbation", "realization") or d.startswith("perturbation_") + ] + + @property + def tmc_dims(self): + """Names of TMC dimensions (realization / perturbation_*).""" + return tuple(self._tmc_dims) + + @property + def mean(self): + """ + TMC mean across all TMC dimensions (realization / perturbation_*). + """ + if not self._tmc_dims: + return self._da.values + return self._da.mean(dim=self._tmc_dims).values + + @property + def std_dev(self): + """ + TMC standard deviation across all TMC dimensions. + + This is the propagated parametric uncertainty from the ensemble, + not the MC sampling error within each run. + """ + if not self._tmc_dims: + return np.zeros_like(self._da.values) + return self._da.std(dim=self._tmc_dims).values + + @property + def per_realization_mean(self): + """ + Raw mean value for each TMC point (all TMC dims retained). + Shape: (TMC dims..., filters..., nuclide, score). + """ + return self._da.values + + @property + def per_realization_std_dev(self): + """ + Monte Carlo standard deviation for each run/combo. + + This is the statistical uncertainty from particle sampling within each + individual OpenMC run, shaped like self._da. + """ + if self._da_mc_std is not None: + return self._da_mc_std.values + return np.zeros_like(self._da.values) + + @property + def per_perturbation_mean(self): + """ + Mean value for each perturbation type (averaging over all realizations). + + For sequential mode: shape (n_perturbations, filters..., nuclide, score) + - Averages over realizations, keeping separate perturbation results + + For matrix mode: shape (n_perturbations, filters..., nuclide, score) + - Averages over all perturbation_i dimensions (all realization grids) + - Returns one value per perturbation type + + For diagonal mode: returns overall mean (single realization dimension) + """ + dims = tuple(self._da.dims) + has_pert = "perturbation" in dims + has_real = "realization" in dims + + if has_pert and has_real: + # Sequential mode: average over realizations only + result = self._da.mean(dim="realization") + return result.values + elif has_pert: + # Edge case: perturbation dim exists but no realization dim + return self._da.values + else: + # Matrix mode: average over all perturbation_i dimensions + pert_dims = [d for d in self._tmc_dims if d.startswith("perturbation_")] + if pert_dims: + # Average over all perturbation dimensions (all realization grids) + result = self._da.mean(dim=pert_dims) + # Result shape: (filters..., nuclide, score) + # Expand to add perturbation axis: (n_perturbations, filters..., nuclide, score) + n_perturbations = len(pert_dims) + # Repeat the result for each perturbation + result_expanded = np.tile(result.values, (n_perturbations,) + (1,) * (result.ndim)) + return result_expanded + else: + # Diagonal or other: collapse all TMC dims + return self.mean + + @property + def per_perturbation_std_dev(self): + """ + Standard deviation for each perturbation type (across all realizations). + + For sequential mode: shape (n_perturbations, filters..., nuclide, score) + - Std deviation across realizations, keeping separate perturbation results + + For matrix mode: shape (n_perturbations, filters..., nuclide, score) + - Std deviation over all perturbation_i dimensions (all realization grids) + - Returns one value per perturbation type + + For diagonal mode: returns overall std_dev (single realization dimension) + """ + dims = tuple(self._da.dims) + has_pert = "perturbation" in dims + has_real = "realization" in dims + + if has_pert and has_real: + # Sequential mode: std over realizations only + result = self._da.std(dim="realization") + return result.values + elif has_pert: + # Edge case: perturbation dim exists but no realization dim + return np.zeros_like(self._da.values) + else: + # Matrix mode: std over all perturbation_i dimensions + pert_dims = [d for d in self._tmc_dims if d.startswith("perturbation_")] + if pert_dims: + # Std over all perturbation dimensions (all realization grids) + result = self._da.std(dim=pert_dims) + # Result shape: (filters..., nuclide, score) + # Expand to add perturbation axis: (n_perturbations, filters..., nuclide, score) + n_perturbations = len(pert_dims) + # Repeat the result for each perturbation + result_expanded = np.tile(result.values, (n_perturbations,) + (1,) * (result.ndim)) + return result_expanded + else: + # Diagonal or other: collapse all TMC dims + return self.std_dev + + @property + def perturbation_dims(self): + """Names of perturbation dimensions in matrix mode (e.g. 'perturbation_0', ...).""" + return tuple(d for d in self._tmc_dims if d.startswith("perturbation_")) + + def __repr__(self): + return f"" diff --git a/src/openmc_fusion_benchmarks/validate_results.py b/src/openmc_fusion_benchmarks/validate_results.py new file mode 100644 index 00000000..6786becf --- /dev/null +++ b/src/openmc_fusion_benchmarks/validate_results.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import numpy as np + + +def normalize_filter_type(type_name: str) -> str: + """Normalize filter type naming across spec and backend representations.""" + name = str(type_name).strip().lower() + if name.endswith("filter"): + name = name[:-6] + return name + + +def _filter_bins_match(spec_filter: dict, observed_axis: dict) -> bool: + """Compare filter bin definitions from spec and observed metadata.""" + expected = spec_filter.get("values") + observed = observed_axis.get("bins") + + if expected is None or observed is None: + return True + + ftype = normalize_filter_type(spec_filter.get("type", "")) + if ftype == "energy": + try: + return np.allclose(np.asarray(expected, dtype=float), np.asarray(observed, dtype=float)) + except Exception: + return False + + try: + return list(expected) == list(observed) + except Exception: + return False + + +def validate_tally_consistency(spec_tally: dict, observed_tally: dict) -> tuple[bool, list[str]]: + """Validate observed tally metadata against repository specification.""" + issues: list[str] = [] + + observed_scores = [str(s) for s in observed_tally.get("scores", [])] + observed_nuclides = [str(n) for n in observed_tally.get("nuclides", [])] + observed_filters = list(observed_tally.get("filters", [])) + + spec_scores = [str(s) for s in spec_tally.get("scores", [])] + if spec_scores and spec_scores != observed_scores: + issues.append(f"scores mismatch: expected {spec_scores}, observed {observed_scores}") + + spec_nuclides = [str(n) for n in spec_tally.get("nuclides", [])] + if spec_nuclides and spec_nuclides != observed_nuclides: + issues.append(f"nuclides mismatch: expected {spec_nuclides}, observed {observed_nuclides}") + + expected_particle = spec_tally.get("particle") + if expected_particle is not None: + particle_filters = [a for a in observed_filters if normalize_filter_type(a.get("name", "")) == "particle"] + if not particle_filters: + issues.append("missing ParticleFilter in observed tally") + else: + bins = particle_filters[0].get("bins", []) + observed_particle = str(bins[0]) if bins else None + if str(expected_particle) != str(observed_particle): + issues.append( + f"particle mismatch: expected {expected_particle}, observed {observed_particle}" + ) + + spec_filters = spec_tally.get("filters", []) + observed_non_particle = [ + a for a in observed_filters if normalize_filter_type(a.get("name", "")) != "particle" + ] + + expected_types = [normalize_filter_type(f.get("type", "")) for f in spec_filters] + observed_types = [normalize_filter_type(a.get("name", "")) for a in observed_non_particle] + if expected_types != observed_types: + issues.append(f"filter type/order mismatch: expected {expected_types}, observed {observed_types}") + + if len(spec_filters) == len(observed_non_particle): + for i, (spec_filter, observed_axis) in enumerate(zip(spec_filters, observed_non_particle)): + ftype = normalize_filter_type(spec_filter.get("type", "")) + if ftype == "energy": + expected_closure = spec_filter.get("closure", "[low, high)") + observed_closure = observed_axis.get("closure", "[low, high)") + if expected_closure != observed_closure: + issues.append( + "energy closure mismatch at index " + f"{i}: expected {expected_closure}, observed {observed_closure}" + ) + + if not _filter_bins_match(spec_filter, observed_axis): + issues.append( + f"filter bins mismatch at index {i} ({spec_filter.get('type')}): " + f"expected {spec_filter.get('values')}, observed {observed_axis.get('bins')}" + ) + + return len(issues) == 0, issues diff --git a/src/openmc_fusion_benchmarks/validate.py b/src/openmc_fusion_benchmarks/validate_spec.py similarity index 75% rename from src/openmc_fusion_benchmarks/validate.py rename to src/openmc_fusion_benchmarks/validate_spec.py index dbd51933..0c8e1bba 100644 --- a/src/openmc_fusion_benchmarks/validate.py +++ b/src/openmc_fusion_benchmarks/validate_spec.py @@ -23,20 +23,20 @@ def validate_benchmark(benchmark_name: str): registry = Registry().with_resources( [(schema.get("$id", "benchmark_schema"), schema)]) - # CHANGED: Wrap the schema in a Resource and use the referencing registry + # Wrap the schema in a Resource and use the referencing registry schema_id = schema.get( - "$id", "https://openmc-fusion/schemas/benchmark_schema") # <-- CHANGED + "$id", "https://openmc-fusion/schemas/benchmark_schema") registry = Registry().with_resources( - [(schema_id, Resource.from_contents(schema))]) # <-- CHANGED + [(schema_id, Resource.from_contents(schema))]) # Load the YAML file to validate with open(benchmark_path, "r") as yaml_file: yaml_data = yaml.safe_load(yaml_file) - # CHANGED: Use jsonschema 2020-12 validator with registry - validator_cls = jsonschema.validators.validator_for(schema) # <-- ADDED - validator_cls.check_schema(schema) # <-- ADDED - validator = validator_cls(schema, registry=registry) # <-- CHANGED + # Use jsonschema 2020-12 validator with registry + validator_cls = jsonschema.validators.validator_for(schema) + validator_cls.check_schema(schema) + validator = validator_cls(schema, registry=registry) # Validate the YAML file errors = sorted(validator.iter_errors(yaml_data), key=lambda e: e.path) diff --git a/test/test_benchmark_results_report_format.py b/test/test_benchmark_results_report_format.py new file mode 100644 index 00000000..937bb902 --- /dev/null +++ b/test/test_benchmark_results_report_format.py @@ -0,0 +1,58 @@ +import json +from pathlib import Path + +import xarray as xr +import numpy as np + +from openmc_fusion_benchmarks.benchmark_results import Results + + +def _write_group(path: Path, group: str, consistent: int, issues): + ds = xr.Dataset( + { + "mean": xr.DataArray( + np.ones((1, 1), dtype=float), + dims=("nuclide", "score"), + coords={ + "nuclide": np.array(["total"], dtype="U"), + "score": np.array(["current"], dtype="U"), + }, + ), + "mc_std": xr.DataArray( + np.full((1, 1), 0.1, dtype=float), + dims=("nuclide", "score"), + coords={ + "nuclide": np.array(["total"], dtype="U"), + "score": np.array(["current"], dtype="U"), + }, + ), + } + ) + + ds["mean"].attrs["tally_id"] = 1 + ds["mean"].attrs["tally_name"] = group + ds["mc_std"].attrs["tally_id"] = 1 + ds["mc_std"].attrs["tally_name"] = group + + ds.attrs["spec_consistent"] = int(consistent) + ds.attrs["spec_consistency_issues"] = json.dumps(issues) + ds.attrs["observed_tally"] = json.dumps({"name": group}) + + mode = "a" if path.exists() else "w" + ds.to_netcdf(path, mode=mode, group=group, engine="h5netcdf") + + +def test_format_spec_consistency_report(tmp_path): + p = tmp_path / "benchmark_results.h5" + _write_group(p, "neutron_leakage", 1, []) + _write_group(p, "photon_leakage", 0, ["scores mismatch"]) + + r = Results.from_file(p) + text = r.format_spec_consistency_report() + + assert "Spec Consistency Report" in text + assert "Tally neutron_leakage (group=neutron_leakage)" in text + assert "Status" in text and "OK" in text + assert "Tally photon_leakage (group=photon_leakage)" in text + assert "MISMATCH" in text + assert "scores mismatch" in text diff --git a/test/test_tmc_statepoint.py b/test/test_tmc_statepoint.py new file mode 100644 index 00000000..7eac073a --- /dev/null +++ b/test/test_tmc_statepoint.py @@ -0,0 +1,80 @@ +from pathlib import Path + +import numpy as np +import pytest +import xarray as xr + +from openmc_fusion_benchmarks.uq.tmc_statepoint import TMCStatePoint + + +def _write_tally_group(path: Path, group: str, tally_id: int, tally_name: str): + data = np.ones((2, 1, 1), dtype=float) + ds = xr.Dataset( + { + "mean": xr.DataArray( + data, + dims=("realization", "nuclide", "score"), + coords={ + "realization": np.arange(2), + "nuclide": np.array(["total"], dtype="U"), + "score": np.array(["current"], dtype="U"), + }, + ), + "mc_std": xr.DataArray( + np.full_like(data, 0.1), + dims=("realization", "nuclide", "score"), + coords={ + "realization": np.arange(2), + "nuclide": np.array(["total"], dtype="U"), + "score": np.array(["current"], dtype="U"), + }, + ), + } + ) + ds["mean"].attrs["tally_id"] = tally_id + ds["mean"].attrs["tally_name"] = tally_name + ds["mc_std"].attrs["tally_id"] = tally_id + ds["mc_std"].attrs["tally_name"] = tally_name + + mode = "a" if path.exists() else "w" + ds.to_netcdf(path, mode=mode, group=group, engine="h5netcdf") + + +def test_tmc_statepoint_discovers_name_groups_and_gets_by_name(tmp_path): + fpath = tmp_path / "tmc_statepoint.2.h5" + _write_tally_group(fpath, group="neutron_leakage", tally_id=1, tally_name="neutron_leakage") + _write_tally_group(fpath, group="photon_leakage", tally_id=2, tally_name="photon_leakage") + + sp = TMCStatePoint(fpath) + + assert set(sp.tallies) == {"neutron_leakage", "photon_leakage"} + assert set(sp.tallies_by_id.keys()) == {1, 2} + + t1 = sp.get_tally(name="neutron_leakage") + assert t1.id == 1 + assert t1.name == "neutron_leakage" + + # Positional string convenience should resolve to name lookup. + t2 = sp.get_tally("photon_leakage") + assert t2.id == 2 + assert t2.name == "photon_leakage" + + # ID lookup remains supported. + t_by_id = sp.get_tally(tally_id=1) + assert t_by_id.name == "neutron_leakage" + + +def test_tmc_statepoint_get_tally_errors(tmp_path): + fpath = tmp_path / "tmc_statepoint.1.h5" + _write_tally_group(fpath, group="neutron_leakage", tally_id=11, tally_name="neutron_leakage") + + sp = TMCStatePoint(fpath) + + with pytest.raises(ValueError, match="name 'missing'"): + sp.get_tally(name="missing") + + with pytest.raises(ValueError, match="id '99'"): + sp.get_tally(tally_id=99) + + with pytest.raises(ValueError, match="Must specify either 'name' or 'tally_id'"): + sp.get_tally() diff --git a/test/test_validate.py b/test/test_validate.py index 01dc7271..74385025 100644 --- a/test/test_validate.py +++ b/test/test_validate.py @@ -2,7 +2,7 @@ import yaml from pathlib import Path from unittest.mock import patch, mock_open, MagicMock -from openmc_fusion_benchmarks import validate_benchmark +from openmc_fusion_benchmarks.validate_spec import validate_benchmark @pytest.fixture @@ -24,9 +24,9 @@ def test_file_not_found(mock_paths): validate_benchmark("dummy_benchmark") -@patch("openmc_fusion_benchmarks.validate.yaml.safe_load") -@patch("openmc_fusion_benchmarks.validate.open") -@patch("openmc_fusion_benchmarks.validate.Path.is_file") +@patch("openmc_fusion_benchmarks.validate_spec.yaml.safe_load") +@patch("openmc_fusion_benchmarks.validate_spec.open") +@patch("openmc_fusion_benchmarks.validate_spec.Path.is_file") def test_valid_yaml_schema_validation_passes(mock_is_file, mock_open_file, mock_safe_load, capsys): # Arrange mock_is_file.return_value = True @@ -59,9 +59,9 @@ def test_valid_yaml_schema_validation_passes(mock_is_file, mock_open_file, mock_ assert "✅ dummy_benchmark is valid!" in out -@patch("openmc_fusion_benchmarks.validate.yaml.safe_load") -@patch("openmc_fusion_benchmarks.validate.open") -@patch("openmc_fusion_benchmarks.validate.Path.is_file") +@patch("openmc_fusion_benchmarks.validate_spec.yaml.safe_load") +@patch("openmc_fusion_benchmarks.validate_spec.open") +@patch("openmc_fusion_benchmarks.validate_spec.Path.is_file") def test_invalid_yaml_schema_validation_fails(mock_is_file, mock_open_file, mock_safe_load, capsys): # Arrange mock_is_file.return_value = True