Speed up apply_time_correction by reducing file I/O and deepcopies (#3617)

openmc/deplete/d1s.py

@@ -5,7 +5,7 @@ shutdown dose rate calculations.
 
 """
 
-from copy import deepcopy
+from copy import copy
 from typing import Sequence
 from math import log, prod
 
@@ -164,8 +164,12 @@ def apply_time_correction(
     radionuclides = [str(x) for x in tally.filters[i_filter].bins]
     tcf = np.array([time_correction_factors[x][index] for x in radionuclides])
 
-    # Create copy of tally
-    new_tally = deepcopy(tally)
+    # Force tally results to be read and std_dev to be computed
+    tally.std_dev
+
+    # Create shallow copy of tally
+    new_tally = copy(tally)
+    new_tally._filters = copy(tally._filters)
 
     # Determine number of bins in other filters
     n_bins_before = prod([f.num_bins for f in tally.filters[:i_filter]])
@@ -177,32 +181,33 @@ def apply_time_correction(
     shape = (n_bins_before, n_radionuclides, n_bins_after, n_nuclides, n_scores)
     tally_sum = new_tally.sum.reshape(shape)
     tally_sum_sq = new_tally.sum_sq.reshape(shape)
+    tally_mean = new_tally.mean.reshape(shape)
+    tally_std_dev = new_tally.std_dev.reshape(shape)
 
     # Apply TCF, broadcasting to the correct dimensions
     tcf.shape = (1, -1, 1, 1, 1)
     new_tally._sum = tally_sum * tcf
     new_tally._sum_sq = tally_sum_sq * (tcf*tcf)
-    new_tally._mean = None
-    new_tally._std_dev = None
+    new_tally._mean = tally_mean * tcf
+    new_tally._std_dev = tally_std_dev * tcf
 
     shape = (-1, n_nuclides, n_scores)
 
     if sum_nuclides:
-        # Query the mean and standard deviation
-        mean = new_tally.mean
-        std_dev = new_tally.std_dev
-
         # Sum over parent nuclides (note that when combining different bins for
         # parent nuclide, we can't work directly on sum_sq)
-        new_tally._mean = mean.sum(axis=1).reshape(shape)
-        new_tally._std_dev = np.linalg.norm(std_dev, axis=1).reshape(shape)
+        new_tally._mean = new_tally.mean.sum(axis=1).reshape(shape)
+        new_tally._std_dev = np.linalg.norm(new_tally.std_dev, axis=1).reshape(shape)
         new_tally._derived = True
 
         # Remove ParentNuclideFilter
         new_tally.filters.pop(i_filter)
     else:
+        # Change shape back to (filter combinations, nuclides, scores)
         new_tally._sum.shape = shape
         new_tally._sum_sq.shape = shape
+        new_tally._mean.shape = shape
+        new_tally._std_dev.shape = shape
 
     return new_tally
 

tests/unit_tests/test_d1s.py

@@ -120,6 +120,13 @@ def test_apply_time_correction(run_in_tmpdir):
         tally = sp.tallies[tally.id]
         flux = tally.mean.flatten()
 
+    # Copy attributes from original tally
+    tally_filters = list(tally.filters)
+    tally_sum = tally.sum.copy()
+    tally_sum_sq = tally.sum_sq.copy()
+    tally_mean = tally.mean.copy()
+    tally_std_dev = tally.std_dev.copy()
+
     # Apply TCF and make sure results are consistent
     result = d1s.apply_time_correction(tally, factors, sum_nuclides=False)
     tcf = np.array([factors[nuc][-1] for nuc in nuclides])
@@ -129,6 +136,13 @@ def test_apply_time_correction(run_in_tmpdir):
     result_summed = d1s.apply_time_correction(tally, factors)
     assert result_summed.mean.flatten()[0] == pytest.approx(result.mean.sum())
 
+    # Make sure original tally is unchanged
+    assert tally.filters == tally_filters
+    assert np.all(tally.sum == tally_sum)
+    assert np.all(tally.sum_sq == tally_sum_sq)
+    assert np.all(tally.mean == tally_mean)
+    assert np.all(tally.std_dev == tally_std_dev)
+
     # Make sure various tally methods work
     result.get_values()
     result_summed.get_values()