Simple rejection based on a vector of cell IDs

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
@@ -82,6 +86,7 @@ private:
   UPtrAngle angle_;                               //!< Angular distribution
   UPtrDist energy_;                               //!< Energy distribution
   UPtrDist time_;                                 //!< Time distribution
+  std::unordered_set<int32_t> cells_;             //!< Cells to reject from
 };
 
 //==============================================================================

openmc/source.py

@@ -5,6 +5,7 @@ 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 +37,8 @@ class Source:
         Strength of the source
     particle : {'neutron', 'photon'}
         Source particle type
+    cells : iterable of int or openmc.Cell
+        Cells to reject based on
 
     Attributes
     ----------
@@ -57,11 +60,14 @@ class Source:
         Strength of the source
     particle : {'neutron', 'photon'}
         Source particle type
+    cells : iterable of int or openmc.Cell
+        Cells to reject based on
 
     """
 
     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',
+                 cells=None):
         self._space = None
         self._angle = None
         self._energy = None
@@ -69,6 +75,7 @@ class Source:
         self._file = None
         self._library = None
         self._parameters = None
+        self._cells = []
 
         if space is not None:
             self.space = space
@@ -86,6 +93,8 @@ class Source:
             self.parameters = parameters
         self.strength = strength
         self.particle = particle
+        if cells is not None:
+            self.cells = cells
 
     @property
     def file(self):
@@ -123,6 +132,10 @@ class Source:
     def particle(self):
         return self._particle
 
+    @property
+    def cells(self):
+        return self._cells
+
     @file.setter
     def file(self, filename):
         cv.check_type('source file', filename, str)
@@ -169,6 +182,13 @@ class Source:
         cv.check_value('source particle', particle, ['neutron', 'photon'])
         self._particle = particle
 
+    @cells.setter
+    def cells(self, cells):
+        self._cells = [
+            cell.id if isinstance(cell, openmc.Cell) else cell
+            for cell in cells
+        ]
+
     def to_xml_element(self):
         """Return XML representation of the source
 
@@ -196,6 +216,9 @@ 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.cells:
+            cells_elem = ET.SubElement(element, "cells")
+            cells_elem.text = ' '.join(str(x) for x in self.cells)
         return element
 
     @classmethod
@@ -251,6 +274,10 @@ class Source:
         if time is not None:
             source.time = Univariate.from_xml_element(time)
 
+        cells = elem.find('cells')
+        if cells is not None:
+            source.cells = [int(x) for x in cells.text.split()]
+
         return source
 
 

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,36 @@ IndependentSource::IndependentSource(pugi::xml_node node)
       double p[] {1.0};
       time_ = UPtrDist {new Discrete {T, p, 1}};
     }
+
+    // Check for cells to reject from
+    if (check_for_node(node, "cells")) {
+      auto cells = get_node_array<int>(node, "cells");
+      cells_.insert(cells.cbegin(), cells.cend());
+    }
   }
 }
 
 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 +190,22 @@ 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
+      if (!cells_.empty()) {
+        found = false;
+        for (const auto& coord : p.coord()) {
+          if (contains(cells_, model::cells[coord.cell]->id_)) {
+            found = true;
+            break;
           }
         }
       }
@@ -205,6 +221,8 @@ SourceSite IndependentSource::sample(uint64_t* seed) const
                     "definition.");
       }
     }
+
+    site.r = p.r();
   }
 
   // Sample angle