Merge pull request #2235 from paulromano/source-rejection

Source rejection by domain (cell, material, universe)

include/openmc/source.h

@@ -4,6 +4,8 @@
 #ifndef OPENMC_SOURCE_H
 #define OPENMC_SOURCE_H
 
+#include <unordered_set>
+
 #include "pugixml.hpp"
 
 #include "openmc/distribution_multi.h"
@@ -51,13 +53,15 @@ public:
 };
 
 //==============================================================================
-//! Source composed of independent spatial, angle, energy, and time distributions
+//! Source composed of independent spatial, angle, energy, and time
+//! distributions
 //==============================================================================
 
 class IndependentSource : public Source {
 public:
   // Constructors
-  IndependentSource(UPtrSpace space, UPtrAngle angle, UPtrDist energy, UPtrDist time);
+  IndependentSource(
+    UPtrSpace space, UPtrAngle angle, UPtrDist energy, UPtrDist time);
   explicit IndependentSource(pugi::xml_node node);
 
   //! Sample from the external source distribution
@@ -76,12 +80,18 @@ public:
   Distribution* time() const { return time_.get(); }
 
 private:
+  // Domain types
+  enum class DomainType { UNIVERSE, MATERIAL, CELL };
+
+  // Data members
   ParticleType particle_ {ParticleType::neutron}; //!< Type of particle emitted
   double strength_ {1.0};                         //!< Source strength
   UPtrSpace space_;                               //!< Spatial distribution
   UPtrAngle angle_;                               //!< Angular distribution
   UPtrDist energy_;                               //!< Energy distribution
   UPtrDist time_;                                 //!< Time distribution
+  DomainType domain_type_;                        //!< Domain type for rejection
+  std::unordered_set<int32_t> domain_ids_;        //!< Domains to reject from
 };
 
 //==============================================================================

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 @@ class Source:
         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 @@ class Source:
     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 @@ class Source:
             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 @@ class Source:
             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:

src/source.cpp

@@ -16,8 +16,10 @@
 #include "openmc/bank.h"
 #include "openmc/capi.h"
 #include "openmc/cell.h"
+#include "openmc/container_util.h"
 #include "openmc/error.h"
 #include "openmc/file_utils.h"
+#include "openmc/geometry.h"
 #include "openmc/hdf5_interface.h"
 #include "openmc/material.h"
 #include "openmc/memory.h"
@@ -151,29 +153,48 @@ IndependentSource::IndependentSource(pugi::xml_node node)
       double p[] {1.0};
       time_ = UPtrDist {new Discrete {T, p, 1}};
     }
+
+    // Check for domains to reject from
+    if (check_for_node(node, "domain_type")) {
+      std::string domain_type = get_node_value(node, "domain_type");
+      if (domain_type == "cell") {
+        domain_type_ = DomainType::CELL;
+      } else if (domain_type == "material") {
+        domain_type_ = DomainType::MATERIAL;
+      } else if (domain_type == "universe") {
+        domain_type_ = DomainType::UNIVERSE;
+      } else {
+        fatal_error(std::string(
+          "Unrecognized domain type for source rejection: " + domain_type));
+      }
+
+      auto ids = get_node_array<int>(node, "domain_ids");
+      domain_ids_.insert(ids.begin(), ids.end());
+    }
   }
 }
 
 SourceSite IndependentSource::sample(uint64_t* seed) const
 {
   SourceSite site;
+  site.particle = particle_;
 
   // Repeat sampling source location until a good site has been found
   bool found = false;
   int n_reject = 0;
   static int n_accept = 0;
+
   while (!found) {
     // Set particle type
-    site.particle = particle_;
+    Particle p;
+    p.type() = particle_;
+    p.u() = {0.0, 0.0, 1.0};
 
     // Sample spatial distribution
-    site.r = space_->sample(seed);
+    p.r() = space_->sample(seed);
 
     // Now search to see if location exists in geometry
-    int32_t cell_index, instance;
-    double xyz[] {site.r.x, site.r.y, site.r.z};
-    int err = openmc_find_cell(xyz, &cell_index, &instance);
-    found = (err != OPENMC_E_GEOMETRY);
+    found = exhaustive_find_cell(p);
 
     // Check if spatial site is in fissionable material
     if (found) {
@@ -181,15 +202,29 @@ SourceSite IndependentSource::sample(uint64_t* seed) const
       if (space_box) {
         if (space_box->only_fissionable()) {
           // Determine material
-          const auto& c = model::cells[cell_index];
-          auto mat_index =
-            c->material_.size() == 1 ? c->material_[0] : c->material_[instance];
-
+          auto mat_index = p.material();
           if (mat_index == MATERIAL_VOID) {
             found = false;
           } else {
-            if (!model::materials[mat_index]->fissionable_)
-              found = false;
+            found = model::materials[mat_index]->fissionable_;
+          }
+        }
+      }
+
+      // Rejection based on cells/materials/universes
+      if (!domain_ids_.empty()) {
+        found = false;
+        if (domain_type_ == DomainType::MATERIAL) {
+          auto mat_index = p.material();
+          if (mat_index != MATERIAL_VOID) {
+            found = contains(domain_ids_, model::materials[mat_index]->id());
+          }
+        } else {
+          for (const auto& coord : p.coord()) {
+            auto id = (domain_type_ == DomainType::CELL)
+                        ? model::cells[coord.cell]->id_
+                        : model::universes[coord.universe]->id_;
+            if (found = contains(domain_ids_, id)) break;
           }
         }
       }
@@ -205,6 +240,8 @@ SourceSite IndependentSource::sample(uint64_t* seed) const
                     "definition.");
       }
     }
+
+    site.r = p.r();
   }
 
   // Sample angle

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()