From 2178566c51792e5b3c2d83f22329aa6ba2c67918 Mon Sep 17 00:00:00 2001 From: SteSeg Date: Fri, 8 May 2026 12:05:15 -0400 Subject: [PATCH 1/4] TMCTally inherits from BaseTally --- src/openmc_fusion_benchmarks/uq/tmc_manager.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index 6414c668..0af55dbf 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -9,6 +9,8 @@ import itertools import h5py +from ..tallies import BaseTally + class TMCManager: def __init__(self, base_model: openmc.Model, perturbations: List[Callable], @@ -571,7 +573,7 @@ def __repr__(self): return f"" -class TMCTally: +class TMCTally(BaseTally): """ Wrapper for a single TMC tally providing an OpenMC Tally-like interface. @@ -586,9 +588,7 @@ class TMCTally: """ 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 + super().__init__(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 = [ From bedf1caa8eb0d625bee092f16e3ada1a457d9181 Mon Sep 17 00:00:00 2001 From: SteSeg Date: Fri, 8 May 2026 12:09:26 -0400 Subject: [PATCH 2/4] unify dimension naming --- src/openmc_fusion_benchmarks/uq/tmc_manager.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index 0af55dbf..e8103a3a 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -10,6 +10,7 @@ import h5py from ..tallies import BaseTally +from ..backends.openmc.tallies import _unique_filter_dims, _build_filter_axis_metadata class TMCManager: @@ -283,10 +284,7 @@ def resolve_statepoint_path(sp_str: str) -> Path: tally_filters[tid] = filters axis_info = { - "filter_axes": [ - {"name": type(f).__name__, "num_bins": f.num_bins} - for f in filters - ], + "filter_axes": _build_filter_axis_metadata(filters, _unique_filter_dims(filters)), "nuclides": [str(n) for n in tally.nuclides] if tally.nuclides else ["total"], "scores": list(tally.scores), } @@ -373,11 +371,8 @@ def resolve_statepoint_path(sp_str: str) -> Path: 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 consistent with backend serializer + filter_dims = _unique_filter_dims(filters) # within each tally group, we can use generic "nuclide" and "score" dims = extra_dims + tuple(filter_dims) + ("nuclide", "score") From f8aa0bacf52c2fb611b66fc7b53cf17e2ab9b51b Mon Sep 17 00:00:00 2001 From: SteSeg Date: Fri, 8 May 2026 12:22:32 -0400 Subject: [PATCH 3/4] harmonization of tmc types --- src/openmc_fusion_benchmarks/uq/tmc_manager.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index e8103a3a..55f4bf40 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -219,7 +219,7 @@ def resolve_statepoint_path(sp_str: str) -> Path: extra_dims = ("perturbation", "realization") extra_coords = { "perturbation": np.arange(n_perturbations), - "realization": np.arange(n_realizations) + "realization": np.arange(n_realizations), } elif mode == "diagonal": @@ -241,7 +241,6 @@ def resolve_statepoint_path(sp_str: str) -> Path: # Infer per-dimension sizes from the data: indices_array = np.array([rec["indices"] for rec in records], dtype=int) - # assume indices run from 0..(n_i-1) along each axis per_dim_sizes = indices_array.max(axis=0) + 1 # length per perturbation dim extra_shape = tuple(int(n) for n in per_dim_sizes) @@ -728,7 +727,7 @@ def per_perturbation_mean(self): 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") @@ -768,7 +767,7 @@ def per_perturbation_std_dev(self): 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") From 3f394e2a6ce5ae37cb658bbbb9bc0f82dc9def1b Mon Sep 17 00:00:00 2001 From: SteSeg Date: Fri, 8 May 2026 12:32:06 -0400 Subject: [PATCH 4/4] removed code duplicates --- .../uq/tmc_manager.py | 78 ++++++------------- 1 file changed, 22 insertions(+), 56 deletions(-) diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index 55f4bf40..18eef8a0 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -10,7 +10,7 @@ import h5py from ..tallies import BaseTally -from ..backends.openmc.tallies import _unique_filter_dims, _build_filter_axis_metadata +from ..backends.openmc.tallies import openmc_tally_to_dataset class TMCManager: @@ -257,39 +257,28 @@ def resolve_statepoint_path(sp_str: str) -> Path: # ---- 3. Use first statepoint as reference for tallies ---- first_sp_path = resolve_statepoint_path(first_rec["statepoint"]) - tally_names = {} # tid -> name - tally_shapes = {} # tid -> nd_shape (filters..., nuclide, score) - tally_filters = {} # tid -> list of filters - tally_axisinfo = {} # tid -> axis_info dict + tally_names = {} # tid -> name + tally_shapes = {} # tid -> nd_shape (filters..., nuclide, score) + tally_templates = {} # tid -> xarray.Dataset template + tally_dims = {} # tid -> tuple of dims (filters..., nuclide, score) + tally_coords = {} # tid -> dim -> coord values with openmc.StatePoint(str(first_sp_path)) as sp0: for tally in sp0.tallies.values(): tid = tally.id - 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 + ds_template = openmc_tally_to_dataset(tally) + da_template = ds_template["mean"] tally_names[tid] = tally.name - tally_shapes[tid] = nd_shape - tally_filters[tid] = filters - - axis_info = { - "filter_axes": _build_filter_axis_metadata(filters, _unique_filter_dims(filters)), - "nuclides": [str(n) for n in tally.nuclides] if tally.nuclides else ["total"], - "scores": list(tally.scores), + tally_shapes[tid] = tuple(int(s) for s in da_template.shape) + tally_templates[tid] = ds_template + tally_dims[tid] = tuple(da_template.dims) + tally_coords[tid] = { + dim: np.asarray(da_template.coords[dim].values) + for dim in da_template.dims } - tally_axisinfo[tid] = axis_info - # ---- 4. Allocate arrays: one per tally ---- tmc_data = {} # tid -> ndarray (extra_shape + nd_shape) tmc_mc_std = {} # tid -> ndarray (extra_shape + nd_shape) @@ -366,31 +355,15 @@ def resolve_statepoint_path(sp_str: str) -> Path: # ---- 6. Build per-tally Datasets and write each into its own group ---- for tid, arr in tmc_data.items(): - nd_shape = tally_shapes[tid] - filters = tally_filters[tid] - axisinfo = tally_axisinfo[tid] + template = tally_templates[tid] + template_dims = tally_dims[tid] - # filter dims consistent with backend serializer - filter_dims = _unique_filter_dims(filters) + dims = extra_dims + template_dims - # within each tally group, we can use generic "nuclide" and "score" - dims = extra_dims + tuple(filter_dims) + ("nuclide", "score") - - # coords: TMC dims coords = dict(extra_coords) + for dim in template_dims: + coords[dim] = (dim, tally_coords[tid][dim]) - # coords: filter dimensions (add integer indices for each filter) - for i, (f, fdim) in enumerate(zip(filters, filter_dims)): - coords[fdim] = (fdim, np.arange(f.num_bins)) - - # coords: nuclide / score for this tally - nuclides = axisinfo["nuclides"] - scores = axisinfo["scores"] - - coords["nuclide"] = ("nuclide", np.array(nuclides, dtype="U")) - coords["score"] = ("score", np.array(scores, dtype="U")) - - # Per-tally dataset ds_tid = xr.Dataset() da_mean = xr.DataArray( @@ -413,16 +386,9 @@ def resolve_statepoint_path(sp_str: str) -> Path: target_da.attrs["tally_id"] = tid target_da.attrs["tally_name"] = tally_name - # serialize complex axisinfo at dataset level - for k, v in axisinfo.items(): - if isinstance(v, (int, float, bool, str, np.number)): - ds_tid.attrs[k] = v - else: - if isinstance(v, np.ndarray): - to_dump = v.tolist() - else: - to_dump = v - ds_tid.attrs[k] = json.dumps(to_dump) + # copy dataset-level metadata from template + for key, value in template.attrs.items(): + ds_tid.attrs[key] = value ds_tid["mean"] = da_mean ds_tid["mc_std"] = da_mc_std