Support safety and reentrant boundaries for multi-universe ORANGE (#1140

)

* Add support for reentrant boundary crossing with multi-universe tracking
* Fixing set_dir such that it queries the correct universe for finding the normal
* Update safety distance for multiple levels
* Add testem3 style unit test for safety distance calculation
* Fix safety distance on multiple levels
* Address comments from merge request
* TestEM3 model updated to include world volume

src/orange/OrangeData.hh

@@ -17,13 +17,19 @@
 #include "geocel/BoundingBox.hh"
 
 #include "OrangeTypes.hh"
-#include "detail/BIHData.hh"
 #include "univ/detail/Types.hh"
 
+#include "detail/BIHData.hh"
+
 namespace celeritas
 {
 //---------------------------------------------------------------------------//
 // PARAMS
+//---------------------------------------------------------------------------//
+
+// VolumeId of exterior volume for all universe types
+static inline constexpr LocalVolumeId local_orange_outside_volume{0};
+
 //---------------------------------------------------------------------------//
 /*!
  * Scalar values particular to an ORANGE geometry instance.

src/orange/OrangeParams.cc

@@ -28,11 +28,12 @@
 #include "OrangeData.hh"  // IWYU pragma: associated
 #include "OrangeInput.hh"
 #include "OrangeTypes.hh"
+#include "univ/detail/LogicStack.hh"
+
 #include "detail/DepthCalculator.hh"
 #include "detail/RectArrayInserter.hh"
 #include "detail/UnitInserter.hh"
 #include "detail/UniverseInserter.hh"
-#include "univ/detail/LogicStack.hh"
 
 #if CELERITAS_USE_JSON
 #    include <nlohmann/json.hpp>
@@ -151,9 +152,14 @@ OrangeParams::OrangeParams(OrangeInput input)
         vol_labels_ = LabelIdMultiMap<VolumeId>{std::move(volume_labels)};
     }
 
-    // Safety currently only works for the single-universe case
-    supports_safety_ = host_data.simple_units.size() == 1
-                       && host_data.simple_units[SimpleUnitId{0}].simple_safety;
+    // Simple safety if all SimpleUnits have simple safety and no RectArrays
+    // are present
+    supports_safety_
+        = std::all_of(
+              host_data.simple_units[AllItems<SimpleUnitRecord>()].begin(),
+              host_data.simple_units[AllItems<SimpleUnitRecord>()].end(),
+              [](SimpleUnitRecord const& su) { return su.simple_safety; })
+          && host_data.rect_arrays.empty();
 
     CELER_VALIDATE(std::holds_alternative<UnitInput>(input.universes.front()),
                    << "global universe is not a SimpleUnit");

src/orange/OrangeTrackView.hh

@@ -439,7 +439,7 @@ CELER_FUNCTION bool OrangeTrackView::is_outside() const
     // Zeroth volume in outermost universe is always the exterior by
     // construction in ORANGE
     auto lsa = this->make_lsa(LevelId{0});
-    return lsa.vol() == LocalVolumeId{0};
+    return lsa.vol() == local_orange_outside_volume;
 }
 
 //---------------------------------------------------------------------------//
@@ -596,12 +596,6 @@ CELER_FUNCTION void OrangeTrackView::cross_boundary()
 
     if (CELER_UNLIKELY(this->boundary() == BoundaryResult::reentrant))
     {
-        if (sl != LevelId{0})
-        {
-            CELER_NOT_IMPLEMENTED(
-                "reentrant surfaces with multilevel ORANGE geometry");
-        }
-
         // Direction changed while on boundary leading to no change in
         // volume/surface. This is logically equivalent to a reflection.
         this->boundary(BoundaryResult::exiting);
@@ -705,17 +699,16 @@ CELER_FUNCTION void OrangeTrackView::set_dir(Real3 const& newdir)
 
     if (this->is_on_boundary())
     {
-        // Changing direction on a boundary is dangerous, as it could mean we
-        // don't leave the volume after all.
-        detail::UniverseIndexer ui(params_.universe_indexer_data);
+        // Changing direction on a boundary, which may result in not leaving
+        // current volume upon the cross_surface call
+        auto lsa = this->make_lsa(this->surface_level());
 
         TrackerVisitor visit_tracker{params_};
-        auto pos = this->pos();
         auto normal = visit_tracker(
-            [&pos, local_surface = this->surf()](auto&& t) {
+            [pos = lsa.pos(), local_surface = this->surf()](auto&& t) {
                 return t.normal(pos, local_surface);
             },
-            UniverseId{0});
+            lsa.universe());
 
         // Normal is in *local* coordinates but newdir is in *global*: rotate
         // up to check
@@ -922,18 +915,28 @@ OrangeTrackView::find_next_step_impl(detail::Intersection isect)
 //---------------------------------------------------------------------------//
 /*!
  * Find the distance to the nearest boundary in any direction.
+ *
+ * The safety distance at a given point is the minimum safety distance over all
+ * levels, since surface deduplication can potentionally elide bounding
+ * surfaces at more deeply embedded levels.
  */
 CELER_FUNCTION real_type OrangeTrackView::find_safety()
 {
     CELER_EXPECT(!this->is_on_boundary());
-    auto lsa = this->make_lsa();
-
-    CELER_ASSERT(lsa.universe() == UniverseId{0});
 
     TrackerVisitor visit_tracker{params_};
-    return visit_tracker(
-        [&lsa](auto&& t) { return t.safety(lsa.pos(), lsa.vol()); },
-        lsa.universe());
+
+    real_type min_safety_dist = numeric_limits<real_type>::infinity();
+
+    for (auto lev : range(LevelId{this->level() + 1}))
+    {
+        auto lsa = this->make_lsa(lev);
+        auto sd = visit_tracker(
+            [&lsa](auto&& t) { return t.safety(lsa.pos(), lsa.vol()); },
+            lsa.universe());
+        min_safety_dist = celeritas::min(min_safety_dist, sd);
+    }
+    return min_safety_dist;
 }
 
 //---------------------------------------------------------------------------//

src/orange/detail/UnitInserter.cc

@@ -285,8 +285,12 @@ UniverseId UnitInserter::operator()(UnitInput const& inp)
     {
         unit.background = LocalVolumeId(inp.volumes.size() - 1);
     }
+
+    // Simple safety if all volumes provide support, excluding the external
+    // volume, which appears first in the list
+    static_assert(local_orange_outside_volume == LocalVolumeId{0});
     unit.simple_safety = std::all_of(
-        vol_records.begin(), vol_records.end(), [](VolumeRecord const& v) {
+        vol_records.begin() + 1, vol_records.end(), [](VolumeRecord const& v) {
             return supports_simple_safety(v.flags);
         });
 
@@ -375,7 +379,7 @@ void UnitInserter::process_daughter(VolumeRecord* vol_record,
     daughter.transform_id = insert_transform_(daughter_input.transform);
 
     vol_record->daughter_id = daughters_.push_back(daughter);
-    vol_record->flags &= VolumeRecord::embedded_universe;
+    vol_record->flags |= VolumeRecord::embedded_universe;
 }
 
 //---------------------------------------------------------------------------//

test/orange/Orange.test.cc

@@ -95,6 +95,12 @@ class HexArrayTest : public OrangeTest
     void SetUp() override { this->build_geometry("hex_array.org.json"); }
 };
 
+#define TestEM3Test TEST_IF_CELERITAS_JSON(TestEM3Test)
+class TestEM3Test : public OrangeTest
+{
+    void SetUp() override { this->build_geometry("testem3.org.json"); }
+};
+
 #define ShiftTrackerTest TEST_IF_CELERITAS_JSON(ShiftTrackerTest)
 class ShiftTrackerTest : public OrangeTest
 {
@@ -926,6 +932,41 @@ TEST_F(UniversesTest, cross_between_daughters)
     EXPECT_EQ("bob.mz", this->params().id_to_label(geo.surface_id()).name);
 }
 
+// Change direction on a universe boundary to reenter the cell
+TEST_F(UniversesTest, reentrant)
+{
+    auto geo = this->make_track_view();
+
+    // Initialize in innermost universe
+    geo = Initializer_t{{0.25, -3.7, 0.7}, {0, 1, 0}};
+    auto next = geo.find_next_step();
+    EXPECT_SOFT_EQ(0.2, next.distance);
+
+    // Move to universe boundary
+    geo.move_to_boundary();
+    EXPECT_EQ("inner_c.py", this->params().id_to_label(geo.surface_id()).name);
+    EXPECT_EQ("patty", this->params().id_to_label(geo.volume_id()).name);
+    EXPECT_VEC_SOFT_EQ(Real3({0.25, -3.5, 0.7}), geo.pos());
+
+    // Change direction on the universe boundary such that we are no longer
+    // exiting the universe
+    geo.set_dir({0, -1, 0});
+
+    // Remain in same cell after crossing boundary
+    geo.cross_boundary();
+    EXPECT_EQ("inner_c.py", this->params().id_to_label(geo.surface_id()).name);
+    EXPECT_EQ("patty", this->params().id_to_label(geo.volume_id()).name);
+    EXPECT_VEC_SOFT_EQ(Real3({0.25, -3.5, 0.7}), geo.pos());
+
+    // Make sure we can take another step after calling cross_boundary
+    next = geo.find_next_step();
+    EXPECT_SOFT_EQ(0.5, next.distance);
+    geo.move_to_boundary();
+    EXPECT_EQ("inner_a.my", this->params().id_to_label(geo.surface_id()).name);
+    EXPECT_EQ("patty", this->params().id_to_label(geo.volume_id()).name);
+    EXPECT_VEC_SOFT_EQ(Real3({0.25, -4, 0.7}), geo.pos());
+}
+
 TEST_F(RectArrayTest, params)
 {
     OrangeParams const& geo = this->params();
@@ -1081,6 +1122,22 @@ TEST_F(HexArrayTest, track_out)
     EXPECT_VEC_CLOSE(d2b, refd2b, real_type(1e-5), real_type(1e-5));
 }
 
+// Test safety distance within a geometry that supports simple safety
+TEST_F(TestEM3Test, safety)
+{
+    EXPECT_FALSE(this->params().supports_safety());
+
+    auto geo = this->make_track_view();
+
+    // Initialize in innermost universe, near the universe boundary
+    geo = Initializer_t{{19.99, 19.9, 19.9}, {0, 1, 0}};
+    EXPECT_SOFT_EQ(0.01, geo.find_safety());
+
+    // Initialize on the other side of the same volume
+    geo = Initializer_t{{19.42, 19.9, 19.9}, {0, 1, 0}};
+    EXPECT_SOFT_EQ(0.01, geo.find_safety());
+}
+
 TEST_F(ShiftTrackerTest, host)
 {
     std::vector<Real3> pos{