Fix sparse storage of multidimensional tally results (#4030)

openmc/tallies.py

@@ -479,10 +479,12 @@ class Tally(IDManagerMixin):
 
             # Convert NumPy arrays to SciPy sparse LIL matrices
             if self.sparse:
-                self._sum = lil_array(self._sum.flatten(), self._sum.shape)
-                self._sum_sq = lil_array(self._sum_sq.flatten(), self._sum_sq.shape)
-                self._sum_third = lil_array(self._sum_third.flatten(), self._sum_third.shape)
-                self._sum_fourth = lil_array(self._sum_fourth.flatten(), self._sum_fourth.shape)
+                self._sum = lil_array(self._sum.reshape(1, -1))
+                self._sum_sq = lil_array(self._sum_sq.reshape(1, -1))
+                if self._sum_third is not None:
+                    self._sum_third = lil_array(self._sum_third.reshape(1, -1))
+                if self._sum_fourth is not None:
+                    self._sum_fourth = lil_array(self._sum_fourth.reshape(1, -1))
 
             # Read simulation time (needed for figure of merit)
             self._simulation_time = f["runtime"]["simulation"][()]
@@ -578,7 +580,7 @@ class Tally(IDManagerMixin):
 
             # Convert NumPy array to SciPy sparse LIL matrix
             if self.sparse:
-                self._mean = lil_array(self._mean.flatten(), self._mean.shape)
+                self._mean = lil_array(self._mean.reshape(1, -1))
 
         if self.sparse:
             return np.reshape(self._mean.toarray(), self.shape)
@@ -599,7 +601,7 @@ class Tally(IDManagerMixin):
 
             # Convert NumPy array to SciPy sparse LIL matrix
             if self.sparse:
-                self._std_dev = lil_array(self._std_dev.flatten(), self._std_dev.shape)
+                self._std_dev = lil_array(self._std_dev.reshape(1, -1))
 
             self.with_batch_statistics = True
 
@@ -630,7 +632,7 @@ class Tally(IDManagerMixin):
             self._vov[mask] = numerator[mask]/denominator[mask] - 1.0/n
 
             if self.sparse:
-                self._vov = lil_array(self._vov.flatten(), self._vov.shape)
+                self._vov = lil_array(self._vov.reshape(1, -1))
 
         if self.sparse:
             return np.reshape(self._vov.toarray(), self.shape)
@@ -1005,17 +1007,19 @@ class Tally(IDManagerMixin):
         # Convert NumPy arrays to SciPy sparse LIL matrices
         if sparse and not self.sparse:
             if self._sum is not None:
-                self._sum = lil_array(self._sum.flatten(), self._sum.shape)
+                self._sum = lil_array(self._sum.reshape(1, -1))
             if self._sum_sq is not None:
-                self._sum_sq = lil_array(self._sum_sq.flatten(), self._sum_sq.shape)
+                self._sum_sq = lil_array(self._sum_sq.reshape(1, -1))
             if self._sum_third is not None:
-                self._sum_third = lil_array(self._sum_third.flatten(), self._sum_third.shape)
+                self._sum_third = lil_array(self._sum_third.reshape(1, -1))
             if self._sum_fourth is not None:
-                self._sum_fourth = lil_array(self._sum_fourth.flatten(), self._sum_fourth.shape)
+                self._sum_fourth = lil_array(self._sum_fourth.reshape(1, -1))
             if self._mean is not None:
-                self._mean = lil_array(self._mean.flatten(), self._mean.shape)
+                self._mean = lil_array(self._mean.reshape(1, -1))
             if self._std_dev is not None:
-                self._std_dev = lil_array(self._std_dev.flatten(), self._std_dev.shape)
+                self._std_dev = lil_array(self._std_dev.reshape(1, -1))
+            if self._vov is not None:
+                self._vov = lil_array(self._vov.reshape(1, -1))
 
             self._sparse = True
 
@@ -1033,6 +1037,8 @@ class Tally(IDManagerMixin):
                 self._mean = np.reshape(self._mean.toarray(), self.shape)
             if self._std_dev is not None:
                 self._std_dev = np.reshape(self._std_dev.toarray(), self.shape)
+            if self._vov is not None:
+                self._vov = np.reshape(self._vov.toarray(), self.shape)
             self._sparse = False
 
     def remove_score(self, score):

tests/unit_tests/test_tallies.py

@@ -184,6 +184,40 @@ def _tally_from_data(x, *, higher_moments=True, normality=True):
         t._sum_fourth = np.array([[[np.sum(x**4)]]], dtype=float)
     return t
 
+
+def test_sparse_tally_data_roundtrip():
+    tally = _tally_from_data(np.arange(1.0, 21.0))
+    expected_sum = tally._sum.copy()
+    expected_sum_sq = tally._sum_sq.copy()
+    expected_sum_third = tally._sum_third.copy()
+    expected_sum_fourth = tally._sum_fourth.copy()
+    expected_mean = expected_sum / tally.num_realizations
+    expected_std_dev = np.sqrt(
+        (expected_sum_sq / tally.num_realizations - expected_mean**2)
+        / (tally.num_realizations - 1)
+    )
+
+    tally.sparse = True
+
+    assert tally.sum == pytest.approx(expected_sum)
+    assert tally.sum_sq == pytest.approx(expected_sum_sq)
+    assert tally.sum_third == pytest.approx(expected_sum_third)
+    assert tally.sum_fourth == pytest.approx(expected_sum_fourth)
+    assert tally.mean == pytest.approx(expected_mean)
+    assert tally.std_dev == pytest.approx(expected_std_dev)
+    expected_vov = tally.vov.copy()
+
+    tally.sparse = False
+
+    assert tally.sum == pytest.approx(expected_sum)
+    assert tally.sum_sq == pytest.approx(expected_sum_sq)
+    assert tally.sum_third == pytest.approx(expected_sum_third)
+    assert tally.sum_fourth == pytest.approx(expected_sum_fourth)
+    assert tally.mean == pytest.approx(expected_mean)
+    assert tally.std_dev == pytest.approx(expected_std_dev)
+    assert tally.vov == pytest.approx(expected_vov)
+
+
 @pytest.mark.parametrize(
     "x, skew_true, kurt_true",
     [   # Rademacher distribution