From 4e8def97789ce3e4c4122b228266b15366c1d623 Mon Sep 17 00:00:00 2001 From: SteSeg Date: Tue, 6 Jan 2026 18:08:15 -0500 Subject: [PATCH 1/6] + tmc_manager draft --- .../uq/tmc_manager.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 src/openmc_fusion_benchmarks/uq/tmc_manager.py diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py new file mode 100644 index 00000000..2442ca62 --- /dev/null +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -0,0 +1,57 @@ +from typing import List, Callable +from pathlib import Path +import openmc +import numpy as np +import xarray as xr + + +class TMCManager: + def __init__(self, base_model: openmc.Model, perturbations: List[Callable], + n_samples, rng:np.random._generator.Generator=None): + self.base_model = base_model + self.perturbations = perturbations + self.n_samples = n_samples + self.results = [] + + # Example of setting rng for reproducibility + if rng is None: + self.rng = np.random.default_rng() + else: + self.rng = rng + + def run(self, cwd='.', *args, **kwargs): + + tmc_statepoint = Path(cwd) / "tmc_statepoint.h5" + + for i,p in enumerate(self.perturbations): + for n in range(self.n_samples): + perturbed_model = p(self.base_model, rng=self.rng) + # build the cwd/tmc/perturbation/ path + pert_cwd = Path(cwd) / "tmc" / f"perturbation_{i}" / f"sample_{n}" + sp_path = perturbed_model.run(cwd=pert_cwd, *args, **kwargs) + # store in the manager the sp_path info to build results path structure + + # Save results in tmc_statepoint + sp = openmc.StatePoint(sp_path) + # open tally and push to netcdf ofb style + for t in sp.tallies: + tally = sp.get_tally(id=t) + df = tally.get_pandas_dataframe() + + # df = df.drop(columns=['surface', 'cell', 'particle', 'nuclide', + # 'score', 'energyfunction'], errors='ignore') + + # Convert to xarray and add dimensions + da = xr.DataArray( + df.values[np.newaxis, :, :], # shape: (1, r, c) + dims=["realization", "row", "column"], + coords={ + "realization": [realization_label], + "column": df.columns, + "row": np.arange(df.shape[0]), + }, + name=tally.name + ) + + def _extract_results(self): + pass \ No newline at end of file From 5f3f8a10a92de5aa8ddc4be3fb69cc4fd6e22ea3 Mon Sep 17 00:00:00 2001 From: SteSeg Date: Thu, 8 Jan 2026 17:19:02 -0500 Subject: [PATCH 2/6] + process_tmc method --- .../uq/tmc_manager.py | 267 +++++++++++++++--- 1 file changed, 231 insertions(+), 36 deletions(-) diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index 2442ca62..b01aa05b 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -2,15 +2,16 @@ from pathlib import Path import openmc import numpy as np +import json import xarray as xr class TMCManager: def __init__(self, base_model: openmc.Model, perturbations: List[Callable], - n_samples, rng:np.random._generator.Generator=None): + realizations:int, rng:np.random._generator.Generator=None): self.base_model = base_model self.perturbations = perturbations - self.n_samples = n_samples + self.realizations = realizations self.results = [] # Example of setting rng for reproducibility @@ -21,37 +22,231 @@ def __init__(self, base_model: openmc.Model, perturbations: List[Callable], def run(self, cwd='.', *args, **kwargs): - tmc_statepoint = Path(cwd) / "tmc_statepoint.h5" - - for i,p in enumerate(self.perturbations): - for n in range(self.n_samples): - perturbed_model = p(self.base_model, rng=self.rng) - # build the cwd/tmc/perturbation/ path - pert_cwd = Path(cwd) / "tmc" / f"perturbation_{i}" / f"sample_{n}" - sp_path = perturbed_model.run(cwd=pert_cwd, *args, **kwargs) - # store in the manager the sp_path info to build results path structure - - # Save results in tmc_statepoint - sp = openmc.StatePoint(sp_path) - # open tally and push to netcdf ofb style - for t in sp.tallies: - tally = sp.get_tally(id=t) - df = tally.get_pandas_dataframe() - - # df = df.drop(columns=['surface', 'cell', 'particle', 'nuclide', - # 'score', 'energyfunction'], errors='ignore') - - # Convert to xarray and add dimensions - da = xr.DataArray( - df.values[np.newaxis, :, :], # shape: (1, r, c) - dims=["realization", "row", "column"], - coords={ - "realization": [realization_label], - "column": df.columns, - "row": np.arange(df.shape[0]), - }, - name=tally.name - ) - - def _extract_results(self): - pass \ No newline at end of file + cwd = Path(cwd).resolve() + + # Prepare TMC manifest file + manifest = cwd / "tmc_manifest.jsonl" + manifest.parent.mkdir(parents=True, exist_ok=True) + + # Open TMC manifest for folder structure + with manifest.open("a") as f_manifest: + # Run TMC engine + for p_idx, p in enumerate(self.perturbations): + for r_idx in range(self.realizations): + # Build perturbed model + perturbed_model = p(self.base_model, rng=self.rng) + + # Build run directory + run_dir = cwd / "tmc" / f"perturbation_{p_idx}" / f"realization_{r_idx}" + run_dir.mkdir(parents=True, exist_ok=True) + + # Run openmc in that directory + sp_path = perturbed_model.run(cwd=run_dir, *args, **kwargs) + sp_path = Path(sp_path).resolve() + + # Record in manifest (relative path from cwd) + rec = { + "perturbation": int(p_idx), + "realization": int(r_idx), + "statepoint": str(sp_path.relative_to(cwd)), + # "params": perturbation_params_if_any, + } + f_manifest.write(json.dumps(rec) + "\n") + + # Postprocess the whole TMC set + self._process_tmc(filepath=manifest) + + def _process_tmc(self, filepath="tmc_manifest.jsonl"): + + manifest_path = Path(filepath).resolve() + tmc_dir = manifest_path.parent + + # tmc_statepoint..h5 + tmc_statepoint = tmc_dir / f"tmc_statepoint.{int(self.realizations)}.h5" + + # Read TMC manifest file + records = [] + with manifest_path.open() as f: + for line in f: + if not line.strip(): + continue + rec = json.loads(line) + records.append(rec) + + # Sort records by perturbation and realization, not really necessary though + records.sort(key=lambda r: (r["perturbation"], r["realization"])) + + # Loop over all records and open the statepoints + for rec in records: + p_idx = rec["perturbation"] + r_idx = rec["realization"] + + # statepoint is stored relative to the project cwd when you wrote it + sp_path = Path(rec["statepoint"]).resolve() + + + # 1: Allocate the TMC array: add a sample axis + n_samples = 100 + pattern = "statepoint_{:04d}.h5" + tally_name = "my_tally" + + # --- Reference statepoint to determine shape and metadata --- + with openmc.StatePoint(pattern.format(1)) as sp: + t = sp.get_tally(name=tally_name) + filters = t.filters + filter_bins = [f.num_bins for f in filters] + n_nuclides = max(len(t.nuclides), 1) + n_scores = len(t.scores) + + flat_shape = t.mean.shape + nd_shape = tuple(filter_bins) + (n_nuclides, n_scores) + + # Sanity check + assert np.prod(filter_bins) == flat_shape[0] + assert flat_shape[1] == n_nuclides + assert flat_shape[2] == n_scores + + # Allocate big array: (sample, filter1, filter2, ..., nuclide, score) + tmc_data = np.empty((n_samples,) + nd_shape, dtype=float) + + # 2: Fill the TMC array + for i in range(n_samples): + fname = pattern.format(i+1) + with openmc.StatePoint(fname) as sp: + t = sp.get_tally(name=tally_name) + + # Flat (prod(filter_bins), n_nuclides, n_scores) + mean_flat = t.mean + + # Reshape to N‑D (filter1, filter2, ..., nuclide, score) + mean_nd = mean_flat.reshape(nd_shape) + + # Store into TMC array + tmc_data[i, ...] = mean_nd + + # 3: Keeping track of what each axis means + # The numeric array doesn’t store “cell vs energy vs mesh” labels by itself, so you should retain some metadata: + axis_info = { + "sample_axis": 0, + "filter_axes": [ + {"name": type(f).__name__, "num_bins": f.num_bins} + for f in filters + ], + "nuclide_axis": len(filter_bins) + 1, + "score_axis": len(filter_bins) + 2, + "nuclides": [str(n) for n in t.nuclides] if t.nuclides else ["total"], + "scores": list(t.scores), + } + # Optionally, for certain filters you can also store bin edges / mesh indices: + for f in filters: + if isinstance(f, openmc.EnergyFilter): + axis_info["energy_edges_eV"] = f.bins + # similarly for MeshFilter, CellFilter, etc. + + + def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): + manifest_path = Path(manifest_path).resolve() + tmc_dir = manifest_path.parent + + # Write the TMC statepoint file + tmc_statepoint = tmc_dir / f"tmc_statepoint.{int(self.realizations)}.h5" + + # Read TMC manifest + records = [] + with manifest_path.open() as f: + for line in f: + line = line.strip() + if not line: + continue + rec = json.loads(line) + records.append(rec) + + # Sort for deterministic order + records.sort(key=lambda r: (r["perturbation"], r["realization"])) + n_samples = len(records) + if n_samples == 0: + raise RuntimeError("TMC manifest is empty; no runs to process") + + # Use first statepoint as reference to determine shape & metadata + first_sp_path = Path(records[0]["statepoint"]).resolve() + tally_shapes = {} # key: tally id, value: nd_shape + tally_filters = {} # key: tally id, value: list of filters + tally_axisinfo = {} # key: tally id, value: axis_info dict + + with openmc.StatePoint(str(first_sp_path)) as sp0: + for tally in sp0.tallies.values(): + tid = tally.id # you can also use tally.name, but id is unambiguous + + filters = tally.filters + filter_bins = [f.num_bins for f in filters] + n_nuclides = max(len(tally.nuclides), 1) + n_scores = len(tally.scores) + + flat_shape = tally.mean.shape # (prod_bins, n_nuclides, n_scores) + nd_shape = tuple(filter_bins) + (n_nuclides, n_scores) + + assert flat_shape[0] == np.prod(filter_bins) + assert flat_shape[1] == n_nuclides + assert flat_shape[2] == n_scores + + tally_shapes[tid] = nd_shape + tally_filters[tid] = filters + + axis_info = { + "sample_axis": 0, + "filter_axes": [ + {"name": type(f).__name__, "num_bins": f.num_bins} + for f in filters + ], + "nuclide_axis": len(filter_bins) + 1, + "score_axis": len(filter_bins) + 2, + "nuclides": [str(n) for n in tally.nuclides] if tally.nuclides else ["total"], + "scores": list(tally.scores), + } + for f in filters: + if isinstance(f, openmc.EnergyFilter): + axis_info["energy_edges_eV"] = f.bins + + tally_axisinfo[tid] = axis_info + + # Allocate TMC arrays, one per tally + tmc_data = {} # key: tally id, value: ndarray (n_samples, ...) + + for tid, nd_shape in tally_shapes.items(): + tmc_data[tid] = np.empty((n_samples,) + nd_shape, dtype=float) + + # Fill arrays by looping over all statepoints + for i, rec in enumerate(records): + sp_path = Path(rec["statepoint"]).resolve() + with openmc.StatePoint(str(sp_path)) as sp: + for tid, arr in tmc_data.items(): + # assumes same tally IDs exist in every statepoint + tally = sp.tallies[tid] + mean_flat = tally.mean + nd_shape = tally_shapes[tid] + mean_nd = mean_flat.reshape(nd_shape) + arr[i, ...] = mean_nd + + # Build xarray Dataset and write to disk + ds = xr.Dataset() + sample_coord = np.arange(n_samples) + + for tid, arr in tmc_data.items(): + nd_shape = tally_shapes[tid] + dims = ("sample",) + tuple(f"dim_{k}" for k in range(len(nd_shape))) + + da = xr.DataArray( + arr, + dims=dims, + coords={"sample": sample_coord}, + name=f"tally_{tid}", + ) + + # attach axis info as attrs + for k, v in tally_axisinfo[tid].items(): + da.attrs[k] = v + + ds[da.name] = da + + ds.to_netcdf(tmc_statepoint) + From 0d0785f2712d63345b9a335ec44d51b207f82c0c Mon Sep 17 00:00:00 2001 From: SteSeg Date: Thu, 8 Jan 2026 17:20:19 -0500 Subject: [PATCH 3/6] - old process tmc method --- .../uq/tmc_manager.py | 91 +------------------ 1 file changed, 1 insertion(+), 90 deletions(-) diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index b01aa05b..c5eccfe7 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -56,94 +56,6 @@ def run(self, cwd='.', *args, **kwargs): # Postprocess the whole TMC set self._process_tmc(filepath=manifest) - def _process_tmc(self, filepath="tmc_manifest.jsonl"): - - manifest_path = Path(filepath).resolve() - tmc_dir = manifest_path.parent - - # tmc_statepoint..h5 - tmc_statepoint = tmc_dir / f"tmc_statepoint.{int(self.realizations)}.h5" - - # Read TMC manifest file - records = [] - with manifest_path.open() as f: - for line in f: - if not line.strip(): - continue - rec = json.loads(line) - records.append(rec) - - # Sort records by perturbation and realization, not really necessary though - records.sort(key=lambda r: (r["perturbation"], r["realization"])) - - # Loop over all records and open the statepoints - for rec in records: - p_idx = rec["perturbation"] - r_idx = rec["realization"] - - # statepoint is stored relative to the project cwd when you wrote it - sp_path = Path(rec["statepoint"]).resolve() - - - # 1: Allocate the TMC array: add a sample axis - n_samples = 100 - pattern = "statepoint_{:04d}.h5" - tally_name = "my_tally" - - # --- Reference statepoint to determine shape and metadata --- - with openmc.StatePoint(pattern.format(1)) as sp: - t = sp.get_tally(name=tally_name) - filters = t.filters - filter_bins = [f.num_bins for f in filters] - n_nuclides = max(len(t.nuclides), 1) - n_scores = len(t.scores) - - flat_shape = t.mean.shape - nd_shape = tuple(filter_bins) + (n_nuclides, n_scores) - - # Sanity check - assert np.prod(filter_bins) == flat_shape[0] - assert flat_shape[1] == n_nuclides - assert flat_shape[2] == n_scores - - # Allocate big array: (sample, filter1, filter2, ..., nuclide, score) - tmc_data = np.empty((n_samples,) + nd_shape, dtype=float) - - # 2: Fill the TMC array - for i in range(n_samples): - fname = pattern.format(i+1) - with openmc.StatePoint(fname) as sp: - t = sp.get_tally(name=tally_name) - - # Flat (prod(filter_bins), n_nuclides, n_scores) - mean_flat = t.mean - - # Reshape to N‑D (filter1, filter2, ..., nuclide, score) - mean_nd = mean_flat.reshape(nd_shape) - - # Store into TMC array - tmc_data[i, ...] = mean_nd - - # 3: Keeping track of what each axis means - # The numeric array doesn’t store “cell vs energy vs mesh” labels by itself, so you should retain some metadata: - axis_info = { - "sample_axis": 0, - "filter_axes": [ - {"name": type(f).__name__, "num_bins": f.num_bins} - for f in filters - ], - "nuclide_axis": len(filter_bins) + 1, - "score_axis": len(filter_bins) + 2, - "nuclides": [str(n) for n in t.nuclides] if t.nuclides else ["total"], - "scores": list(t.scores), - } - # Optionally, for certain filters you can also store bin edges / mesh indices: - for f in filters: - if isinstance(f, openmc.EnergyFilter): - axis_info["energy_edges_eV"] = f.bins - # similarly for MeshFilter, CellFilter, etc. - - def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): manifest_path = Path(manifest_path).resolve() tmc_dir = manifest_path.parent @@ -175,8 +87,7 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): with openmc.StatePoint(str(first_sp_path)) as sp0: for tally in sp0.tallies.values(): - tid = tally.id # you can also use tally.name, but id is unambiguous - + tid = tally.id filters = tally.filters filter_bins = [f.num_bins for f in filters] n_nuclides = max(len(tally.nuclides), 1) From 243567d82b11bb9d2a78942d06209b6cd1a26cfe Mon Sep 17 00:00:00 2001 From: SteSeg Date: Mon, 12 Jan 2026 11:16:34 -0500 Subject: [PATCH 4/6] + draft tmc _process, StatePoint & TMCTally --- src/openmc_fusion_benchmarks/uq/__init__.py | 1 + .../uq/tmc_manager.py | 304 +++++++++++++++--- 2 files changed, 268 insertions(+), 37 deletions(-) diff --git a/src/openmc_fusion_benchmarks/uq/__init__.py b/src/openmc_fusion_benchmarks/uq/__init__.py index fcd7ced3..17660be2 100644 --- a/src/openmc_fusion_benchmarks/uq/__init__.py +++ b/src/openmc_fusion_benchmarks/uq/__init__.py @@ -1,2 +1,3 @@ from .uq_utils import * from .tmc_engine import * +from .tmc_manager import * diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index c5eccfe7..d45be6fc 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -3,22 +3,23 @@ import openmc import numpy as np import json +import copy import xarray as xr +import inspect class TMCManager: def __init__(self, base_model: openmc.Model, perturbations: List[Callable], realizations:int, rng:np.random._generator.Generator=None): self.base_model = base_model - self.perturbations = perturbations + # self.perturbations = perturbations self.realizations = realizations - self.results = [] + self.rng = rng or np.random.default_rng() + + # perturbations is a list of factories: factory(rng) -> perturb(model) + # call each factory once to get a closure perturb(model) -> model + self.perturbations = [factory(self.rng) for factory in perturbations] - # Example of setting rng for reproducibility - if rng is None: - self.rng = np.random.default_rng() - else: - self.rng = rng def run(self, cwd='.', *args, **kwargs): @@ -33,14 +34,14 @@ def run(self, cwd='.', *args, **kwargs): # Run TMC engine for p_idx, p in enumerate(self.perturbations): for r_idx in range(self.realizations): - # Build perturbed model - perturbed_model = p(self.base_model, rng=self.rng) - # Build run directory + # fresh copy so perturbations do not accumulate + model_copy = copy.deepcopy(self.base_model) + perturbed_model = p(model_copy) + run_dir = cwd / "tmc" / f"perturbation_{p_idx}" / f"realization_{r_idx}" run_dir.mkdir(parents=True, exist_ok=True) - # Run openmc in that directory sp_path = perturbed_model.run(cwd=run_dir, *args, **kwargs) sp_path = Path(sp_path).resolve() @@ -54,16 +55,16 @@ def run(self, cwd='.', *args, **kwargs): f_manifest.write(json.dumps(rec) + "\n") # Postprocess the whole TMC set - self._process_tmc(filepath=manifest) + self._process_tmc(manifest_path=manifest) def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): manifest_path = Path(manifest_path).resolve() tmc_dir = manifest_path.parent - # Write the TMC statepoint file + # xarray will write NetCDF (HDF5-backed); extension is up to you tmc_statepoint = tmc_dir / f"tmc_statepoint.{int(self.realizations)}.h5" - # Read TMC manifest + # ---- 1. Read TMC manifest ---- records = [] with manifest_path.open() as f: for line in f: @@ -75,19 +76,21 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): # Sort for deterministic order records.sort(key=lambda r: (r["perturbation"], r["realization"])) - n_samples = len(records) - if n_samples == 0: + n_realizations = len(records) + if n_realizations == 0: raise RuntimeError("TMC manifest is empty; no runs to process") - # Use first statepoint as reference to determine shape & metadata + # ---- 2. Use first statepoint as reference to determine shape & metadata ---- first_sp_path = Path(records[0]["statepoint"]).resolve() - tally_shapes = {} # key: tally id, value: nd_shape - tally_filters = {} # key: tally id, value: list of filters + tally_names = {} # key: tally id, value: tally name + tally_shapes = {} # key: tally id, value: nd_shape + tally_filters = {} # key: tally id, value: list of filters tally_axisinfo = {} # key: tally id, value: axis_info dict with openmc.StatePoint(str(first_sp_path)) as sp0: for tally in sp0.tallies.values(): - tid = tally.id + tid = tally.id + filters = tally.filters filter_bins = [f.num_bins for f in filters] n_nuclides = max(len(tally.nuclides), 1) @@ -100,6 +103,7 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): assert flat_shape[1] == n_nuclides assert flat_shape[2] == n_scores + tally_names[tid] = tally.name tally_shapes[tid] = nd_shape tally_filters[tid] = filters @@ -110,23 +114,20 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): for f in filters ], "nuclide_axis": len(filter_bins) + 1, - "score_axis": len(filter_bins) + 2, - "nuclides": [str(n) for n in tally.nuclides] if tally.nuclides else ["total"], - "scores": list(tally.scores), + "score_axis": len(filter_bins) + 2, + "nuclides": [str(n) for n in tally.nuclides] if tally.nuclides else ["total"], + "scores": list(tally.scores), } - for f in filters: - if isinstance(f, openmc.EnergyFilter): - axis_info["energy_edges_eV"] = f.bins tally_axisinfo[tid] = axis_info - # Allocate TMC arrays, one per tally + # ---- 3. Allocate TMC arrays, one per tally ---- tmc_data = {} # key: tally id, value: ndarray (n_samples, ...) for tid, nd_shape in tally_shapes.items(): - tmc_data[tid] = np.empty((n_samples,) + nd_shape, dtype=float) + tmc_data[tid] = np.empty((n_realizations,) + nd_shape, dtype=float) - # Fill arrays by looping over all statepoints + # ---- 4. Fill arrays by looping over all statepoints ---- for i, rec in enumerate(records): sp_path = Path(rec["statepoint"]).resolve() with openmc.StatePoint(str(sp_path)) as sp: @@ -138,26 +139,255 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): mean_nd = mean_flat.reshape(nd_shape) arr[i, ...] = mean_nd - # Build xarray Dataset and write to disk + # ---- 5. Build xarray Dataset and write to disk ---- ds = xr.Dataset() - sample_coord = np.arange(n_samples) + sample_coord = np.arange(n_realizations) for tid, arr in tmc_data.items(): nd_shape = tally_shapes[tid] - dims = ("sample",) + tuple(f"dim_{k}" for k in range(len(nd_shape))) + filters = tally_filters[tid] + + # Build dimension names following OpenMC conventions + filter_dims = [] + for f in filters: + # Use filter type name without "Filter" suffix, lowercase + filter_type = type(f).__name__.replace("Filter", "").lower() + filter_dims.append(filter_type) + + dims = ("realization",) + tuple(filter_dims) + ("nuclide", "score") + # Use tally name if available, otherwise fall back to ID + tally_name = tally_names[tid] or f"tally_{tid}" + da = xr.DataArray( arr, dims=dims, - coords={"sample": sample_coord}, - name=f"tally_{tid}", + coords={"realization": sample_coord}, + name=tally_name, ) - # attach axis info as attrs - for k, v in tally_axisinfo[tid].items(): - da.attrs[k] = v + # Attach tally metadata + da.attrs["tally_id"] = tid + da.attrs["tally_name"] = tally_names[tid] + + # Attach axis info as attrs; serialize non-scalar things to JSON strings + axisinfo = tally_axisinfo[tid] + for k, v in axisinfo.items(): + # scalars are fine + if isinstance(v, (int, float, bool, str, np.number)): + da.attrs[k] = v + else: + # lists, dicts, numpy arrays -> JSON string + if isinstance(v, np.ndarray): + to_dump = v.tolist() + else: + to_dump = v + da.attrs[k] = json.dumps(to_dump) ds[da.name] = da + # If you have netcdf4/h5netcdf installed, you can also specify engine explicitly: + # ds.to_netcdf(tmc_statepoint, engine="h5netcdf") ds.to_netcdf(tmc_statepoint) + + # Store path for later retrieval + self.tmc_statepoint_path = tmc_statepoint + def get_tmc_statepoint(self, path=None): + """ + Load and return a TMCStatePoint wrapper for the TMC results. + + Parameters + ---------- + path : str or Path, optional + Path to the TMC statepoint file. If not provided, uses the path + from the last run() call. + + Returns + ------- + TMCStatePoint + Wrapper object providing OpenMC StatePoint-like interface to TMC data. + """ + if path is None: + if not hasattr(self, 'tmc_statepoint_path'): + raise RuntimeError("No TMC statepoint path available. Either run TMC first or provide path.") + path = self.tmc_statepoint_path + else: + path = Path(path).resolve() + + return TMCStatePoint(path) + + +class TMCStatePoint: + """ + Wrapper for TMC statepoint providing 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() + self._ds = xr.open_dataset(self.path) + 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 = {} + for var_name in self._ds.data_vars: + da = self._ds[var_name] + tally_id = da.attrs.get('tally_id') + if tally_id is not None: + self._tallies[tally_id] = TMCTally(da) + return self._tallies + + def get_tally(self, id=None, name=None): + """ + Get a tally by ID or name (mimics openmc.StatePoint.get_tally). + + Parameters + ---------- + id : int, optional + Tally ID + name : str, optional + Tally name + + Returns + ------- + TMCTally + The requested tally + """ + if id is not None: + return self.tallies[id] + 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 'id' or 'name'") + + def close(self): + """Close the underlying NetCDF file.""" + self._ds.close() + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + def __repr__(self): + n_tallies = len(self.tallies) + n_realizations = self._ds.dims.get('realization', 0) + return f"" + + +class TMCTally: + """ + Wrapper for a single TMC tally providing OpenMC Tally-like interface. + + Parameters + ---------- + data_array : xarray.DataArray + The DataArray containing the TMC tally data + """ + + def __init__(self, data_array): + self._da = data_array + + @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 mean(self): + """Mean values across all realizations.""" + return self._da.mean(dim='realization').values + + @property + def std_dev(self): + """Standard deviation across realizations.""" + return self._da.std(dim='realization').values + + @property + def data(self): + """Full TMC data array (all realizations).""" + return self._da.values + + @property + def scores(self): + """List of score names.""" + scores_json = self._da.attrs.get('scores') + if scores_json: + return json.loads(scores_json) + return [] + + @property + def nuclides(self): + """List of nuclide names.""" + nuclides_json = self._da.attrs.get('nuclides') + if nuclides_json: + return json.loads(nuclides_json) + return [] + + @property + def shape(self): + """Shape of the data array.""" + return self._da.shape + + @property + def dims(self): + """Dimension names.""" + return self._da.dims + + 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 data + """ + da = self._da + + # Apply filter dimension selections + if filter_kwargs: + da = da.sel(**filter_kwargs) + + # Apply score selection + if scores is not None: + 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: + 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 From c4d065210b9c059316841b271f4153c543d89e4a Mon Sep 17 00:00:00 2001 From: SteSeg Date: Mon, 12 Jan 2026 12:05:18 -0500 Subject: [PATCH 5/6] + improvements in TMCTally --- .../uq/tmc_manager.py | 90 +++++++++++++++++-- 1 file changed, 82 insertions(+), 8 deletions(-) diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index d45be6fc..f6e30ac3 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -123,9 +123,11 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): # ---- 3. Allocate TMC arrays, one per tally ---- tmc_data = {} # key: tally id, value: ndarray (n_samples, ...) + tmc_mc_std = {} # key: tally id, value: MC std_dev per realization for tid, nd_shape in tally_shapes.items(): tmc_data[tid] = np.empty((n_realizations,) + nd_shape, dtype=float) + tmc_mc_std[tid] = np.empty((n_realizations,) + nd_shape, dtype=float) # ---- 4. Fill arrays by looping over all statepoints ---- for i, rec in enumerate(records): @@ -135,9 +137,12 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): # assumes same tally IDs exist in every statepoint tally = sp.tallies[tid] mean_flat = tally.mean + std_flat = tally.std_dev nd_shape = tally_shapes[tid] mean_nd = mean_flat.reshape(nd_shape) + std_nd = std_flat.reshape(nd_shape) arr[i, ...] = mean_nd + tmc_mc_std[tid][i, ...] = std_nd # ---- 5. Build xarray Dataset and write to disk ---- ds = xr.Dataset() @@ -148,13 +153,14 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): filters = tally_filters[tid] # Build dimension names following OpenMC conventions + # Make nuclide and score dimensions unique per tally to avoid conflicts filter_dims = [] for f in filters: # Use filter type name without "Filter" suffix, lowercase filter_type = type(f).__name__.replace("Filter", "").lower() filter_dims.append(filter_type) - dims = ("realization",) + tuple(filter_dims) + ("nuclide", "score") + dims = ("realization",) + tuple(filter_dims) + (f"nuclide_{tid}", f"score_{tid}") # Use tally name if available, otherwise fall back to ID tally_name = tally_names[tid] or f"tally_{tid}" @@ -165,10 +171,20 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): coords={"realization": sample_coord}, name=tally_name, ) + + # Create corresponding MC std_dev DataArray + da_mc_std = xr.DataArray( + tmc_mc_std[tid], + dims=dims, + coords={"realization": sample_coord}, + name=f"{tally_name}_mc_std", + ) # Attach tally metadata da.attrs["tally_id"] = tid da.attrs["tally_name"] = tally_names[tid] + da_mc_std.attrs["tally_id"] = tid + da_mc_std.attrs["tally_name"] = tally_names[tid] # Attach axis info as attrs; serialize non-scalar things to JSON strings axisinfo = tally_axisinfo[tid] @@ -183,8 +199,20 @@ def _process_tmc(self, manifest_path="tmc_manifest.jsonl"): else: to_dump = v da.attrs[k] = json.dumps(to_dump) + + # Copy axis info to MC std DataArray + for k, v in axisinfo.items(): + if isinstance(v, (int, float, bool, str, np.number)): + da_mc_std.attrs[k] = v + else: + if isinstance(v, np.ndarray): + to_dump = v.tolist() + else: + to_dump = v + da_mc_std.attrs[k] = json.dumps(to_dump) ds[da.name] = da + ds[da_mc_std.name] = da_mc_std # If you have netcdf4/h5netcdf installed, you can also specify engine explicitly: # ds.to_netcdf(tmc_statepoint, engine="h5netcdf") @@ -239,13 +267,16 @@ def tallies(self): if self._tallies is None: self._tallies = {} for var_name in self._ds.data_vars: + # Skip MC std arrays (they're accessed via the main tally) + if var_name.endswith('_mc_std'): + continue da = self._ds[var_name] tally_id = da.attrs.get('tally_id') if tally_id is not None: - self._tallies[tally_id] = TMCTally(da) + self._tallies[tally_id] = TMCTally(da, parent_ds=self._ds) return self._tallies - def get_tally(self, id=None, name=None): + def get_tally(self, tally_id=None, name=None): """ Get a tally by ID or name (mimics openmc.StatePoint.get_tally). @@ -261,8 +292,8 @@ def get_tally(self, id=None, name=None): TMCTally The requested tally """ - if id is not None: - return self.tallies[id] + if tally_id is not None: + return self.tallies[tally_id] elif name is not None: for tally in self.tallies.values(): if tally.name == name: @@ -295,10 +326,13 @@ class TMCTally: ---------- data_array : xarray.DataArray The DataArray containing the TMC tally data + parent_ds : xarray.Dataset, optional + Parent dataset containing MC uncertainty data """ - def __init__(self, data_array): + def __init__(self, data_array, parent_ds=None): self._da = data_array + self._parent_ds = parent_ds @property def id(self): @@ -312,14 +346,41 @@ def name(self): @property def mean(self): - """Mean values across all realizations.""" + """TMC mean: mean value across all realizations.""" return self._da.mean(dim='realization').values @property def std_dev(self): - """Standard deviation across realizations.""" + """TMC standard deviation: propagated parametric uncertainty across realizations.""" return self._da.std(dim='realization').values + @property + def realization_means(self): + """Mean value for each individual realization (shape: n_realizations x ...).""" + return self._da.values + + @property + def realization_mc_stds(self): + """ + Monte Carlo standard deviation for each individual realization. + + This is the statistical uncertainty from particle sampling within each + individual OpenMC run (shape: n_realizations x ...). + """ + # Get the corresponding MC std DataArray from the parent dataset + mc_std_name = f"{self._da.name}_mc_std" + if hasattr(self._da, '_parent_ds') and mc_std_name in self._da._parent_ds: + return self._da._parent_ds[mc_std_name].values + # Fallback: try to find it in the same file + try: + ds = xr.open_dataset(self._da.encoding.get('source', '')) + if mc_std_name in ds: + return ds[mc_std_name].values + except: + pass + # If not found, return zeros as fallback + return np.zeros_like(self._da.values) + @property def data(self): """Full TMC data array (all realizations).""" @@ -341,6 +402,19 @@ def nuclides(self): return json.loads(nuclides_json) return [] + @property + def filters(self): + """List of filter information (type and number of bins).""" + filters_json = self._da.attrs.get('filter_axes') + if filters_json: + return json.loads(filters_json) + return [] + + @property + def realizations(self): + """Number of TMC realizations.""" + return self._da.sizes.get('realization', 0) + @property def shape(self): """Shape of the data array.""" From ef5f1bce4f8509da5a615e962909deb55d0f283b Mon Sep 17 00:00:00 2001 From: SteSeg Date: Mon, 12 Jan 2026 13:06:43 -0500 Subject: [PATCH 6/6] small fix in TMCTally --- src/openmc_fusion_benchmarks/uq/tmc_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index f6e30ac3..379bd92e 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -360,7 +360,7 @@ def realization_means(self): return self._da.values @property - def realization_mc_stds(self): + def realization_stds(self): """ Monte Carlo standard deviation for each individual realization.