Fix access order issues after applying tally results from Model.run. (#3313)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

openmc/statepoint.py

@@ -6,6 +6,7 @@ import warnings
 
 import h5py
 import numpy as np
+from pathlib import Path
 from uncertainties import ufloat
 
 import openmc
@@ -415,7 +416,7 @@ class StatePoint:
 
                     # Create Tally object and assign basic properties
                     tally = openmc.Tally(tally_id)
-                    tally._sp_filename = self._f.filename
+                    tally._sp_filename = Path(self._f.filename)
                     tally.name = group['name'][()].decode() if 'name' in group else ''
 
                     # Check if tally has multiply_density attribute

openmc/tallies.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from collections.abc import Iterable, MutableSequence
 import copy
-from functools import partial, reduce
+from functools import partial, reduce, wraps
 from itertools import product
 from numbers import Integral, Real
 import operator
@@ -179,6 +179,25 @@ class Tally(IDManagerMixin):
         parts.append('{: <15}=\t{}'.format('Multiply dens.', self.multiply_density))
         return '\n\t'.join(parts)
 
+    @staticmethod
+    def ensure_results(f):
+        """A decorator to be applied to any method that might use tally results.
+           Results will be loaded if appropriate based on the tally properties.
+
+        Args:
+            f function: Tally method to wrap
+
+        Returns:
+            function: Wrapped function that reads tally results before calling
+            the methodif necessary
+        """
+        @wraps(f)
+        def read(self):
+            if self._sp_filename is not None and not self.derived:
+                self._read_results()
+            return f(self)
+        return read
+
     @property
     def name(self):
         return self._name
@@ -218,6 +237,7 @@ class Tally(IDManagerMixin):
         self._filters = cv.CheckedList(_FILTER_CLASSES, 'tally filters', filters)
 
     @property
+    @ensure_results
     def nuclides(self):
         return self._nuclides
 
@@ -314,6 +334,7 @@ class Tally(IDManagerMixin):
                                         triggers)
 
     @property
+    @ensure_results
     def num_realizations(self):
         return self._num_realizations
 
@@ -340,11 +361,11 @@ class Tally(IDManagerMixin):
         with h5py.File(self._sp_filename, 'r') as f:
             # Set number of realizations
             group = f[f'tallies/tally {self.id}']
-            self.num_realizations = int(group['n_realizations'][()])
+            self._num_realizations = int(group['n_realizations'][()])
 
             # Update nuclides
             nuclide_names = group['nuclides'][()]
-            self.nuclides = [name.decode().strip() for name in nuclide_names]
+            self._nuclides = [name.decode().strip() for name in nuclide_names]
 
             # Extract Tally data from the file
             data = group['results']
@@ -368,13 +389,11 @@ class Tally(IDManagerMixin):
         self._results_read = True
 
     @property
+    @ensure_results
     def sum(self):
         if not self._sp_filename or self.derived:
             return None
 
-        # Make sure results have been read
-        self._read_results()
-
         if self.sparse:
             return np.reshape(self._sum.toarray(), self.shape)
         else:
@@ -386,13 +405,11 @@ class Tally(IDManagerMixin):
         self._sum = sum
 
     @property
+    @ensure_results
     def sum_sq(self):
         if not self._sp_filename or self.derived:
             return None
 
-        # Make sure results have been read
-        self._read_results()
-
         if self.sparse:
             return np.reshape(self._sum_sq.toarray(), self.shape)
         else:
@@ -935,10 +952,24 @@ class Tally(IDManagerMixin):
         statepoint : openmc.PathLike or openmc.StatePoint
             Statepoint used to update tally results
         """
+        # derived tallies are populated with data based on combined tallies
+        # and should not be modified
+        if self.derived:
+            return
+
         if isinstance(statepoint, openmc.StatePoint):
-            self._sp_filename = statepoint._f.filename
+            self._sp_filename = Path(statepoint._f.filename)
         else:
-            self._sp_filename = str(statepoint)
+            self._sp_filename = Path(str(statepoint))
+
+        # reset these properties to ensure that any results access after this
+        # point are based on the current statepoint file
+        self._sum = None
+        self._sum_sq = None
+        self._mean = None
+        self._std_dev = None
+        self._num_realizations = 0
+        self._results_read = False
 
     @classmethod
     def from_xml_element(cls, elem, **kwargs):

tests/unit_tests/test_tallies.py

@@ -106,14 +106,44 @@ def test_tally_application(sphere_model, run_in_tmpdir):
     tally.scores = ['flux', 'absorption', 'fission', 'scatter']
     sphere_model.tallies = [tally]
 
+    # FIRST RUN
     # run the simulation and apply results
     sp_file = sphere_model.run(apply_tally_results=True)
+    # before calling for any property requiring results (including the equivalence check below),
+    # the following internal attributes of the original should be unset
+    assert tally._mean is None
+    assert tally._std_dev is None
+    assert tally._sum is None
+    assert tally._sum_sq is None
+    assert tally._num_realizations == 0
+    # the statepoint file property should be set, however
+    assert tally._sp_filename == sp_file
+
     with openmc.StatePoint(sp_file) as sp:
         assert tally in sp.tallies.values()
         sp_tally = sp.tallies[tally.id]
 
     # at this point the tally information regarding results should be the same
-    assert (sp_tally.mean == tally.mean).all()
     assert (sp_tally.std_dev == tally.std_dev).all()
+    assert (sp_tally.mean == tally.mean).all()
     assert sp_tally.nuclides == tally.nuclides
 
+    # SECOND RUN
+    # change the number of particles and ensure that the results are different
+    sphere_model.settings.particles += 1
+    sp_file = sphere_model.run(apply_tally_results=True)
+
+    assert (sp_tally.std_dev != tally.std_dev).any()
+    assert (sp_tally.mean != tally.mean).any()
+
+    # now re-read data from the new stateopint file and
+    # ensure that the new results match those in
+    # the latest statepoint
+    with openmc.StatePoint(sp_file) as sp:
+        assert tally in sp.tallies.values()
+        sp_tally = sp.tallies[tally.id]
+
+    # at this point the tally information regarding results should be the same
+    assert (sp_tally.std_dev == tally.std_dev).all()
+    assert (sp_tally.mean == tally.mean).all()
+    assert sp_tally.nuclides == tally.nuclides