Skip to content

Harmoniza tmctallies #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SteSeg merged 4 commits into from
May 11, 2026
Merged

Harmoniza tmctallies #7

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2178566
TMCTally inherits from BaseTally
SteSeg May 8, 2026
bedf1ca
unify dimension naming
SteSeg May 8, 2026
f8aa0ba
harmonization of tmc types
SteSeg May 8, 2026
3f394e2
removed code duplicates
SteSeg May 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 29 additions & 69 deletions src/openmc_fusion_benchmarks/uq/tmc_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down Expand Up @@ -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":
Expand All @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -571,7 +533,7 @@ def __repr__(self):
return f"<TMCStatePoint: {n_realizations} TMC combinations, {n_tallies} tallies>"


class TMCTally:
class TMCTally(BaseTally):
"""
Wrapper for a single TMC tally providing an OpenMC Tally-like interface.

Expand All @@ -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 = [
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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")
Expand Down
Loading