Refactor Geant4 utilities in preparation for g4vg (celeritas-project#…

…1095)

* Reuse daughter placer in g4vg
* Refactor Geant4 volume visitor to an actual volume visitor
* Fix cxx target
* Split generic geo test into multiple
* Fix build errors
* Fix edge case in importer

src/celeritas/CMakeLists.txt

@@ -135,7 +135,6 @@ if(CELERITAS_USE_Geant4)
     ext/detail/GeantModelImporter.cc
     ext/detail/GeantPhysicsList.cc
     ext/detail/GeantProcessImporter.cc
-    ext/detail/GeantVolumeVisitor.cc
   )
   set(_cg4_libs Celeritas::corecel XercesC::XercesC ${Geant4_LIBRARIES})
 

src/celeritas/ext/GeantGeoParams.cc

@@ -33,10 +33,40 @@
 #include "GeantUtils.hh"
 #include "ScopedGeantExceptionHandler.hh"
 #include "ScopedGeantLogger.hh"
-#include "detail/GeantVolumeVisitor.hh"
+#include "VisitGeantVolumes.hh"
 
 namespace celeritas
 {
+namespace
+{
+//---------------------------------------------------------------------------//
+std::vector<Label>
+get_volume_labels(G4LogicalVolume const& world, bool unique_volumes)
+{
+    std::vector<Label> labels;
+    visit_geant_volumes(
+        [&](G4LogicalVolume const& lv) {
+            auto i = static_cast<std::size_t>(lv.GetInstanceID());
+            if (i >= labels.size())
+            {
+                labels.resize(i + 1);
+            }
+            if (unique_volumes)
+            {
+                labels[i] = Label::from_geant(make_gdml_name(lv));
+            }
+            else
+            {
+                labels[i] = Label::from_geant(lv.GetName());
+            }
+        },
+        world);
+    return labels;
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace
+
 //---------------------------------------------------------------------------//
 /*!
  * Construct from a GDML input.
@@ -243,25 +273,18 @@ void GeantGeoParams::build_metadata()
     CELER_EXPECT(host_ref_);
 
     ScopedMem record_mem("GeantGeoParams.build_metadata");
+
+    auto const* world_lv = host_ref_.world->GetLogicalVolume();
+    CELER_ASSERT(world_lv);
+
     // Construct volume labels
     vol_labels_ = LabelIdMultiMap<VolumeId>(
-        [world = host_ref_.world, unique_volumes = !loaded_gdml_] {
-            G4LogicalVolume const* lv = world->GetLogicalVolume();
-            CELER_ASSERT(lv);
-
-            // Recursive loop over all logical volumes to populate map
-            detail::GeantVolumeVisitor visitor(unique_volumes);
-            visitor.visit(*lv);
-            return visitor.build_label_vector();
-        }());
+        get_volume_labels(*world_lv, !loaded_gdml_));
 
     // Save world bbox (NOTE: assumes no transformation on PV?)
-    bbox_ = [world = host_ref_.world] {
-        G4LogicalVolume const* lv = world->GetLogicalVolume();
-        CELER_ASSERT(lv);
-        G4VSolid const* solid = lv->GetSolid();
+    bbox_ = [world_lv] {
+        G4VSolid const* solid = world_lv->GetSolid();
         CELER_ASSERT(solid);
-
         G4VisExtent const& extent = solid->GetExtent();
 
         return BBox({convert_from_geant(G4ThreeVector(extent.GetXmin(),

src/celeritas/ext/GeantGeoUtils.cc

@@ -12,6 +12,7 @@
 #include <string_view>
 #include <unordered_set>
 #include <G4GDMLParser.hh>
+#include <G4GDMLWriteStructure.hh>
 #include <G4LogicalVolume.hh>
 #include <G4LogicalVolumeStore.hh>
 #include <G4PhysicalVolumeStore.hh>
@@ -255,5 +256,31 @@ find_geant_volumes(std::unordered_set<std::string> names)
     return result;
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Generate the GDML name for a Geant4 logical volume.
+ */
+std::string make_gdml_name(G4LogicalVolume const& lv)
+{
+    // Run the LV through the GDML export name generator so that the volume is
+    // uniquely identifiable in VecGeom. Reuse the same instance to reduce
+    // overhead: note that the method isn't const correct.
+    static G4GDMLWriteStructure temp_writer;
+
+    auto const* refl_factory = G4ReflectionFactory::Instance();
+    if (auto const* unrefl_lv
+        = refl_factory->GetConstituentLV(const_cast<G4LogicalVolume*>(&lv)))
+    {
+        // If this is a reflected volume, add the reflection extension after
+        // the final pointer to match the converted VecGeom name
+        std::string name
+            = temp_writer.GenerateName(unrefl_lv->GetName(), unrefl_lv);
+        name += refl_factory->GetVolumesNameExtension();
+        return name;
+    }
+
+    return temp_writer.GenerateName(lv.GetName(), &lv);
+}
+
 //---------------------------------------------------------------------------//
 }  // namespace celeritas

src/celeritas/ext/GeantGeoUtils.hh

@@ -76,10 +76,15 @@ void reset_geant_geometry();
 // Get a view to the Geant4 LV store
 Span<G4LogicalVolume*> geant_logical_volumes();
 
+//---------------------------------------------------------------------------//
 // Find Geant4 logical volumes corresponding to a list of names
 std::unordered_set<G4LogicalVolume const*>
     find_geant_volumes(std::unordered_set<std::string>);
 
+//---------------------------------------------------------------------------//
+// Generate the GDML name for a Geant4 logical volume
+std::string make_gdml_name(G4LogicalVolume const&);
+
 //---------------------------------------------------------------------------//
 // INLINE DEFINITIONS
 //---------------------------------------------------------------------------//
@@ -103,6 +108,17 @@ inline Span<G4LogicalVolume*> geant_logical_volumes()
 {
     CELER_NOT_CONFIGURED("Geant4");
 }
+
+inline std::unordered_set<G4LogicalVolume const*>
+find_geant_volumes(std::unordered_set<std::string>)
+{
+    CELER_NOT_CONFIGURED("Geant4");
+}
+
+inline std::string make_gdml_name(G4LogicalVolume const&)
+{
+    CELER_NOT_CONFIGURED("Geant4");
+}
 #endif
 
 //---------------------------------------------------------------------------//

src/celeritas/ext/GeantImporter.cc

@@ -23,6 +23,7 @@
 #include <G4ElementVector.hh>
 #include <G4EmParameters.hh>
 #include <G4GammaGeneralProcess.hh>
+#include <G4LogicalVolumeStore.hh>
 #include <G4Material.hh>
 #include <G4MaterialCutsCouple.hh>
 #include <G4Navigator.hh>
@@ -66,10 +67,11 @@
 #include "celeritas/io/SeltzerBergerReader.hh"
 #include "celeritas/phys/PDGNumber.hh"
 
+#include "GeantGeoUtils.hh"
 #include "ScopedGeantExceptionHandler.hh"
+#include "VisitGeantVolumes.hh"
 #include "detail/AllElementReader.hh"
 #include "detail/GeantProcessImporter.hh"
-#include "detail/GeantVolumeVisitor.hh"
 
 inline constexpr double mev_scale = 1 / CLHEP::MeV;
 
@@ -814,15 +816,49 @@ ImportData GeantImporter::operator()(DataSelection const& selected)
 std::vector<ImportVolume>
 GeantImporter::import_volumes(bool unique_volumes) const
 {
-    detail::GeantVolumeVisitor visitor(unique_volumes);
-    // Recursive loop over all logical volumes to populate map
-    visitor.visit(*world_->GetLogicalVolume());
+    // Note: if the LV has been purged (i.e. by trying to run multiple
+    // geometries in the same execution), the instance ID's won't correspond to
+    // the location in the vector.
+    G4LogicalVolumeStore* lv_store = G4LogicalVolumeStore::GetInstance();
+    CELER_ASSERT(lv_store);
+    std::vector<ImportVolume> result;
+    result.reserve(lv_store->size());
+
+    // Recursive loop over all logical volumes to populate volumes
+    visit_geant_volumes(
+        [unique_volumes, &result](G4LogicalVolume const& lv) {
+            auto i = static_cast<std::size_t>(lv.GetInstanceID());
+            if (i >= result.size())
+            {
+                result.resize(i + 1);
+            }
+
+            ImportVolume& volume = result[lv.GetInstanceID()];
+            if (auto* cuts = lv.GetMaterialCutsCouple())
+            {
+                volume.material_id = cuts->GetIndex();
+            }
+            volume.name = lv.GetName();
+            volume.solid_name = lv.GetSolid()->GetName();
 
-    auto volumes = visitor.build_volume_vector();
-    CELER_LOG(debug) << "Loaded " << volumes.size() << " volumes with "
+            if (volume.name.empty())
+            {
+                CELER_LOG(warning)
+                    << "No logical volume name specified for instance ID " << i
+                    << " (material " << volume.material_id << ")";
+            }
+            else if (unique_volumes)
+            {
+                // Add pointer as GDML writer does
+                volume.name = make_gdml_name(lv);
+            }
+        },
+        *world_->GetLogicalVolume());
+
+    CELER_LOG(debug) << "Loaded " << result.size() << " volumes with "
                      << (unique_volumes ? "uniquified" : "original")
                      << " names";
-    return volumes;
+    return result;
 }
 
 //---------------------------------------------------------------------------//

src/celeritas/ext/GeantVolumeMapper.cc

@@ -13,7 +13,7 @@
 #include "corecel/io/Join.hh"
 #include "corecel/io/Logger.hh"
 
-#include "detail/GeantVolumeVisitor.hh"
+#include "GeantGeoUtils.hh"
 
 namespace celeritas
 {
@@ -38,8 +38,7 @@ VolumeId GeantVolumeMapper::operator()(G4LogicalVolume const& lv) const
     {
         // Label doesn't have a pointer address attached: we probably need
         // to regenerate to match the exported GDML file
-        label
-            = Label::from_geant(detail::GeantVolumeVisitor::generate_name(lv));
+        label = Label::from_geant(make_gdml_name(lv));
     }
 
     if (auto id = geo.find_volume(label))
@@ -63,7 +62,7 @@ VolumeId GeantVolumeMapper::operator()(G4LogicalVolume const& lv) const
     // Try regenerating the name even if we *did* have a pointer
     // address attached (in case an original GDML volume name already
     // had a pointer suffix and LoadGdml added another)
-    label = Label::from_geant(detail::GeantVolumeVisitor::generate_name(lv));
+    label = Label::from_geant(make_gdml_name(lv));
     all_ids = geo.find_volumes(label.name);
     if (all_ids.size() > 1)
     {

src/celeritas/ext/VisitGeantVolumes.hh

@@ -0,0 +1,52 @@
+//----------------------------------*-C++-*----------------------------------//
+// Copyright 2023-2024 UT-Battelle, LLC, and other Celeritas developers.
+// See the top-level COPYRIGHT file for details.
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+//---------------------------------------------------------------------------//
+//! \file celeritas/ext/VisitGeantVolumes.hh
+//---------------------------------------------------------------------------//
+#pragma once
+
+#include <unordered_set>
+#include <vector>
+#include <G4LogicalVolume.hh>
+
+#include "corecel/Assert.hh"
+#include "corecel/cont/Range.hh"
+
+namespace celeritas
+{
+//---------------------------------------------------------------------------//
+/*!
+ * Do a recursive depth-first listing of Geant4 logical volumes.
+ *
+ * This will visit each volume exactly once based on when it's encountered in
+ * the hierarchy. The visitor function F should have the signature
+ * \code void(*)(G4LogicalVolume const&) \endcode .
+ */
+template<class F>
+void visit_geant_volumes(F&& vis, G4LogicalVolume const& parent_vol)
+{
+    std::unordered_set<G4LogicalVolume const*> visited;
+    std::vector<G4LogicalVolume const*> stack{&parent_vol};
+
+    while (!stack.empty())
+    {
+        G4LogicalVolume const* lv = stack.back();
+        stack.pop_back();
+        vis(*lv);
+        for (auto const i : range(lv->GetNoDaughters()))
+        {
+            G4LogicalVolume* daughter = lv->GetDaughter(i)->GetLogicalVolume();
+            CELER_ASSERT(daughter);
+            auto&& [iter, inserted] = visited.insert(daughter);
+            if (inserted)
+            {
+                stack.push_back(daughter);
+            }
+        }
+    }
+}
+
+//---------------------------------------------------------------------------//
+}  // namespace celeritas