Fix for issue #3158 - Streamline use of CompositeSurface with SurfaceFilter

openmc/filter.py

@@ -428,22 +428,33 @@ def get_pandas_dataframe(self, data_size, stride, **kwargs):
 
 class WithIDFilter(Filter):
     """Abstract parent for filters of types with IDs (Cell, Material, etc.)."""
-    def __init__(self, bins, filter_id=None):
-        bins = np.atleast_1d(bins)
 
+    def clean_bins(self, bins):
+        """Clean up bins if needed (e.g. expand CompositeSurface's into their component surfaces)
+        """
+        return np.atleast_1d(bins)
+
+    def check_bins(self, bins):
+        # Ensure bins are unique, warn if not
+        if len(bins) != len(set(bins)):
+            msg = f'Duplicate bins found in {self.short_name} filter: {bins}'
+        # Check the bin values
+        for edge in bins:
+            cv.check_greater_than('filter bin', edge, 0, equality=True)
+
+    @Filter.bins.setter
+    def bins(self, bins):
+        # Expand objects and clean up bins if needed
+        bins = self.clean_bins(bins)
         # Make sure bins are either integers or appropriate objects
         cv.check_iterable_type('filter bins', bins,
                                (Integral, self.expected_type))
-
         # Extract ID values
         bins = np.array([b if isinstance(b, Integral) else b.id
                          for b in bins])
-        super().__init__(bins, filter_id)
-
-    def check_bins(self, bins):
-        # Check the bin values.
-        for edge in bins:
-            cv.check_greater_than('filter bin', edge, 0, equality=True)
+        # Validate bins
+        self.check_bins(bins)
+        self._bins = bins
 
 
 class UniverseFilter(WithIDFilter):
@@ -730,6 +741,21 @@ class SurfaceFilter(WithIDFilter):
     """
     expected_type = Surface
 
+    def clean_bins(self, bins):
+        """"If composite surfaces are present, expand add the component surfaces to the bins.
+        """
+        bins = np.atleast_1d(bins)
+        composite_surfaces = list(filter(lambda x: isinstance(x, openmc.model.CompositeSurface), bins))
+        if composite_surfaces:
+            msg = f'In SurfaceFilter {len(composite_surfaces)} bins will be added for the ' \
+                  f'CompositeSurface {composite_surfaces}.'
+            warnings.warn(msg, UserWarning)
+            bins = list(filter(lambda x: isinstance(x, openmc.Surface), bins))
+            composite_surface_bins = [s for cs in composite_surfaces for s in cs.component_surfaces if s not in bins]
+            bins = np.concatenate((np.atleast_1d(bins), composite_surface_bins))
+
+        return bins
+
 
 class ParticleFilter(Filter):
     """Bins tally events based on the Particle type.

openmc/model/surface_composite.py

@@ -31,6 +31,14 @@ def rotate(self, rotation, pivot=(0., 0., 0.), order='xyz', inplace=False):
             s = getattr(surf, name)
             setattr(surf, name, s.rotate(rotation, pivot, order, inplace))
         return surf
+
+    @property
+    def component_surfaces(self):
+        return [getattr(self, name) for name in self._surface_names]
+
+    @property
+    def component_surface_ids(self):
+        return [s.id for s in self.component_surfaces]
 
     @property
     def boundary_type(self):

tests/unit_tests/test_filters.py

@@ -294,3 +294,51 @@ def test_tabular_from_energyfilter():
 
     tab = efilter.get_tabular(values=np.array([10, 10, 5]), interpolation='linear-linear')
     assert tab.interpolation == 'linear-linear'
+
+
+def test_surfacefilter_w_compositesurface(run_in_tmpdir, box_model):
+    m = openmc.Material()
+    m.add_nuclide('U235', 1.0)
+    m.set_density('g/cm3', 1.0)
+
+    box = openmc.model.RectangularPrism(10., 10., boundary_type='vacuum')
+    c = openmc.Cell(fill=m, region=-box)
+    box_model.geometry.root_universe = openmc.Universe(cells=[c])
+
+    tally = openmc.Tally()
+    tally.filters = [openmc.SurfaceFilter([box])]
+    tally.scores = ['current']
+    assert tally.filters[0].num_bins == 4
+
+    box_model.tallies = [tally]
+
+    sp_name = box_model.run()
+
+    with openmc.StatePoint(sp_name) as sp:
+        current = sp.tallies[tally.id]
+        assert len(current.get_pandas_dataframe()['surface']) == 4
+        assert np.all(current.get_pandas_dataframe()['surface'] == box.component_surface_ids)
+
+    box = openmc.model.RectangularParallelepiped(*[-10, 10]*3, boundary_type='vacuum')
+    c = openmc.Cell(fill=m, region=-box)
+    box_model.geometry.root_universe = openmc.Universe(cells=[c])
+
+
+    tally = openmc.Tally()
+    surface_filter = openmc.SurfaceFilter([box])
+    tally.filters = [surface_filter]
+    tally.scores = ['current']
+    assert tally.filters[0].num_bins == 6
+
+    # the redundant surface added here should not change the number of bins
+    surface_filter.bins = [box, box.xmin]
+    assert surface_filter.num_bins == 6
+
+    box_model.tallies = [tally]
+
+    sp_name = box_model.run()
+
+    with openmc.StatePoint(sp_name) as sp:
+        current = sp.tallies[tally.id]
+        assert len(current.get_pandas_dataframe()['surface']) == 6
+        assert np.all(current.get_pandas_dataframe()['surface'] == box.component_surface_ids)