Add multilevel "volume instance" (#1461)

* Add volume instance ID and move level ID * Add level and volume instance ID accessors * Add visit-volumes test * Add PV depth-first search * Switch to breadth-first * Add PV labels to G4 geometry * Add PV to VecGeom * Add test of PV stack * Delete materials from ml test * Add max depth * Use depth-first search
celeritas-project · Oct 23, 2024 · 1b6b8b2 · 1b6b8b2
1 parent 77ffa18
commit 1b6b8b2
Show file tree

Hide file tree

Showing 24 changed files with 806 additions and 102 deletions.
diff --git a/src/celeritas/ext/GeantImporter.cc b/src/celeritas/ext/GeantImporter.cc
@@ -1304,7 +1304,7 @@ GeantImporter::import_volumes(bool unique_volumes) const
                 volume.name = make_gdml_name(lv);
             }
         },
-        *world_->GetLogicalVolume());
+        *world_);
 
     CELER_LOG(debug) << "Loaded " << count << " volumes with "
                      << (unique_volumes ? "uniquified" : "original")

diff --git a/src/geocel/GeoParamsInterface.hh b/src/geocel/GeoParamsInterface.hh
@@ -16,6 +16,7 @@
 #include "Types.hh"
 
 class G4LogicalVolume;
+class G4VPhysicalVolume;
 
 namespace celeritas
 {
@@ -34,6 +35,7 @@ class GeoParamsInterface
     //! \name Type aliases
     using SpanConstVolumeId = Span<VolumeId const>;
     using VolumeMap = LabelIdMultiMap<VolumeId>;
+    using VolInstanceMap = LabelIdMultiMap<VolumeInstanceId>;
     //!@}
 
   public:
@@ -43,14 +45,23 @@ class GeoParamsInterface
     //! Outer bounding box of geometry
     virtual BBox const& bbox() const = 0;
 
+    //! Maximum nested scene/volume depth
+    virtual LevelId::size_type max_depth() const = 0;
+
     //// VOLUMES ////
 
     //! Get volume metadata
     virtual VolumeMap const& volumes() const = 0;
 
+    //! Get volume instance metadata
+    virtual VolInstanceMap const& volume_instances() const = 0;
+
     //! Get the volume ID corresponding to a Geant4 logical volume
     virtual VolumeId find_volume(G4LogicalVolume const* volume) const = 0;
 
+    //! Get the Geant4 PV corresponding to a volume instance
+    virtual G4VPhysicalVolume const* id_to_pv(VolumeInstanceId id) const = 0;
+
     //// DEPRECATED: remove in v0.6 ////
 
     //! Number of volumes

diff --git a/src/geocel/GeoParamsOutput.cc b/src/geocel/GeoParamsOutput.cc
@@ -43,6 +43,7 @@ void GeoParamsOutput::output(JsonPimpl* j) const
 
     obj["supports_safety"] = geo_->supports_safety();
     obj["bbox"] = geo_->bbox();
+    obj["max_depth"] = geo_->max_depth();
 
     // Save volume names
     {

diff --git a/src/geocel/Types.hh b/src/geocel/Types.hh
@@ -36,15 +36,21 @@ using SquareMatrixReal3 = SquareMatrix<real_type, 3>;
 
 //---------------------------------------------------------------------------//
 
+//! Type-safe "level", i.e., depth of embedded unit/scene/volume
+using LevelId = OpaqueId<struct Level_>;
+
 //! Identifier for a material fill
 using GeoMaterialId = OpaqueId<struct GeoMaterial_>;
 
 //! Identifier for a surface (for surface-based geometries)
 using SurfaceId = OpaqueId<struct Surface_>;
 
-//! Identifier for a geometry volume
+//! Identifier for a geometry volume that may be repeated
 using VolumeId = OpaqueId<struct Volume_>;
 
+//! Identifier for an instance of a geometry volume (aka physical/placed)
+using VolumeInstanceId = OpaqueId<struct VolumeInstance_>;
+
 //---------------------------------------------------------------------------//
 // ENUMERATIONS
 //---------------------------------------------------------------------------//

diff --git a/src/geocel/g4/GeantGeoParams.cc b/src/geocel/g4/GeantGeoParams.cc
@@ -11,8 +11,10 @@
 #include <G4GeometryManager.hh>
 #include <G4LogicalVolume.hh>
 #include <G4LogicalVolumeStore.hh>
+#include <G4PhysicalVolumeStore.hh>
 #include <G4Transportation.hh>
 #include <G4TransportationManager.hh>
+#include <G4VPhysicalVolume.hh>
 #include <G4VSolid.hh>
 #include <G4Version.hh>
 #include <G4VisExtent.hh>
@@ -42,7 +44,7 @@ namespace
 {
 //---------------------------------------------------------------------------//
 std::vector<Label>
-get_volume_labels(G4LogicalVolume const& world, bool unique_volumes)
+get_volume_labels(G4VPhysicalVolume const& world, bool unique_volumes)
 {
     std::vector<Label> labels;
     visit_geant_volumes(
@@ -65,6 +67,57 @@ get_volume_labels(G4LogicalVolume const& world, bool unique_volumes)
     return labels;
 }
 
+//---------------------------------------------------------------------------//
+std::vector<Label> get_pv_labels(G4VPhysicalVolume const& world)
+{
+    std::vector<Label> labels;
+    std::unordered_map<G4VPhysicalVolume const*, int> max_depth;
+
+    visit_geant_volume_instances(
+        [&labels, &max_depth](G4VPhysicalVolume const& pv, int depth) {
+            auto&& [iter, inserted] = max_depth.insert({&pv, depth});
+            if (!inserted)
+            {
+                if (iter->second >= depth)
+                {
+                    // Already visited PV at this depth or more
+                    return false;
+                }
+                // Update the max depth
+                iter->second = depth;
+            }
+
+            auto i = static_cast<std::size_t>(pv.GetInstanceID());
+            if (i >= labels.size())
+            {
+                labels.resize(i + 1);
+            }
+            if (labels[i].empty())
+            {
+                labels[i] = Label::from_geant(pv.GetName());
+                CELER_ASSERT(!labels[i].empty());
+            }
+            return true;
+        },
+        world);
+    return labels;
+}
+
+//---------------------------------------------------------------------------//
+LevelId::size_type get_max_depth(G4VPhysicalVolume const& world)
+{
+    int result{0};
+    visit_geant_volume_instances(
+        [&result](G4VPhysicalVolume const&, int level) {
+            result = max(level, result);
+            return true;
+        },
+        world);
+    CELER_ENSURE(result >= 0);
+    // Maximum "depth" is one greater than "highest level"
+    return static_cast<LevelId::size_type>(result + 1);
+}
+
 //---------------------------------------------------------------------------//
 }  // namespace
 
@@ -179,11 +232,30 @@ VolumeId GeantGeoParams::find_volume(G4LogicalVolume const* volume) const
     return result;
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Get the Geant4 physical volume corresponding to a volume instance ID.
+ *
+ * If the input ID is false, a null pointer will be returned.
+ */
+G4VPhysicalVolume const* GeantGeoParams::id_to_pv(VolumeInstanceId id) const
+{
+    CELER_EXPECT(!id || id < vol_instances_.size());
+    if (!id)
+    {
+        return nullptr;
+    }
+
+    G4PhysicalVolumeStore* pv_store = G4PhysicalVolumeStore::GetInstance();
+    CELER_ASSERT(id < pv_store->size());
+    return (*pv_store)[id.unchecked_get()];
+}
+
 //---------------------------------------------------------------------------//
 /*!
  * Get the Geant4 logical volume corresponding to a volume ID.
  *
- * If the input volume ID is false, a null pointer will be returned.
+ * If the input volume ID is unassigned, a null pointer will be returned.
  */
 G4LogicalVolume const* GeantGeoParams::id_to_lv(VolumeId id) const
 {
@@ -225,15 +297,16 @@ void GeantGeoParams::build_metadata()
 
     ScopedMem record_mem("GeantGeoParams.build_metadata");
 
-    auto const* world_lv = host_ref_.world->GetLogicalVolume();
-    CELER_ASSERT(world_lv);
-
     // Construct volume labels
-    volumes_ = LabelIdMultiMap<VolumeId>(
-        "volume", get_volume_labels(*world_lv, !loaded_gdml_));
+    volumes_ = VolumeMap{"volume",
+                         get_volume_labels(*host_ref_.world, !loaded_gdml_)};
+    vol_instances_
+        = VolInstanceMap{"volume instance", get_pv_labels(*host_ref_.world)};
+    max_depth_ = get_max_depth(*host_ref_.world);
 
     // Save world bbox (NOTE: assumes no transformation on PV?)
-    bbox_ = [world_lv] {
+    bbox_ = [world_lv = host_ref_.world->GetLogicalVolume()] {
+        CELER_EXPECT(world_lv);
         G4VSolid const* solid = world_lv->GetSolid();
         CELER_ASSERT(solid);
         G4VisExtent const& extent = solid->GetExtent();

diff --git a/src/geocel/g4/GeantGeoParams.hh b/src/geocel/g4/GeantGeoParams.hh
@@ -57,14 +57,23 @@ class GeantGeoParams final : public GeoParamsInterface,
     //! Outer bounding box of geometry
     BBox const& bbox() const final { return bbox_; }
 
+    // Maximum nested scene/volume depth
+    LevelId::size_type max_depth() const final { return max_depth_; }
+
     //// VOLUMES ////
 
-    // Get volume metadata
+    // Get (logical) volume metadata
     inline VolumeMap const& volumes() const final;
 
+    // Get (physical) volume instance metadata
+    inline VolInstanceMap const& volume_instances() const final;
+
     // Get the volume ID corresponding to a Geant4 logical volume
     VolumeId find_volume(G4LogicalVolume const* volume) const final;
 
+    // Get the Geant4 physical volume corresponding to a volume instance ID
+    G4VPhysicalVolume const* id_to_pv(VolumeInstanceId vol_id) const final;
+
     // Get the Geant4 logical volume corresponding to a volume ID
     G4LogicalVolume const* id_to_lv(VolumeId vol_id) const;
 
@@ -92,7 +101,9 @@ class GeantGeoParams final : public GeoParamsInterface,
 
     // Host metadata/access
     VolumeMap volumes_;
+    VolInstanceMap vol_instances_;
     BBox bbox_;
+    LevelId::size_type max_depth_{0};
 
     // Host/device storage and reference
     HostRef host_ref_;
@@ -111,11 +122,24 @@ class GeantGeoParams final : public GeoParamsInterface,
 //---------------------------------------------------------------------------//
 /*!
  * Get volume metadata.
+ *
+ * Volumes correspond directly to Geant4 logical volumes.
  */
 auto GeantGeoParams::volumes() const -> VolumeMap const&
 {
     return volumes_;
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Get volume instance metadata.
+ *
+ * Volume instances correspond directly to Geant4 physical volumes.
+ */
+auto GeantGeoParams::volume_instances() const -> VolInstanceMap const&
+{
+    return vol_instances_;
+}
+
 //---------------------------------------------------------------------------//
 }  // namespace celeritas
diff --git a/src/geocel/g4/GeantGeoTrackView.hh b/src/geocel/g4/GeantGeoTrackView.hh
@@ -85,9 +85,14 @@ class GeantGeoTrackView
     CELER_FORCEINLINE Real3 const& dir() const { return dir_; }
     //!@}
 
-    // Get the volume ID in the current cell.
-    CELER_FORCEINLINE VolumeId volume_id() const;
-    CELER_FORCEINLINE int volume_physid() const;
+    // Get the volume ID in the lowest level volume.
+    inline VolumeId volume_id() const;
+    // Get the physical volume ID in the current cell
+    inline VolumeInstanceId volume_instance_id() const;
+    // Get the depth in the geometry hierarchy
+    inline LevelId level() const;
+    // Get the volume instance ID for all levels
+    inline void volume_instance_id(Span<VolumeInstanceId> levels) const;
 
     //!@{
     //! VecGeom states are never "on" a surface
@@ -96,9 +101,9 @@ class GeantGeoTrackView
     //!@}
 
     // Whether the track is outside the valid geometry region
-    CELER_FORCEINLINE bool is_outside() const;
+    inline bool is_outside() const;
     // Whether the track is exactly on a surface
-    CELER_FORCEINLINE bool is_on_boundary() const;
+    inline bool is_on_boundary() const;
     //! Whether the last operation resulted in an error
     CELER_FORCEINLINE bool failed() const { return false; }
 
@@ -255,20 +260,52 @@ VolumeId GeantGeoTrackView::volume_id() const
 /*!
  * Get the physical volume ID in the current cell.
  */
-int GeantGeoTrackView::volume_physid() const
+VolumeInstanceId GeantGeoTrackView::volume_instance_id() const
 {
     CELER_EXPECT(!this->is_outside());
     G4VPhysicalVolume* pv = touch_handle_()->GetVolume(0);
     if (!pv)
-        return -1;
-    return pv->GetInstanceID();
+        return {};
+    return id_cast<VolumeInstanceId>(pv->GetInstanceID());
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Get the depth in the geometry hierarchy.
+ */
+LevelId GeantGeoTrackView::level() const
+{
+    auto* touch = touch_handle_();
+    return id_cast<LevelId>(touch->GetHistoryDepth());
+}
+
+//---------------------------------------------------------------------------//
+/*!
+ * Get the volume instance ID at every level.
+ *
+ * The input span size must be equal to the value of "level" plus one. The
+ * top-most level ("world" or level zero) starts at index zero and moves
+ * downward. Note that Geant4 uses the \em reverse nomenclature.
+ */
+void GeantGeoTrackView::volume_instance_id(Span<VolumeInstanceId> levels) const
+{
+    CELER_EXPECT(levels.size() == this->level().get() + 1);
+
+    auto* touch = touch_handle_();
+    auto const max_depth = static_cast<size_type>(touch->GetHistoryDepth());
+    for (auto lev : range(levels.size()))
+    {
+        G4VPhysicalVolume* pv = touch->GetVolume(max_depth - lev);
+        CELER_ASSERT(pv);
+        levels[lev] = id_cast<VolumeInstanceId>(pv->GetInstanceID());
+    }
 }
 
 //---------------------------------------------------------------------------//
 /*!
  * Whether the track is outside the valid geometry region.
  */
-bool GeantGeoTrackView::is_outside() const
+CELER_FORCEINLINE bool GeantGeoTrackView::is_outside() const
 {
     return this->volume() == nullptr;
 }
@@ -277,7 +314,7 @@ bool GeantGeoTrackView::is_outside() const
 /*!
  * Whether the track is on the boundary of a volume.
  */
-bool GeantGeoTrackView::is_on_boundary() const
+CELER_FORCEINLINE bool GeantGeoTrackView::is_on_boundary() const
 {
     return safety_radius_ == 0.0;
 }
@@ -286,7 +323,7 @@ bool GeantGeoTrackView::is_on_boundary() const
 /*!
  * Find the distance to the next geometric boundary.
  */
-Propagation GeantGeoTrackView::find_next_step()
+CELER_FORCEINLINE Propagation GeantGeoTrackView::find_next_step()
 {
     return this->find_next_step(numeric_limits<real_type>::infinity());
 }
@@ -348,7 +385,7 @@ Propagation GeantGeoTrackView::find_next_step(real_type max_step)
 /*!
  * Find the safety at the current position.
  */
-auto GeantGeoTrackView::find_safety() -> real_type
+CELER_FORCEINLINE auto GeantGeoTrackView::find_safety() -> real_type
 {
     return this->find_safety(numeric_limits<real_type>::infinity());
 }
@@ -472,7 +509,7 @@ void GeantGeoTrackView::set_dir(Real3 const& newdir)
 /*!
  * Whether a next step has been calculated.
  */
-bool GeantGeoTrackView::has_next_step() const
+CELER_FORCEINLINE bool GeantGeoTrackView::has_next_step() const
 {
     return next_step_ != 0;
 }