Repository: openmc-dev/openmc

Issue:
Generic filter for rejecting sources points when sampling
As originally discussed on Discourse  https://openmc.discourse.group/t/source-uniform-sampling-in-cell-of-given-material/1062

It would be great to add generic filter to allow the rejection of source points when being sampled during the creation of particles.

Source sampling should allow for rejection of the proposed sampled coordinates if the coordinates are not in a user specified material or not in a user specified cell.

While the generic source sampling filter would be separate to the distributions themselves it could be passed to the distributions and the distributions should make use of the filter to ascertain if the sampling point results in a neutron or not.

Perhaps the best distribution to demonstrate this on is the ```openmc.stats.Box``` distribution as it already has a similar feature.

This will allow us to make interesting shaped sources like the example picture below which shows source sampling randomly across a openmc.stats.Box rejecting sample points that don't appear within a certain material / cell

![Screenshot from 2021-04-27 20-16-20](https://user-images.githubusercontent.com/8583900/116299325-73d16300-a795-11eb-8566-1cee34602f41.png)

There are some related lines in OpenMC relating to a similar feature that allows sampling points to be rejected if they are not within a fissionable material with the ```only_fissionable``` keyword that might be useful for anyone considering this feature.

 https://github.com/openmc-dev/openmc/blob/efcc0c158edcc3e389eb20a5c5ffc4e0a65ef554/src/distribution_spatial.cpp#L187-L204

https://github.com/openmc-dev/openmc/blob/efcc0c158edcc3e389eb20a5c5ffc4e0a65ef554/src/source.cpp#L171-L173

https://github.com/openmc-dev/openmc/blob/efcc0c158edcc3e389eb20a5c5ffc4e0a65ef554/openmc/stats/multivariate.py#L617-L716


Discussion:
I've issued a bounty reward for this issue https://www.bountysource.com/issues/98138076-generic-filter-for-rejecting-sources-points-when-sampling
Could the spawn probability be inhomogeneous and a function of x, y, z?

Resolving patch:
diff --git a/openmc/source.py b/openmc/source.py
--- a/openmc/source.py
+++ b/openmc/source.py
@@ -1,10 +1,13 @@
+from collections.abc import Iterable
 from enum import Enum
 from numbers import Real
+import warnings
 from xml.etree import ElementTree as ET
 
 import numpy as np
 import h5py
 
+import openmc
 import openmc.checkvalue as cv
 from openmc.stats.multivariate import UnitSphere, Spatial
 from openmc.stats.univariate import Univariate
@@ -36,6 +39,9 @@ class Source:
         Strength of the source
     particle : {'neutron', 'photon'}
         Source particle type
+    domains : iterable of openmc.Cell, openmc.Material, or openmc.Universe
+        Domains to reject based on, i.e., if a sampled spatial location is not
+        within one of these domains, it will be rejected.
 
     Attributes
     ----------
@@ -57,11 +63,16 @@ class Source:
         Strength of the source
     particle : {'neutron', 'photon'}
         Source particle type
+    ids : Iterable of int
+        IDs of domains to use for rejection
+    domain_type : {'cell', 'material', 'universe'}
+        Type of domain to use for rejection
 
     """
 
     def __init__(self, space=None, angle=None, energy=None, time=None, filename=None,
-                 library=None, parameters=None, strength=1.0, particle='neutron'):
+                 library=None, parameters=None, strength=1.0, particle='neutron',
+                 domains=None):
         self._space = None
         self._angle = None
         self._energy = None
@@ -87,6 +98,17 @@ def __init__(self, space=None, angle=None, energy=None, time=None, filename=None
         self.strength = strength
         self.particle = particle
 
+        self._domain_ids = []
+        self._domain_type = None
+        if domains is not None:
+            if isinstance(domains[0], openmc.Cell):
+                self.domain_type = 'cell'
+            elif isinstance(domains[0], openmc.Material):
+                self.domain_type = 'material'
+            elif isinstance(domains[0], openmc.Universe):
+                self.domain_type = 'universe'
+            self.domain_ids = [d.id for d in domains]
+
     @property
     def file(self):
         return self._file
@@ -123,6 +145,24 @@ def strength(self):
     def particle(self):
         return self._particle
 
+    @property
+    def domain_ids(self):
+        return self._domain_ids
+
+    @property
+    def domain_type(self):
+        return self._domain_type
+
+    @domain_ids.setter
+    def domain_ids(self, ids):
+        cv.check_type('domain IDs', ids, Iterable, Real)
+        self._domain_ids = ids
+
+    @domain_type.setter
+    def domain_type(self, domain_type):
+        cv.check_value('domain type', domain_type, ('cell', 'material', 'universe'))
+        self._domain_type = domain_type
+
     @file.setter
     def file(self, filename):
         cv.check_type('source file', filename, str)
@@ -196,6 +236,11 @@ def to_xml_element(self):
             element.append(self.energy.to_xml_element('energy'))
         if self.time is not None:
             element.append(self.time.to_xml_element('time'))
+        if self.domain_ids:
+            dt_elem = ET.SubElement(element, "domain_type")
+            dt_elem.text = self.domain_type
+            id_elem = ET.SubElement(element, "domain_ids")
+            id_elem.text = ' '.join(str(uid) for uid in self.domain_ids)
         return element
 
     @classmethod
@@ -213,7 +258,24 @@ def from_xml_element(cls, elem):
             Source generated from XML element
 
         """
-        source = cls()
+        domain_type = get_text(elem, "domain_type")
+        if domain_type is not None:
+            domain_ids = [int(x) for x in get_text(elem, "domain_ids").split()]
+
+            # Instantiate some throw-away domains that are used by the
+            # constructor to assign IDs
+            with warnings.catch_warnings():
+                warnings.simplefilter('ignore', openmc.IDWarning)
+                if domain_type == 'cell':
+                    domains = [openmc.Cell(uid) for uid in domain_ids]
+                elif domain_type == 'material':
+                    domains = [openmc.Material(uid) for uid in domain_ids]
+                elif domain_type == 'universe':
+                    domains = [openmc.Universe(uid) for uid in domain_ids]
+        else:
+            domains = None
+
+        source = cls(domains=domains)
 
         strength = get_text(elem, 'strength')
         if strength is not None:


Test patch:
diff --git a/tests/unit_tests/test_source.py b/tests/unit_tests/test_source.py
--- a/tests/unit_tests/test_source.py
+++ b/tests/unit_tests/test_source.py
@@ -1,6 +1,7 @@
 from math import pi
 
 import openmc
+import openmc.lib
 import openmc.stats
 import numpy as np
 from pytest import approx
@@ -96,3 +97,40 @@ def test_source_xml_roundtrip():
     np.testing.assert_allclose(new_src.angle.reference_uvw, src.angle.reference_uvw)
     assert new_src.particle == src.particle
     assert new_src.strength == approx(src.strength)
+
+
+def test_rejection(run_in_tmpdir):
+    # Model with two spheres inside a box
+    mat = openmc.Material()
+    mat.add_nuclide('H1', 1.0)
+    sph1 = openmc.Sphere(x0=3, r=1.0)
+    sph2 = openmc.Sphere(x0=-3, r=1.0)
+    cube = openmc.model.RectangularParallelepiped(
+        -5., 5., -5., 5., -5., 5., boundary_type='reflective'
+    )
+    cell1 = openmc.Cell(fill=mat, region=-sph1)
+    cell2 = openmc.Cell(fill=mat, region=-sph2)
+    non_source_region = +sph1 & +sph2 & -cube
+    cell3 = openmc.Cell(region=non_source_region)
+    model = openmc.Model()
+    model.geometry = openmc.Geometry([cell1, cell2, cell3])
+    model.settings.particles = 100
+    model.settings.batches = 10
+    model.settings.run_mode = 'fixed source'
+
+    # Set up a box source with rejection on the spherical cell
+    space = openmc.stats.Box(*cell3.bounding_box)
+    model.settings.source = openmc.Source(space=space, domains=[cell1, cell2])
+
+    # Load up model via openmc.lib and sample source
+    model.export_to_xml()
+    openmc.lib.init()
+    particles = openmc.lib.sample_external_source(1000)
+
+    # Make sure that all sampled sources are within one of the spheres
+    joint_region = cell1.region | cell2.region
+    for p in particles:
+        assert p.r in joint_region
+        assert p.r not in non_source_region
+
+    openmc.lib.finalize()