added test and expand_dims kwarg

openmc/tallies.py

@@ -1405,7 +1405,7 @@ class Tally(IDManagerMixin):
 
         return df
 
-    def get_reshaped_data(self, value='mean'):
+    def get_reshaped_data(self, value='mean', expand_dims=False):
         """Returns an array of tally data with one dimension per filter.
 
         The tally data in OpenMC is stored as a 3D array with the dimensions
@@ -1417,17 +1417,24 @@ class Tally(IDManagerMixin):
 
         This builds and returns a reshaped version of the tally data array with
         unique dimensions corresponding to each tally filter. For example,
-        suppose this tally has arrays of data with shape (8,5,5) corresponding
-        to two filters (2 and 4 bins, respectively), five nuclides and five
+        suppose this tally has arrays of data with shape (30,5,5) corresponding
+        to two filters (2 and 15 bins, respectively), five nuclides and five
         scores. This method will return a version of the data array with the
-        with a new shape of (2,4,5,5) such that the first two dimensions
-        correspond directly to the two filters with two and four bins.
+        with a new shape of (2,15,5,5) such that the first two dimensions
+        correspond directly to the two filters with two and fifteen bins. If
+        expand_dims is True and our filter above with 15 bins is an instance of
+        :class:`openmc.MeshFilter` with a shape of (3,5,1). The resulting tally
+        data array will have a new shape of (2,3,5,1,5,5).
 
         Parameters
         ----------
         value : str
             A string for the type of value to return  - 'mean' (default),
             'std_dev', 'rel_err', 'sum', or 'sum_sq' are accepted
+        expand_dims : bool, optional
+            Whether or not to expand the dimensions of filters with multiple
+            dimensions. This will result in more than one dimension per filter
+            for the returned data array.
 
         Returns
         -------
@@ -1439,24 +1446,32 @@ class Tally(IDManagerMixin):
         # Get the 3D array of data in filters, nuclides and scores
         data = self.get_values(value=value)
 
-        # Build a new array shape with one dimension per filter
+        # Build a new array shape with one dimension per filter or expand
+        # multidimensional filters if desired
         new_shape = tuple()
         idx0 = None
         for i, f in enumerate(self.filters):
-            # Mesh filter indices are backwards so we need to flip them
-            if isinstance(f, openmc.MeshFilter):
-                fshape = f.shape[::-1]
-                new_shape += fshape
-                idx0, idx1 = i, i + len(fshape) - 1
+            if expand_dims:
+                # Mesh filter indices are backwards so we need to flip them
+                if isinstance(f, openmc.MeshFilter):
+                    fshape = f.shape[::-1]
+                    new_shape += fshape
+                    idx0, idx1 = i, i + len(fshape) - 1
+                else:
+                    new_shape += f.shape
             else:
-                new_shape += f.shape
+                new_shape += (np.prod(f.shape),)
 
         new_shape += (self.num_nuclides, self.num_scores)
 
         # Reshape the data with one dimension for each filter
         data = np.reshape(data, new_shape)
+
+        # If we had a MeshFilter we should swap the axes to have the same shape
+        # for the data and the filter
         if idx0 is not None:
             data = np.swapaxes(data, idx0, idx1)
+
         return data
 
     def hybrid_product(self, other, binary_op, filter_product=None,

tests/unit_tests/test_filter_mesh.py

@@ -1,7 +1,10 @@
+import math
+
 import numpy as np
-import openmc
 from uncertainties import unumpy
 
+import openmc
+
 
 def test_spherical_mesh_estimators(run_in_tmpdir):
     """Test that collision/tracklength estimators agree for SphericalMesh"""
@@ -62,7 +65,8 @@ def test_cylindrical_mesh_estimators(run_in_tmpdir):
     mat.add_nuclide('U235', 1.0)
     mat.set_density('g/cm3', 10.0)
 
-    cyl = openmc.model.RightCircularCylinder((0., 0., -5.), 10., 10.0, boundary_type='vacuum')
+    cyl = openmc.model.RightCircularCylinder((0., 0., -5.), 10., 10.0,
+                                             boundary_type='vacuum')
     cell = openmc.Cell(fill=mat, region=-cyl)
     model = openmc.Model()
     model.geometry = openmc.Geometry([cell])
@@ -107,3 +111,43 @@ def test_cylindrical_mesh_estimators(run_in_tmpdir):
     diff = unumpy.nominal_values(delta)
     std_dev = unumpy.std_devs(delta)
     assert np.all(diff < 3*std_dev)
+
+def test_get_reshaped_data(run_in_tmpdir):
+    """Test that expanding MeshFilter dimensions works as expected"""
+
+    mat = openmc.Material()
+    mat.add_nuclide('U235', 1.0)
+    mat.set_density('g/cm3', 10.0)
+
+    sphere = openmc.Sphere(r=10.0, boundary_type='vacuum')
+    cell = openmc.Cell(fill=mat, region=-sphere)
+    model = openmc.Model()
+    model.geometry = openmc.Geometry([cell])
+    model.settings.particles = 1_000
+    model.settings.inactive = 10
+    model.settings.batches = 20
+
+    sph_mesh = openmc.SphericalMesh()
+    sph_mesh.r_grid = np.linspace(0.0, 5.0**3, 20)**(1/3)
+    sph_mesh.theta_grid = np.linspace(0, math.pi, 4)
+    sph_mesh.phi_grid = np.linspace(0, 2*math.pi, 3)
+    tally1 = openmc.Tally()
+    efilter = openmc.EnergyFilter([0, 1e5, 1e8])
+    meshfilter = openmc.MeshFilter(sph_mesh)
+    assert meshfilter.shape == (19, 3, 2)
+    tally1.filters = [efilter, meshfilter]
+    tally1.scores = ['flux']
+
+    model.tallies = openmc.Tallies([tally1])
+
+    # Run OpenMC
+    sp_filename = model.run()
+
+    # Get flux tally as reshaped data
+    with openmc.StatePoint(sp_filename) as sp:
+        t1 = sp.tallies[tally1.id]
+        data1 = t1.get_reshaped_data()
+        data2 = t1.get_reshaped_data(expand_dims=True)
+
+    assert data1.shape == (2, 19*3*2, 1, 1)
+    assert data2.shape == (2, 19, 3, 2, 1, 1)