Allow spatial constraints on element sources within MeshSource (#3431)

include/openmc/source.h

@@ -139,6 +139,9 @@ public:
   DomainType domain_type() const { return domain_type_; }
   const std::unordered_set<int32_t>& domain_ids() const { return domain_ids_; }
 
+  // Setter for spatial distribution
+  void set_space(UPtrSpace space) { space_ = std::move(space); }
+
 protected:
   // Indicates whether derived class already handles constraints
   bool constraints_applied() const override { return true; }
@@ -205,6 +208,23 @@ typedef unique_ptr<Source> create_compiled_source_t(std::string parameters);
 //! Mesh-based source with different distributions for each element
 //==============================================================================
 
+// Helper class to sample spatial position on a single mesh element
+class MeshElementSpatial : public SpatialDistribution {
+public:
+  MeshElementSpatial(int32_t mesh_index, int elem_index)
+    : mesh_index_(mesh_index), elem_index_(elem_index)
+  {}
+
+  //! Sample a position from the distribution
+  //! \param seed Pseudorandom number seed pointer
+  //! \return Sampled position
+  Position sample(uint64_t* seed) const override;
+
+private:
+  int32_t mesh_index_ {C_NONE}; //!< Index in global meshes array
+  int elem_index_;              //! Index of mesh element
+};
+
 class MeshSource : public Source {
 public:
   // Constructors
@@ -219,18 +239,15 @@ public:
   double strength() const override { return space_->total_strength(); }
 
   // Accessors
-  const std::unique_ptr<Source>& source(int32_t i) const
+  const unique_ptr<IndependentSource>& source(int32_t i) const
   {
     return sources_.size() == 1 ? sources_[0] : sources_[i];
   }
 
-protected:
-  bool constraints_applied() const override { return true; }
-
 private:
   // Data members
-  unique_ptr<MeshSpatial> space_;           //!< Mesh spatial
-  vector<std::unique_ptr<Source>> sources_; //!< Source distributions
+  unique_ptr<MeshSpatial> space_;                 //!< Mesh spatial
+  vector<unique_ptr<IndependentSource>> sources_; //!< Source distributions
 };
 
 //==============================================================================

src/source.cpp

@@ -239,7 +239,9 @@ bool Source::satisfies_spatial_constraints(Position r) const
   if (!domain_ids_.empty()) {
     if (domain_type_ == DomainType::MATERIAL) {
       auto mat_index = geom_state.material();
-      if (mat_index != MATERIAL_VOID) {
+      if (mat_index == MATERIAL_VOID) {
+        accepted = false;
+      } else {
         accepted = contains(domain_ids_, model::materials[mat_index]->id());
       }
     } else {
@@ -525,6 +527,15 @@ CompiledSourceWrapper::~CompiledSourceWrapper()
 #endif
 }
 
+//==============================================================================
+// MeshElementSpatial implementation
+//==============================================================================
+
+Position MeshElementSpatial::sample(uint64_t* seed) const
+{
+  return model::meshes[mesh_index_]->sample_element(elem_index_, seed);
+}
+
 //==============================================================================
 // MeshSource implementation
 //==============================================================================
@@ -539,10 +550,23 @@ MeshSource::MeshSource(pugi::xml_node node) : Source(node)
   // read all source distributions and populate strengths vector for MeshSpatial
   // object
   for (auto source_node : node.children("source")) {
-    sources_.emplace_back(Source::create(source_node));
+    auto src = Source::create(source_node);
+    if (auto ptr = dynamic_cast<IndependentSource*>(src.get())) {
+      src.release();
+      sources_.emplace_back(ptr);
+    } else {
+      fatal_error(
+        "The source assigned to each element must be an IndependentSource.");
+    }
     strengths.push_back(sources_.back()->strength());
   }
 
+  // Set spatial distributions for each mesh element
+  for (int elem_index = 0; elem_index < sources_.size(); ++elem_index) {
+    sources_[elem_index]->set_space(
+      std::make_unique<MeshElementSpatial>(mesh_idx, elem_index));
+  }
+
   // the number of source distributions should either be one or equal to the
   // number of mesh elements
   if (sources_.size() > 1 && sources_.size() != mesh->n_bins()) {
@@ -556,29 +580,12 @@ MeshSource::MeshSource(pugi::xml_node node) : Source(node)
 
 SourceSite MeshSource::sample(uint64_t* seed) const
 {
-  // Sample the CDF defined in initialization above
+  // Sample a mesh element based on the relative strengths
   int32_t element = space_->sample_element_index(seed);
 
-  // Sample position and apply rejection on spatial domains
-  Position r;
-  do {
-    r = space_->mesh()->sample_element(element, seed);
-  } while (!this->satisfies_spatial_constraints(r));
-
-  SourceSite site;
-  while (true) {
-    // Sample source for the chosen element and replace the position
-    site = source(element)->sample_with_constraints(seed);
-    site.r = r;
-
-    // Apply other rejections
-    if (satisfies_energy_constraints(site.E) &&
-        satisfies_time_constraints(site.time)) {
-      break;
-    }
-  }
-
-  return site;
+  // Sample the distribution for the specific mesh element; note that the
+  // spatial distribution has been set for each element using MeshElementSpatial
+  return source(element)->sample_with_constraints(seed);
 }
 
 //==============================================================================

tests/unit_tests/test_source_mesh.py

@@ -1,5 +1,7 @@
 from itertools import product
 from pathlib import Path
+from math import sqrt
+import random
 
 import pytest
 import numpy as np
@@ -334,13 +336,10 @@ def test_umesh_source_independent(run_in_tmpdir, request, void_model, library):
     n_elements = 12_000
     model.settings.source = openmc.MeshSource(uscd_mesh, n_elements*[ind_source])
     model.export_to_model_xml()
-    try:
-        openmc.lib.init()
+    with openmc.lib.run_in_memory():
         openmc.lib.simulation_init()
         sites = openmc.lib.sample_external_source(10)
         openmc.lib.statepoint_write('statepoint.h5')
-    finally:
-        openmc.lib.finalize()
 
     with openmc.StatePoint('statepoint.h5') as sp:
         uscd_mesh = sp.meshes[uscd_mesh.id]
@@ -351,51 +350,53 @@ def test_umesh_source_independent(run_in_tmpdir, request, void_model, library):
         assert site.r in bounding_box
 
 
-def test_mesh_source_file(run_in_tmpdir):
-    # Creating a source file with a single particle
-    source_particle = openmc.SourceParticle(time=10.0)
-    openmc.write_source_file([source_particle], 'source.h5')
-    file_source = openmc.FileSource('source.h5')
+def test_mesh_source_constraints(run_in_tmpdir):
+    """Test application of constraints to underlying mesh element sources"""
 
+    # Create simple model with two cells
+    m1 = openmc.Material()
+    m1.add_nuclide('H1', 1.0)
+    m2 = m1.clone()
+    sph = openmc.Sphere(r=100, boundary_type='vacuum')
+    box1 = openmc.model.RectangularParallelepiped(-1, 0, -1, 1, -1, 1)
+    box2 = openmc.model.RectangularParallelepiped(0, 2, -1, 1, -1, 1)
+    cell1 = openmc.Cell(fill=m1, region=-box1)
+    cell2 = openmc.Cell(fill=m2, region=-box2)
+    outer = openmc.Cell(region=-sph & (+box1 | +box2))
     model = openmc.Model()
+    model.geometry = openmc.Geometry([cell1, cell2, outer])
 
-    rect_prism = openmc.model.RectangularParallelepiped(
-        -5.0, 5.0, -5.0, 5.0, -5.0, 5.0, boundary_type='vacuum')
-
-    mat = openmc.Material()
-    mat.add_nuclide('H1', 1.0)
-
-    model.geometry = openmc.Geometry([openmc.Cell(fill=mat, region=-rect_prism)])
-    model.settings.particles = 1000
-    model.settings.batches = 10
-    model.settings.run_mode = 'fixed source'
-
+    # Define a mesh covering the two cells: the first mesh element contains
+    # cell1 (-1 < x < 0) and the second element contains cells2 (0 < x < 2)
     mesh = openmc.RegularMesh()
-    mesh.lower_left = (-1, -2, -3)
-    mesh.upper_right = (2, 3, 4)
-    mesh.dimension = (1, 1, 1)
+    mesh.lower_left = (-3., -1., -1.)
+    mesh.upper_right = (3., 1., 1.)
+    mesh.dimension = (2, 1, 1)
 
-    model.settings.source = openmc.MeshSource(mesh, [file_source])
+    # Define a mesh source with a randomly chosen probability
+    p = random.random()
+    src1 = openmc.IndependentSource(strength=p, constraints={'domains': [cell1]})
+    src2 = openmc.IndependentSource(strength=1 - p, constraints={'domains': [cell2]})
+    model.settings.source = openmc.MeshSource(mesh, [src1, src2])
 
+    # Finish settings and export
+    model.settings.particles = 100
+    model.settings.batches = 1
     model.export_to_model_xml()
 
-    openmc.lib.init()
-    openmc.lib.simulation_init()
-    sites = openmc.lib.sample_external_source(10)
-    openmc.lib.simulation_finalize()
-    openmc.lib.finalize()
+    with openmc.lib.run_in_memory():
+        # Sample sites from the source
+        sites = openmc.lib.sample_external_source(N := 1000)
 
-    # The mesh bounds do not contain the point of the lone source site in the
-    # file source, so it should not appear in the set of source sites produced
-    # from the mesh source. Additionally, the source should be located within
-    # the mesh
-    bbox = mesh.bounding_box
-    for site in sites:
-        assert site.r != (0, 0, 0)
-        assert site.E == source_particle.E
-        assert site.u == source_particle.u
-        assert site.time == source_particle.time
-        assert site.r in bbox
+        # Check that all sites are either in cell1 or cell2
+        xs = np.array([s.r[0] for s in sites])
+        assert (xs >= -1.0).all()
+        assert (xs <= 2.0).all()
+
+        # Check that the correct percentage of the sites are in cell1
+        sigma = sqrt(p*(1- p)/N)
+        frac = xs[(-1.0 <= xs) & (xs <= 0.0)].size / N
+        assert frac == pytest.approx(p, abs=5*sigma)
 
 
 @pytest.mark.parametrize("mesh_type", ('rectangular', 'cylindrical', 'spherical'))