Merge pull request #2288 from eepeterson/filter_shape

Change to get_reshaped_data for consistency with mesh dimensions

openmc/arithmetic.py

@@ -604,6 +604,10 @@ class AggregateFilter:
     def num_bins(self):
         return len(self.bins) if self.aggregate_filter else 0
 
+    @property
+    def shape(self):
+        return (self.num_bins,)
+
     @type.setter
     def type(self, filter_type):
         if filter_type not in _FILTER_TYPES:

openmc/filter.py

@@ -102,6 +102,8 @@ class Filter(IDManagerMixin, metaclass=FilterMeta):
         Unique identifier for the filter
     num_bins : Integral
         The number of filter bins
+    shape : tuple
+        The shape of the filter
 
     """
 
@@ -205,6 +207,10 @@ class Filter(IDManagerMixin, metaclass=FilterMeta):
     def num_bins(self):
         return len(self.bins)
 
+    @property
+    def shape(self):
+        return (self.num_bins,)
+
     def check_bins(self, bins):
         """Make sure given bins are valid for this filter.
 
@@ -839,6 +845,12 @@ class MeshFilter(Filter):
         else:
             self.bins = list(mesh.indices)
 
+    @property
+    def shape(self):
+        if isinstance(self, MeshSurfaceFilter):
+            return (self.num_bins,)
+        return self.mesh.dimension
+
     @property
     def translation(self):
         return self._translation
@@ -958,8 +970,6 @@ class MeshSurfaceFilter(MeshFilter):
 
     Attributes
     ----------
-    bins : Integral
-        The mesh ID
     mesh : openmc.MeshBase
         The mesh object that events will be tallied onto
     translation : Iterable of float
@@ -968,10 +978,8 @@ class MeshSurfaceFilter(MeshFilter):
     id : int
         Unique identifier for the filter
     bins : list of tuple
-
         A list of mesh indices / surfaces for each filter bin, e.g. [(1, 1,
         'x-min out'), (1, 1, 'x-min in'), ...]
-
     num_bins : Integral
         The number of filter bins
 

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,12 +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
-        new_shape = tuple(f.num_bins for f in self.filters)
+        # 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):
+            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 += (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)