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/src/openmc_fusion_benchmarks/backends/__init__.py b/src/openmc_fusion_benchmarks/backends/__init__.py new file mode 100644 index 00000000..64331ab4 --- /dev/null +++ b/src/openmc_fusion_benchmarks/backends/__init__.py @@ -0,0 +1 @@ +"""Backend-specific interfaces (OpenMC, Serpent, MCNP, ...).""" \ No newline at end of file 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..859c10ba --- /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 * \ No newline at end of file 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..6eacfac6 --- /dev/null +++ b/src/openmc_fusion_benchmarks/backends/openmc/tallies.py @@ -0,0 +1,230 @@ +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 + + +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 _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, +) -> 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 + + ds.attrs["filter_axes"] = json.dumps( + [{"name": type(f).__name__, "num_bins": int(f.num_bins)} for f in filters] + ) + ds.attrs["nuclides"] = json.dumps(nuclides) + ds.attrs["scores"] = json.dumps(scores) + + return ds + + +def save_openmc_statepoint_tallies( + statepoint: openmc.StatePoint, + filename: str | Path, + tally_names: Iterable[str] | 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) + 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, + ) + 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 \ No newline at end of file diff --git a/src/openmc_fusion_benchmarks/benchmark.py b/src/openmc_fusion_benchmarks/benchmark.py index e4725f6a..77039a37 100644 --- a/src/openmc_fusion_benchmarks/benchmark.py +++ b/src/openmc_fusion_benchmarks/benchmark.py @@ -3,10 +3,12 @@ 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 .backends.openmc.tallies import ( + make_default_openmc_normalizer, + save_openmc_statepoint_tallies, +) from .uq.tmc_engine import tmc_engine import openmc @@ -401,16 +403,35 @@ 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 - ) + tally_names = [t["name"] for t in tallies_data] + normalizer = make_default_openmc_normalizer(mesh) + + # Accept both already-open StatePoint objects and statepoint file paths. + if isinstance(statepoint, openmc.StatePoint): + sp = statepoint + save_openmc_statepoint_tallies( + statepoint=sp, + filename="benchmark_results.h5", + tally_names=tally_names, + tmc_coords={"realization": ["baseline"]}, + append_dim="realization", + normalizer=normalizer, + ) + else: + with openmc.StatePoint(str(statepoint)) as sp: + save_openmc_statepoint_tallies( + statepoint=sp, + filename="benchmark_results.h5", + tally_names=tally_names, + tmc_coords={"realization": ["baseline"]}, + append_dim="realization", + normalizer=normalizer, + ) return diff --git a/src/openmc_fusion_benchmarks/benchmark_results.py b/src/openmc_fusion_benchmarks/benchmark_results.py index aa73c073..c3f03d6e 100644 --- a/src/openmc_fusion_benchmarks/benchmark_results.py +++ b/src/openmc_fusion_benchmarks/benchmark_results.py @@ -7,12 +7,17 @@ from .database import _resolve_database_path -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)` +from .tallies import Tally + + +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. + - 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 +26,75 @@ 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 xr.open_dataset(self.filepath, group=candidate, engine="h5netcdf") as ds: + if "mean" not in ds: + continue + if ds["mean"].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 = xr.open_dataset(self.filepath, group=group, engine="h5netcdf") + if "mean" not in ds: + 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) + + +class BenchmarkResults(Results): + """Backward-compatible alias class for benchmark-centric naming.""" + + +# Explicit alias for OFB naming style. +OFBResults = Results \ No newline at end of file diff --git a/src/openmc_fusion_benchmarks/tallies.py b/src/openmc_fusion_benchmarks/tallies.py new file mode 100644 index 00000000..9e5cc541 --- /dev/null +++ b/src/openmc_fusion_benchmarks/tallies.py @@ -0,0 +1,135 @@ +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 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/test/test_benchmark_comprehensive.py b/test/test_benchmark_comprehensive.py index 020b7847..fbdcde24 100644 --- a/test/test_benchmark_comprehensive.py +++ b/test/test_benchmark_comprehensive.py @@ -510,10 +510,16 @@ def test_postprocess(): with patch.object(OpenmcBenchmark, '_build_settings', return_value=openmc.Settings()): bench = OpenmcBenchmark("test") - mock_sp = Mock() - with patch("openmc_fusion_benchmarks.benchmark._openmc_to_ofb") as mock_post: - bench._postprocess(statepoint=mock_sp, mesh="test.h5m") - mock_post.assert_called_once() + class DummyStatePoint: + pass + + with patch("openmc_fusion_benchmarks.benchmark.openmc.StatePoint", new=DummyStatePoint): + mock_sp = DummyStatePoint() + with patch("openmc_fusion_benchmarks.benchmark.make_default_openmc_normalizer", return_value="norm") as mock_norm: + with patch("openmc_fusion_benchmarks.benchmark.save_openmc_statepoint_tallies") as mock_post: + bench._postprocess(statepoint=mock_sp, mesh="test.h5m") + mock_norm.assert_called_once_with("test.h5m") + mock_post.assert_called_once() def test_run_without_uq(): diff --git a/test/test_benchmark_postprocess.py b/test/test_benchmark_postprocess.py new file mode 100644 index 00000000..babcf88a --- /dev/null +++ b/test/test_benchmark_postprocess.py @@ -0,0 +1,103 @@ +import importlib.util +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock, patch + + +def _module_available(name: str) -> bool: + if name in sys.modules: + return True + try: + return importlib.util.find_spec(name) is not None + except ValueError: + return False + + +# The package imports benchmark.py at import time; provide stubs when unavailable. +if not _module_available("openmc"): + openmc_stub = types.ModuleType("openmc") + openmc_stub.__path__ = [] + for cls_name in ( + "StatePoint", + "Tally", + "Filter", + "CellFilter", + "SurfaceFilter", + "MaterialFilter", + "EnergyFilter", + "ParticleFilter", + "Materials", + "Material", + "Geometry", + "Settings", + "Tallies", + "Model", + "DAGMCUniverse", + ): + setattr(openmc_stub, cls_name, type(cls_name, (), {})) + sys.modules.setdefault("openmc", openmc_stub) + + openmc_data_stub = types.ModuleType("openmc.data") + openmc_data_stub.zam = lambda _name: (1, 1, 0) + sys.modules.setdefault("openmc.data", openmc_data_stub) + +try: + import pydagmc as _pydagmc # noqa: F401 +except Exception: + pydagmc_stub = types.ModuleType("pydagmc") + pydagmc_stub.Model = type("Model", (), {}) + sys.modules["pydagmc"] = pydagmc_stub + +try: + import cad_to_dagmc as _cad_to_dagmc # noqa: F401 +except Exception: + cad_stub = types.ModuleType("cad_to_dagmc") + cad_stub.CadToDagmc = type("CadToDagmc", (), {}) + sys.modules["cad_to_dagmc"] = cad_stub + +if not _module_available("sandy"): + sys.modules.setdefault("sandy", types.ModuleType("sandy")) + + +from openmc_fusion_benchmarks.benchmark import OpenmcBenchmark + + +def test_postprocess_with_open_statepoint_object(): + fake_self = SimpleNamespace(_benchmark_spec={"tallies": [{"name": "t1"}, {"name": "t2"}]}) + fake_sp = object() + + with patch("openmc_fusion_benchmarks.benchmark.make_default_openmc_normalizer", return_value="norm") as mk_norm: + with patch("openmc_fusion_benchmarks.benchmark.save_openmc_statepoint_tallies") as save: + with patch("openmc_fusion_benchmarks.benchmark.openmc.StatePoint", new=object): + OpenmcBenchmark._postprocess(fake_self, statepoint=fake_sp, mesh="mesh.h5m") + + mk_norm.assert_called_once_with("mesh.h5m") + save.assert_called_once() + kwargs = save.call_args.kwargs + assert kwargs["statepoint"] is fake_sp + assert kwargs["filename"] == "benchmark_results.h5" + assert kwargs["tally_names"] == ["t1", "t2"] + assert kwargs["tmc_coords"] == {"realization": ["baseline"]} + + +def test_postprocess_with_statepoint_path_uses_context_manager(): + fake_self = SimpleNamespace(_benchmark_spec={"tallies": [{"name": "t1"}]}) + + class DummyStatePoint: + def __init__(self, *_args, **_kwargs): + pass + + def __enter__(self): + return SimpleNamespace(get_tally=lambda *a, **k: None) + + def __exit__(self, *_exc): + return False + + with patch("openmc_fusion_benchmarks.benchmark.make_default_openmc_normalizer", return_value="norm"): + with patch("openmc_fusion_benchmarks.benchmark.save_openmc_statepoint_tallies") as save: + with patch("openmc_fusion_benchmarks.benchmark.openmc.StatePoint", new=DummyStatePoint): + OpenmcBenchmark._postprocess(fake_self, statepoint=Path("statepoint.10.h5")) + + save.assert_called_once() diff --git a/test/test_benchmark_results.py b/test/test_benchmark_results.py index bf2a606a..d27472ea 100644 --- a/test/test_benchmark_results.py +++ b/test/test_benchmark_results.py @@ -1,84 +1,169 @@ +import importlib.util +import json +import sys +import types + +import numpy as np import pytest -import h5py import xarray as xr -import numpy as np -from pathlib import Path -from openmc_fusion_benchmarks.benchmark_results import BenchmarkResults -from openmc_fusion_benchmarks.database import list_database_benchmarks, list_database_files + + +def _module_available(name: str) -> bool: + if name in sys.modules: + return True + try: + return importlib.util.find_spec(name) is not None + except ValueError: + return False + + +# The package __init__ imports benchmark.py, which imports these dependencies. +# Provide minimal stubs in environments where OpenMC stack is not available. +if not _module_available("openmc"): + openmc_stub = types.ModuleType("openmc") + openmc_stub.__path__ = [] + for cls_name in ( + "StatePoint", + "Tally", + "Filter", + "CellFilter", + "SurfaceFilter", + "MaterialFilter", + "EnergyFilter", + "ParticleFilter", + "Materials", + "Material", + "Geometry", + "Settings", + "Tallies", + "Model", + "DAGMCUniverse", + ): + setattr(openmc_stub, cls_name, type(cls_name, (), {})) + sys.modules.setdefault("openmc", openmc_stub) + + openmc_data_stub = types.ModuleType("openmc.data") + openmc_data_stub.zam = lambda _name: (1, 1, 0) + sys.modules.setdefault("openmc.data", openmc_data_stub) + +try: + import pydagmc as _pydagmc # noqa: F401 +except Exception: + pydagmc_stub = types.ModuleType("pydagmc") + pydagmc_stub.Model = type("Model", (), {}) + sys.modules["pydagmc"] = pydagmc_stub + +try: + import cad_to_dagmc as _cad_to_dagmc # noqa: F401 +except Exception: + cad_stub = types.ModuleType("cad_to_dagmc") + cad_stub.CadToDagmc = type("CadToDagmc", (), {}) + sys.modules["cad_to_dagmc"] = cad_stub + +if not _module_available("sandy"): + sys.modules.setdefault("sandy", types.ModuleType("sandy")) + + +from openmc_fusion_benchmarks.benchmark_results import BenchmarkResults, OFBResults, Results +from openmc_fusion_benchmarks.tallies import Tally + + +def _write_structured_group(filepath, group="test_tally", tally_name="mytally"): + ds = xr.Dataset( + { + "mean": xr.DataArray( + np.arange(4.0).reshape(2, 1, 2), + dims=("cell", "nuclide", "score"), + coords={ + "cell": [0, 1], + "nuclide": np.array(["total"], dtype="U"), + "score": np.array(["flux", "heating"], dtype="U"), + }, + ), + "mc_std": xr.DataArray( + np.full((2, 1, 2), 0.1), + dims=("cell", "nuclide", "score"), + coords={ + "cell": [0, 1], + "nuclide": np.array(["total"], dtype="U"), + "score": np.array(["flux", "heating"], dtype="U"), + }, + ), + } + ) + ds["mean"].attrs["tally_id"] = 1 + ds["mean"].attrs["tally_name"] = tally_name + ds["mc_std"].attrs["tally_id"] = 1 + ds["mc_std"].attrs["tally_name"] = tally_name + ds.attrs["observed_tally"] = json.dumps({"name": tally_name}) + ds.to_netcdf(filepath, mode="w", engine="h5netcdf", group=group) @pytest.fixture def temp_results_file(tmp_path): - """Create a temporary HDF5 file with test data.""" filepath = tmp_path / "test_results.h5" - - # Create sample xarray DataArray - data = xr.DataArray( - np.random.rand(3, 2), - dims=["row", "column"], - coords={"row": [0, 1, 2], "column": ["mean", "std. dev."]}, - name="test_tally" - ) - - # Save to HDF5 - data.to_netcdf(filepath, mode="w", engine="netcdf4", group="test_tally") - + _write_structured_group(filepath, group="test_tally", tally_name="mytally") return filepath def test_benchmark_results_from_file(temp_results_file): - """Test loading results from a file path.""" results = BenchmarkResults.from_file(temp_results_file) assert results.filepath == temp_results_file assert results.filepath.exists() def test_benchmark_results_file_not_found(): - """Test that loading a nonexistent file raises FileNotFoundError.""" with pytest.raises(FileNotFoundError): BenchmarkResults.from_file("/nonexistent/path/results.h5") def test_benchmark_results_from_run_dir(tmp_path): - """Test loading results from a run directory.""" - # Create a results file in the temp directory filepath = tmp_path / "results.h5" - data = xr.DataArray( - np.random.rand(2, 2), - dims=["row", "column"], - name="tally" - ) - data.to_netcdf(filepath, mode="w", engine="netcdf4", group="tally") - - # Load from run directory + _write_structured_group(filepath, group="tally", tally_name="tally") + results = BenchmarkResults.from_run_dir(tmp_path, "results.h5") assert results.filepath.exists() assert results.filepath.name == "results.h5" +def test_results_aliases(temp_results_file): + base = Results.from_file(temp_results_file) + alias = OFBResults.from_file(temp_results_file) + assert isinstance(base, Results) + assert isinstance(alias, Results) + + def test_benchmark_results_tallies_property(temp_results_file): - """Test the tallies property returns list of tally names.""" results = BenchmarkResults.from_file(temp_results_file) - tallies = results.tallies - assert isinstance(tallies, list) - assert "test_tally" in tallies + assert results.tallies == ["test_tally"] -def test_benchmark_results_get_tally(temp_results_file): - """Test retrieving a specific tally.""" +def test_get_tally_by_group_name(temp_results_file): results = BenchmarkResults.from_file(temp_results_file) tally = results.get_tally("test_tally") - assert isinstance(tally, xr.DataArray) - assert tally.name == "test_tally" - - -def test_benchmark_results_from_database(): - """Test loading results from the package database.""" - benchmarks = list_database_benchmarks() - if benchmarks: - files = list_database_files(benchmarks[0]) - if files: - # Try to load the first available file - results = BenchmarkResults.from_database(benchmarks[0], files[0]) - assert results.filepath.exists() - assert results.filepath.suffix == '.h5' + assert isinstance(tally, Tally) + assert tally.name == "mytally" + assert tally.id == 1 + + +def test_get_tally_by_logical_name(temp_results_file): + results = BenchmarkResults.from_file(temp_results_file) + tally = results.get_tally("mytally") + assert isinstance(tally, Tally) + assert tally.name == "mytally" + + +def test_get_tally_missing_mean_raises(tmp_path): + filepath = tmp_path / "bad_results.h5" + ds = xr.Dataset({"only_var": xr.DataArray(np.arange(3.0), dims=("row",))}) + ds.to_netcdf(filepath, mode="w", engine="h5netcdf", group="bad") + + results = BenchmarkResults.from_file(filepath) + with pytest.raises(ValueError, match="does not contain a 'mean' dataset"): + results.get_tally("bad") + + +def test_get_tally_unknown_name_raises(temp_results_file): + results = BenchmarkResults.from_file(temp_results_file) + with pytest.raises(ValueError, match="No tally with name or group"): + results.get_tally("missing") diff --git a/test/test_openmc_backend_tallies.py b/test/test_openmc_backend_tallies.py new file mode 100644 index 00000000..bd53a4b2 --- /dev/null +++ b/test/test_openmc_backend_tallies.py @@ -0,0 +1,192 @@ +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import xarray as xr + +openmc = pytest.importorskip("openmc") + +from openmc_fusion_benchmarks.backends.openmc import tallies as backend + + +def _make_filter(filter_cls, bins, num_bins=None): + """Create filter objects for both real OpenMC and lightweight stubs.""" + try: + flt = filter_cls(bins) + # Real OpenMC objects expose read-only num_bins; do not overwrite. + return flt + except TypeError: + flt = filter_cls() + flt.bins = bins + flt.num_bins = int(num_bins if num_bins is not None else len(np.asarray(bins).reshape(-1))) + return flt + + +class DummyStatePoint: + def __init__(self, tallies): + self._selected = list(tallies) + self.tallies = {i + 1: t for i, t in enumerate(self._selected)} + + def get_tally(self, name): + for tally in self._selected: + if tally.name == name: + return tally + raise ValueError(name) + + +class DummyTally: + def __init__(self, tally_id, name, filters, nuclides, scores, mean_nd, std_nd): + self.id = tally_id + self.name = name + self.filters = filters + self.nuclides = nuclides + self.scores = scores + self.mean = np.asarray(mean_nd).reshape(-1) + self.std_dev = np.asarray(std_nd).reshape(-1) + + +def test_unique_filter_dims_and_coord_helper(): + filters = [ + _make_filter(openmc.CellFilter, [1], num_bins=1), + _make_filter(openmc.CellFilter, [2], num_bins=1), + _make_filter(openmc.EnergyFilter, [0.0, 1.0], num_bins=1), + ] + assert backend._unique_filter_dims(filters) == ["cell", "cell_1", "energy"] + + scalar = backend._to_1d_coord(7) + assert scalar.shape == (1,) + assert scalar[0] == 7 + + +def test_make_default_openmc_normalizer_cell_and_surface(): + mesh = SimpleNamespace( + volumes_by_id={1: SimpleNamespace(volume=2.0), 2: SimpleNamespace(volume=4.0)}, + surfaces_by_id={10: SimpleNamespace(area=5.0)}, + ) + normalizer = backend.make_default_openmc_normalizer(mesh) + + tally = DummyTally( + tally_id=1, + name="norm", + filters=[ + _make_filter(openmc.CellFilter, [1, 2], num_bins=2), + _make_filter(openmc.SurfaceFilter, [10], num_bins=1), + ], + nuclides=["total"], + scores=["flux"], + mean_nd=np.ones((2, 1, 1)), + std_nd=np.ones((2, 1, 1)), + ) + + mean, std = normalizer(tally, np.ones((2, 1, 1)), np.ones((2, 1, 1))) + np.testing.assert_allclose(mean[:, 0, 0], np.array([1.0 / 10.0, 1.0 / 20.0])) + np.testing.assert_allclose(std[:, 0, 0], np.array([1.0 / 10.0, 1.0 / 20.0])) + + +def test_make_default_openmc_normalizer_material_filter_raises(): + mesh = SimpleNamespace(volumes_by_id={}, surfaces_by_id={}) + normalizer = backend.make_default_openmc_normalizer(mesh) + tally = DummyTally( + tally_id=1, + name="mat", + filters=[_make_filter(openmc.MaterialFilter, [1], num_bins=1)], + nuclides=["total"], + scores=["flux"], + mean_nd=np.ones((1, 1, 1)), + std_nd=np.ones((1, 1, 1)), + ) + + with pytest.raises(NotImplementedError, match="Material filter normalization"): + normalizer(tally, np.ones((1, 1, 1)), np.ones((1, 1, 1))) + + +def test_make_default_openmc_normalizer_zero_factor_raises(): + mesh = SimpleNamespace( + volumes_by_id={1: SimpleNamespace(volume=0.0)}, + surfaces_by_id={}, + ) + normalizer = backend.make_default_openmc_normalizer(mesh) + tally = DummyTally( + tally_id=1, + name="zero", + filters=[_make_filter(openmc.CellFilter, [1], num_bins=1)], + nuclides=["total"], + scores=["flux"], + mean_nd=np.ones((1, 1, 1)), + std_nd=np.ones((1, 1, 1)), + ) + + with pytest.raises(ValueError, match="Normalization factor contains zero"): + normalizer(tally, np.ones((1, 1, 1)), np.ones((1, 1, 1))) + + +def test_openmc_tally_to_dataset_and_save_append(tmp_path): + filters = [ + _make_filter(openmc.ParticleFilter, ["neutron"], num_bins=1), + _make_filter(openmc.CellFilter, [1, 2], num_bins=2), + ] + mean_nd = np.arange(4.0).reshape(1, 2, 1, 2) + std_nd = np.full_like(mean_nd, 0.2) + tally = DummyTally( + tally_id=7, + name="cell_flux", + filters=filters, + nuclides=["total"], + scores=["flux", "heating"], + mean_nd=mean_nd, + std_nd=std_nd, + ) + + ds = backend.openmc_tally_to_dataset( + tally=tally, + tmc_coords={"realization": ["r0"]}, + normalizer=lambda _t, m, s: (m * 2.0, s * 2.0), + ) + assert tuple(ds["mean"].dims) == ("realization", "particle", "cell", "nuclide", "score") + assert ds["mean"].attrs["tally_name"] == "cell_flux" + + sp = DummyStatePoint([tally]) + fpath = tmp_path / "results.h5" + backend.save_openmc_statepoint_tallies( + statepoint=sp, + filename=fpath, + tally_names=["cell_flux"], + tmc_coords={"realization": ["r0"]}, + append_dim="realization", + ) + backend.save_openmc_statepoint_tallies( + statepoint=sp, + filename=fpath, + tally_names=["cell_flux"], + tmc_coords={"realization": ["r1"]}, + append_dim="realization", + ) + + loaded = xr.open_dataset(fpath, group="cell_flux", engine="h5netcdf") + assert int(loaded.sizes["realization"]) == 2 + loaded.close() + + +def test_save_openmc_statepoint_tallies_group_by_id_and_validation(tmp_path): + tally = DummyTally( + tally_id=42, + name="", + filters=[_make_filter(openmc.CellFilter, [1], num_bins=1)], + nuclides=["total"], + scores=["flux"], + mean_nd=np.ones((1, 1, 1)), + std_nd=np.ones((1, 1, 1)), + ) + sp = DummyStatePoint([tally]) + fpath = tmp_path / "id_results.h5" + + out = backend.save_openmc_statepoint_tallies(statepoint=sp, filename=fpath, group_by="id") + assert Path(out).exists() + + loaded = xr.open_dataset(fpath, group="tally_42", engine="h5netcdf") + assert loaded.attrs["group"] == "tally_42" + loaded.close() + + with pytest.raises(ValueError, match="group_by"): + backend.save_openmc_statepoint_tallies(statepoint=sp, filename=fpath, group_by="bad") diff --git a/test/test_tallies.py b/test/test_tallies.py new file mode 100644 index 00000000..95b03323 --- /dev/null +++ b/test/test_tallies.py @@ -0,0 +1,119 @@ +import importlib.util +import json +import sys +import types + +import numpy as np +import xarray as xr + + +def _module_available(name: str) -> bool: + if name in sys.modules: + return True + try: + return importlib.util.find_spec(name) is not None + except ValueError: + return False + + +# Keep imports stable when OpenMC stack is unavailable locally. +if not _module_available("openmc"): + openmc_stub = types.ModuleType("openmc") + openmc_stub.__path__ = [] + for cls_name in ( + "StatePoint", + "Tally", + "Filter", + "CellFilter", + "SurfaceFilter", + "MaterialFilter", + "EnergyFilter", + "ParticleFilter", + "Materials", + "Material", + "Geometry", + "Settings", + "Tallies", + "Model", + "DAGMCUniverse", + ): + setattr(openmc_stub, cls_name, type(cls_name, (), {})) + sys.modules.setdefault("openmc", openmc_stub) + + openmc_data_stub = types.ModuleType("openmc.data") + openmc_data_stub.zam = lambda _name: (1, 1, 0) + sys.modules.setdefault("openmc.data", openmc_data_stub) + +try: + import pydagmc as _pydagmc # noqa: F401 +except Exception: + pydagmc_stub = types.ModuleType("pydagmc") + pydagmc_stub.Model = type("Model", (), {}) + sys.modules["pydagmc"] = pydagmc_stub + +try: + import cad_to_dagmc as _cad_to_dagmc # noqa: F401 +except Exception: + cad_stub = types.ModuleType("cad_to_dagmc") + cad_stub.CadToDagmc = type("CadToDagmc", (), {}) + sys.modules["cad_to_dagmc"] = cad_stub + +if not _module_available("sandy"): + sys.modules.setdefault("sandy", types.ModuleType("sandy")) + + +from openmc_fusion_benchmarks.tallies import BaseTally, Tally + + +def test_base_tally_properties_and_slice_from_coords(): + mean = xr.DataArray( + np.arange(8.0).reshape(2, 2, 1, 2), + dims=("realization", "cell", "nuclide", "score"), + coords={ + "realization": ["r0", "r1"], + "cell": [10, 20], + "nuclide": np.array(["total"], dtype="U"), + "score": np.array(["flux", "heating"], dtype="U"), + }, + ) + mean.attrs["tally_id"] = 11 + mean.attrs["tally_name"] = "my_tally" + + std = xr.DataArray(np.full(mean.shape, 0.2), dims=mean.dims, coords=mean.coords) + tally = BaseTally(mean, std) + + assert tally.id == 11 + assert tally.name == "my_tally" + assert tally.shape == (2, 2, 1, 2) + assert tally.dims == ("realization", "cell", "nuclide", "score") + assert tally.scores == ["flux", "heating"] + assert tally.nuclides == ["total"] + + sliced = tally.get_slice(scores=["heating"], nuclides=["total"], cell=20) + assert sliced.shape == (2, 1, 1) + + +def test_base_tally_uses_parent_dataset_attrs_and_std_default(): + mean = xr.DataArray(np.ones((1, 1)), dims=("i", "j")) + mean.attrs["scores"] = json.dumps(["score_attr"]) + mean.attrs["nuclides"] = json.dumps(["n_attr"]) + mean.attrs["filter_axes"] = json.dumps([{"name": "CellFilter", "num_bins": 1}]) + + parent = xr.Dataset() + parent.attrs["scores"] = json.dumps(["score_parent"]) + parent.attrs["nuclides"] = json.dumps(["n_parent"]) + parent.attrs["filter_axes"] = json.dumps([{"name": "SurfaceFilter", "num_bins": 2}]) + + tally = BaseTally(mean, parent_ds=parent) + assert tally.scores == ["score_parent"] + assert tally.nuclides == ["n_parent"] + assert tally.filters == [{"name": "SurfaceFilter", "num_bins": 2}] + np.testing.assert_allclose(tally.std_dev, np.zeros_like(mean.values)) + + +def test_tally_repr(): + mean = xr.DataArray(np.ones((1,)), dims=("i",)) + mean.attrs["tally_id"] = 3 + mean.attrs["tally_name"] = "repr_tally" + tally = Tally(mean) + assert "repr_tally" in repr(tally)