Overlap detection for plotter (#3969)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

include/openmc/geometry.h

@@ -3,9 +3,12 @@
 
 #include <cmath>
 #include <cstdint>
+#include <unordered_map>
+#include <vector>
 
 #include "openmc/array.h"
 #include "openmc/constants.h"
+#include "openmc/random_ray/source_region.h" // For hash_combine
 #include "openmc/vector.h"
 
 namespace openmc {
@@ -13,6 +16,34 @@ namespace openmc {
 class BoundaryInfo;
 class GeometryState;
 
+//==============================================================================
+//! OverlapKey to store cell and universe data of a single overlap, along with
+//! a functor for hashing an OverlapKey into an unordered_map.
+//==============================================================================
+
+struct OverlapKey {
+  int universe_id;
+  int cell1_id;
+  int cell2_id;
+
+  bool operator==(const OverlapKey& other) const
+  {
+    return universe_id == other.universe_id && cell1_id == other.cell1_id &&
+           cell2_id == other.cell2_id;
+  }
+};
+
+struct OverlapKeyHash {
+  std::size_t operator()(const OverlapKey& k) const
+  {
+    size_t seed = 0;
+    hash_combine(seed, k.universe_id);
+    hash_combine(seed, k.cell1_id);
+    hash_combine(seed, k.cell2_id);
+    return seed;
+  }
+};
+
 //==============================================================================
 // Global variables
 //==============================================================================
@@ -24,6 +55,10 @@ extern "C" int n_coord_levels; //!< Number of CSG coordinate levels
 
 extern vector<int64_t> overlap_check_count;
 
+// Overlap data structures get cleared every slice_data run
+extern vector<OverlapKey> overlap_keys;
+extern std::unordered_map<OverlapKey, int, OverlapKeyHash> overlap_key_index;
+
 } // namespace model
 
 //==============================================================================
@@ -38,8 +73,7 @@ inline bool coincident(double d1, double d2)
 //==============================================================================
 //! Check for overlapping cells at a particle's position.
 //==============================================================================
-
-bool check_cell_overlap(GeometryState& p, bool error = true);
+int check_cell_overlap(GeometryState& p, bool error = true);
 
 //==============================================================================
 //! Get the cell instance for a particle at the specified universe level

include/openmc/plot.h

@@ -5,6 +5,7 @@
 #include <sstream>
 #include <unordered_map>
 #include <unordered_set>
+#include <vector>
 
 #include "openmc/tensor.h"
 #include "pugixml.hpp"
@@ -155,7 +156,7 @@ struct IdData {
   // Methods
   void set_value(size_t y, size_t x, const Particle& p, int level,
     Filter* filter = nullptr, FilterMatch* match = nullptr);
-  void set_overlap(size_t y, size_t x);
+  void set_overlap(size_t y, size_t x, int overlap_idx);
 
   // Members
   tensor::Tensor<int32_t> data_; //!< 2D array of cell & material ids
@@ -168,7 +169,7 @@ struct PropertyData {
   // Methods
   void set_value(size_t y, size_t x, const Particle& p, int level,
     Filter* filter = nullptr, FilterMatch* match = nullptr);
-  void set_overlap(size_t y, size_t x);
+  void set_overlap(size_t y, size_t x, int overlap_idx);
 
   // Members
   tensor::Tensor<double> data_; //!< 2D array of temperature & density data
@@ -181,7 +182,7 @@ struct RasterData {
   // Methods
   void set_value(size_t y, size_t x, const Particle& p, int level,
     Filter* filter = nullptr, FilterMatch* match = nullptr);
-  void set_overlap(size_t y, size_t x);
+  void set_overlap(size_t y, size_t x, int overlap_idx);
 
   // Members
   tensor::Tensor<int32_t>
@@ -278,8 +279,11 @@ T SlicePlotBase::get_map(int32_t filter_index) const
         if (found_cell) {
           data.set_value(y, x, p, j, filter, &match);
         }
-        if (show_overlaps_ && check_cell_overlap(p, false)) {
-          data.set_overlap(y, x);
+        if (show_overlaps_) {
+          int overlap_idx = check_cell_overlap(p, false);
+          if (overlap_idx >= 0) {
+            data.set_overlap(y, x, overlap_idx);
+          }
         }
       } // inner for
     }

openmc/lib/plot.py

@@ -255,6 +255,34 @@ def property_map(plot):
     return prop_data
 
 
+_dll.openmc_slice_data_overlap_count.argtypes = [POINTER(c_size_t)]
+_dll.openmc_slice_data_overlap_count.restype = c_int
+_dll.openmc_slice_data_overlap_count.errcheck = _error_handler
+
+_dll.openmc_slice_data_overlap_info.argtypes = [c_size_t, POINTER(c_int32)]
+_dll.openmc_slice_data_overlap_info.restype = c_int
+_dll.openmc_slice_data_overlap_info.errcheck = _error_handler
+
+
+# Python wrappings for overlap functions
+def slice_data_overlap_count():
+    count = c_size_t()
+    _dll.openmc_slice_data_overlap_count(count)
+    return count.value
+
+
+def slice_data_overlap_info():
+    n = slice_data_overlap_count()
+    overlap_info = np.empty(n * 3, dtype=np.int32)
+
+    if n > 0:
+        _dll.openmc_slice_data_overlap_info(
+            n,
+            overlap_info.ctypes.data_as(POINTER(c_int32)),
+        )
+    return overlap_info, n
+
+
 _dll.openmc_get_plot_index.argtypes = [c_int32, POINTER(c_int32)]
 _dll.openmc_get_plot_index.restype = c_int
 _dll.openmc_get_plot_index.errcheck = _error_handler

src/geometry.cpp

@@ -26,16 +26,22 @@ int n_coord_levels;
 
 vector<int64_t> overlap_check_count;
 
+vector<OverlapKey> overlap_keys;
+std::unordered_map<OverlapKey, int, OverlapKeyHash> overlap_key_index;
+
 } // namespace model
 
 //==============================================================================
 // Non-member functions
 //==============================================================================
 
-bool check_cell_overlap(GeometryState& p, bool error)
+int check_cell_overlap(GeometryState& p, bool error)
 {
   int n_coord = p.n_coord();
 
+  // If no overlap found, return a nonphysical index
+  int overlap_index = -1;
+
   // Loop through each coordinate level
   for (int j = 0; j < n_coord; j++) {
     Universe& univ = *model::universes[p.coord(j).universe()];
@@ -44,21 +50,40 @@ bool check_cell_overlap(GeometryState& p, bool error)
     for (auto index_cell : univ.cells_) {
       Cell& c = *model::cells[index_cell];
       if (c.contains(p.coord(j).r(), p.coord(j).u(), p.surface())) {
+#pragma omp atomic
+        ++model::overlap_check_count[index_cell];
         if (index_cell != p.coord(j).cell()) {
           if (error) {
             fatal_error(
               fmt::format("Overlapping cells detected: {}, {} on universe {}",
                 c.id_, model::cells[p.coord(j).cell()]->id_, univ.id_));
           }
-          return true;
+
+          // With no fatal error (plotter is calling), now adds overlaps and
+          // ensures order does not matter when making overlap key
+          int cell_a = model::cells[index_cell]->id_;
+          int cell_b = model::cells[p.coord(j).cell()]->id_;
+          int a = std::min(cell_a, cell_b);
+          int b = std::max(cell_a, cell_b);
+          OverlapKey key {univ.id_, a, b};
+#pragma omp critical(overlap_key_update)
+          {
+            auto it = model::overlap_key_index.find(key);
+            if (it != model::overlap_key_index.end()) {
+              overlap_index = it->second; // already exists, reuse index
+            } else {
+              int idx = int(model::overlap_keys.size());
+              model::overlap_keys.push_back(key);
+              model::overlap_key_index[key] = idx;
+              overlap_index = idx;
+            }
+          }
+          break;
         }
-#pragma omp atomic
-        ++model::overlap_check_count[index_cell];
       }
     }
   }
-
-  return false;
+  return overlap_index;
 }
 
 //==============================================================================

src/plot.cpp

@@ -73,7 +73,7 @@ void IdData::set_value(size_t y, size_t x, const Particle& p, int level,
   }
 }
 
-void IdData::set_overlap(size_t y, size_t x)
+void IdData::set_overlap(size_t y, size_t x, int /*overlap_idx*/)
 {
   for (size_t k = 0; k < data_.shape(2); ++k)
     data_(y, x, k) = OVERLAP;
@@ -94,7 +94,7 @@ void PropertyData::set_value(size_t y, size_t x, const Particle& p, int level,
   }
 }
 
-void PropertyData::set_overlap(size_t y, size_t x)
+void PropertyData::set_overlap(size_t y, size_t x, int /*overlap_idx*/)
 {
   data_(y, x) = OVERLAP;
 }
@@ -154,12 +154,12 @@ void RasterData::set_value(size_t y, size_t x, const Particle& p, int level,
   }
 }
 
-void RasterData::set_overlap(size_t y, size_t x)
+void RasterData::set_overlap(size_t y, size_t x, int overlap_idx)
 {
   // Set cell, instance, and material to OVERLAP, but preserve filter bin
   id_data_(y, x, 0) = OVERLAP;
   id_data_(y, x, 1) = OVERLAP;
-  id_data_(y, x, 2) = OVERLAP;
+  id_data_(y, x, 2) = OVERLAP - overlap_idx - 1;
   // Note: id_data_(y, x, 3) is NOT overwritten - preserves filter bin for tally
   // plotting
 
@@ -1991,10 +1991,12 @@ extern "C" int openmc_slice_data(const double origin[3], const double u_span[3],
     plot_params.show_overlaps_ = color_overlaps;
     plot_params.slice_level_ = level;
 
+    // Clear overlap data structures on new slice call
+    model::overlap_keys.clear();
+    model::overlap_key_index.clear();
+
     // Use get_map<RasterData> to generate data
     auto data = plot_params.get_map<RasterData>(filter_index);
-
-    // Copy geometry data
     std::copy(data.id_data_.begin(), data.id_data_.end(), geom_data);
 
     // Copy property data if requested
@@ -2010,6 +2012,32 @@ extern "C" int openmc_slice_data(const double origin[3], const double u_span[3],
   return 0;
 }
 
+// Gets the number of overlaps that we need data for
+extern "C" int openmc_slice_data_overlap_count(size_t* count)
+{
+  if (!count) {
+    set_errmsg("Null pointer passed for overlap count.");
+    return OPENMC_E_INVALID_ARGUMENT;
+  }
+  *count = model::overlap_keys.size();
+
+  return 0;
+}
+
+// Plotter pre-allocates array size based on what is returned with
+// overlap_count; populates an array of size 3*count
+extern "C" int openmc_slice_data_overlap_info(
+  size_t count, int32_t* overlap_info)
+{
+  for (size_t i = 0; i < count; ++i) {
+    overlap_info[i * 3] = model::overlap_keys[i].universe_id;
+    overlap_info[i * 3 + 1] = model::overlap_keys[i].cell1_id;
+    overlap_info[i * 3 + 2] = model::overlap_keys[i].cell2_id;
+  }
+
+  return 0;
+}
+
 extern "C" int openmc_get_plot_index(int32_t id, int32_t* index)
 {
   auto it = model::plot_map.find(id);

tests/unit_tests/test_slice_data_overlap_info.py

@@ -0,0 +1,87 @@
+import pytest
+import numpy as np
+import openmc
+import openmc.lib
+
+# Sentinel value matching _OVERLAP in plotmodel.py and OVERLAP in plot.cpp
+_OVERLAP = -3
+
+
+@pytest.fixture(scope='module')
+def overlap_model():
+    openmc.reset_auto_ids()
+
+    # Three cylinders: cyl1 and cyl2 overlap near x=0, cyl2 and cyl3 overlap
+    # near x=4. This gives us two spatially distinct overlap regions in one model.
+    mat1 = openmc.Material(components={'H1': 1.0})
+    mat2 = openmc.Material(components={'H1': 1.0})
+    mat3 = openmc.Material(components={'H1': 1.0})
+
+    # cyl1 and cyl2 overlap on the left, cyl2 and cyl3 overlap on the right
+    cyl1 = openmc.ZCylinder(x0=-2.0, r=2.5)
+    cyl2 = openmc.ZCylinder(x0=0.0, r=2.5)
+    cyl3 = openmc.ZCylinder(x0=2.0, r=2.5)
+    boundary = openmc.Sphere(r=20.0, boundary_type='vacuum')
+    cell1 = openmc.Cell(region=-cyl1, fill=mat1)
+    cell2 = openmc.Cell(region=-cyl2, fill=mat2)
+    cell3 = openmc.Cell(region=-cyl3, fill=mat3)
+    cell_outside = openmc.Cell(region=+cyl1 & +cyl2 & +cyl3 & -boundary)
+    geometry = openmc.Geometry([cell1, cell2, cell3, cell_outside])
+
+    settings = openmc.Settings()
+    settings.run_mode = 'fixed source'
+    settings.particles = 100
+    settings.batches = 1
+    model = openmc.Model(geometry=geometry, settings=settings)
+
+    with openmc.lib.TemporarySession(model, args=['-s', '1']):
+        yield
+
+
+def run_slice(origin=(0.0, 0.0, 0.0), width=(10.0, 6.0), show_overlaps=True):
+    # Helper that runs a slice over a region covering both overlap zones
+    geom_data, _ = openmc.lib.slice_data(
+        origin=origin,
+        width=width,
+        basis='xy',
+        pixels=(100, 60),
+        show_overlaps=show_overlaps,
+        include_properties=False,
+    )
+    return geom_data
+
+
+def test_overlaps_enabled(overlap_model):
+    # Run a single slice with overlap detection enabled and check all
+    # expected properties in one pass.
+    geom_data = run_slice()
+    overlap_info, n = openmc.lib.slice_data_overlap_info()
+    mat_ids = geom_data[:, :, 2]
+
+    # mat_ids should contain values more negative than _OVERLAP; RasterData
+    # encodes each unique overlap as OVERLAP - overlap_idx - 1 into slot 2.
+    assert np.any(mat_ids < _OVERLAP)
+
+    # overlap_keys should have 2 entries for the two distinct overlapping
+    # cylinder pairs in this model.
+    assert n == 2, f"Expected exactly 2 overlap entries, got {n}"
+
+    # Each entry is a (universe_id, cell1_id, cell2_id) triple; verify values.
+    for i in range(n):
+        universe_id = int(overlap_info[i * 3])
+        cell1_id = int(overlap_info[i * 3 + 1])
+        cell2_id = int(overlap_info[i * 3 + 2])
+        assert universe_id == 1
+        assert cell1_id in {1, 2, 3}
+        assert cell2_id in {1, 2, 3}
+        assert cell1_id != cell2_id
+
+
+def test_overlaps_disabled(overlap_model):
+    # With show_overlaps=False, set_overlap is never called and overlap_keys
+    # is never written to, so the image and map should both be clean.
+    geom_data = run_slice(show_overlaps=False)
+    _, n = openmc.lib.slice_data_overlap_info()
+
+    assert not np.any(geom_data[:, :, 2] < _OVERLAP)
+    assert n == 0