diff --git a/src/openmc_fusion_benchmarks/uq/tmc_manager.py b/src/openmc_fusion_benchmarks/uq/tmc_manager.py index 6414c668..18eef8a0 100644 --- a/src/openmc_fusion_benchmarks/uq/tmc_manager.py +++ b/src/openmc_fusion_benchmarks/uq/tmc_manager.py @@ -9,6 +9,9 @@ import itertools import h5py +from ..tallies import BaseTally +from ..backends.openmc.tallies import openmc_tally_to_dataset + class TMCManager: def __init__(self, base_model: openmc.Model, perturbations: List[Callable], @@ -216,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": @@ -238,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) @@ -255,42 +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": [ - {"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), + 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) @@ -367,34 +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] - - # 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) + template = tally_templates[tid] + template_dims = tally_dims[tid] - # within each tally group, we can use generic "nuclide" and "score" - dims = extra_dims + tuple(filter_dims) + ("nuclide", "score") + dims = extra_dims + template_dims - # 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( @@ -417,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 @@ -571,7 +533,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 +548,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 = [ @@ -733,7 +693,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") @@ -773,7 +733,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")