fix tracking manager in regions, transform g4 state in vecgeom state

examples/common/src/FTFP_BERT_AdePT.cc

@@ -38,8 +38,11 @@ FTFP_BERT_AdePT::FTFP_BERT_AdePT(G4int ver)
   // EM Physics
 
   // Register the EM physics to use for tracking on CPU
+  // Note: The explicit registering of physics on CPU is not needed anymore:
+  // the AdePTTrackingManager takes care of all EM particles and, on CPU, hands them over to
+  // the specialized G4HepEmTrackingManager
   // RegisterPhysics(new G4EmStandardPhysics());
-  RegisterPhysics(new HepEMPhysics(ver));
+  // RegisterPhysics(new HepEMPhysics(ver));
 
   // Register the AdePT physics
   RegisterPhysics(new AdePTPhysics());

include/AdePT/core/AdePTTransport.cuh

@@ -113,7 +113,6 @@ __global__ void InitTracks(adeptint::TrackData *trackinfo, int ntracks, int star
                            Secondaries secondaries, const vecgeom::VPlacedVolume *world, AdeptScoring *userScoring,
                            VolAuxData const *auxDataArray)
 {
-  constexpr double tolerance = 10. * vecgeom::kTolerance;
   for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < ntracks; i += blockDim.x * gridDim.x) {
     adept::TrackManager<Track> *trackmgr = nullptr;
     // These tracks come from Geant4, do not count them here
@@ -149,18 +148,11 @@ __global__ void InitTracks(adeptint::TrackData *trackinfo, int ntracks, int star
     track.localTime  = trackinfo->localTime;
     track.properTime = trackinfo->properTime;
 
+    // setting up the NavState
     track.navState.Clear();
-    // We locate the pushed point because we run the risk that the
-    // point is not located in the GPU region
-#ifndef ADEPT_USE_SURF
-    AdePTNavigator::LocatePointIn(world, track.pos + tolerance * track.dir, track.navState, true);
-#else
-    AdePTNavigator::LocatePointIn(vecgeom::NavigationState::WorldId(), track.pos + tolerance * track.dir,
-                                  track.navState, true);
-#endif
-    // The track must be on boundary at this point
-    track.navState.SetBoundaryState(true);
+    track.navState = trackinfo->navState;
     // nextState is initialized as needed.
+
 #ifndef ADEPT_USE_SURF
     int lvolID = track.navState.Top()->GetLogicalVolume()->id();
 #else
@@ -355,10 +347,15 @@ void ShowerGPU(IntegrationLayer &integration, int event, adeptint::TrackBuffer &
   COPCORE_CUDA_CHECK(cudaMemcpyAsync(gpuState.toDevice_dev, buffer.toDevice.data(),
                                      buffer.toDevice.size() * sizeof(adeptint::TrackData), cudaMemcpyHostToDevice,
                                      gpuState.stream));
-  // Initialize AdePT tracks using the track buffer copied from CPU
+
+#ifndef DEBUG_SINGLE_THREAD
   constexpr int initThreads = 32;
-  int initBlocks            = (buffer.toDevice.size() + initThreads - 1) / initThreads;
+#else
+  constexpr int initThreads = 1;
+#endif
+  int initBlocks = (buffer.toDevice.size() + initThreads - 1) / initThreads;
 
+  // Initialize AdePT tracks using the track buffer copied from CPU
   InitTracks<<<initBlocks, initThreads, 0, gpuState.stream>>>(gpuState.toDevice_dev, buffer.toDevice.size(),
                                                               buffer.startTrack, event, secondaries, world_dev,
                                                               scoring_dev, VolAuxArray::GetInstance().fAuxData_dev);

include/AdePT/core/AdePTTransport.h

@@ -44,7 +44,7 @@ class AdePTTransport : public AdePTTransportInterface {
   /// @brief Adds a track to the buffer
   void AddTrack(int pdg, int parentID, double energy, double x, double y, double z, double dirx, double diry,
                 double dirz, double globalTime, double localTime, double properTime, int threadId, unsigned int eventId,
-                unsigned int trackIndex);
+                unsigned int trackIndex, vecgeom::NavigationState &&state);
 
   void SetTrackCapacity(size_t capacity) { fCapacity = capacity; }
   /// @brief Get the track capacity on GPU
@@ -74,6 +74,11 @@ class AdePTTransport : public AdePTTransportInterface {
   void Cleanup();
   /// @brief Interface for transporting a buffer of tracks in AdePT.
   void Shower(int event, int threadId);
+  /// @brief Gets the VecGeomToG4Map
+  std::unordered_map<size_t, const G4VPhysicalVolume *> GetVecGeomG4Map() const override
+  {
+    return fIntegrationLayer.GetVecGeomG4Map();
+  }
 
 private:
   static inline G4HepEmState *fg4hepem_state{nullptr}; ///< The HepEm state singleton

include/AdePT/core/AdePTTransport.icc

@@ -41,17 +41,17 @@ void ShowerGPU(IntegrationLayer &integration, int event, TrackBuffer &buffer, GP
 template <typename IntegrationLayer>
 AdePTTransport<IntegrationLayer>::AdePTTransport(AdePTConfiguration &configuration)
 {
-  fDebugLevel = 0;
-  fBufferThreshold = configuration.GetTransportBufferThreshold();
-  fMaxBatch = 2 * configuration.GetTransportBufferThreshold();
+  fDebugLevel        = 0;
+  fBufferThreshold   = configuration.GetTransportBufferThreshold();
+  fMaxBatch          = 2 * configuration.GetTransportBufferThreshold();
   fTrackInAllRegions = configuration.GetTrackInAllRegions();
-  fGPURegionNames = configuration.GetGPURegionNames();
-  fCUDAStackLimit = configuration.GetCUDAStackLimit();
-  fCapacity = 1024 * 1024 * configuration.GetMillionsOfTrackSlots() / configuration.GetNumThreads();
+  fGPURegionNames    = configuration.GetGPURegionNames();
+  fCUDAStackLimit    = configuration.GetCUDAStackLimit();
+  fCapacity          = 1024 * 1024 * configuration.GetMillionsOfTrackSlots() / configuration.GetNumThreads();
   fHitBufferCapacity = 1024 * 1024 * configuration.GetMillionsOfHitSlots() / configuration.GetNumThreads();
-  
-  printf( "AdePT Allocated track capacity: %d tracks\n", fCapacity);
-  printf( "AdePT Allocated step buffer capacity: %d tracks\n", fHitBufferCapacity);
+
+  printf("AdePT Allocated track capacity: %d tracks\n", fCapacity);
+  printf("AdePT Allocated step buffer capacity: %d tracks\n", fHitBufferCapacity);
 }
 
 template <typename IntegrationLayer>
@@ -63,10 +63,12 @@ bool AdePTTransport<IntegrationLayer>::InitializeField(double bz)
 template <typename IntegrationLayer>
 void AdePTTransport<IntegrationLayer>::AddTrack(int pdg, int parent_id, double energy, double x, double y, double z,
                                                 double dirx, double diry, double dirz, double globalTime,
-                                                double localTime, double properTime, int /*threadId*/, unsigned int eventId,
-                                                unsigned int /*trackIndex*/)
+                                                double localTime, double properTime, int /*threadId*/,
+                                                unsigned int eventId, unsigned int /*trackIndex*/,
+                                                vecgeom::NavigationState &&state)
 {
-  fBuffer.toDevice.emplace_back(pdg, parent_id, energy, x, y, z, dirx, diry, dirz, globalTime, localTime, properTime);
+  fBuffer.toDevice.emplace_back(pdg, parent_id, energy, x, y, z, dirx, diry, dirz, globalTime, localTime, properTime,
+                                state);
   if (pdg == 11)
     fBuffer.nelectrons++;
   else if (pdg == -11)
@@ -85,9 +87,8 @@ template <typename IntegrationLayer>
 bool AdePTTransport<IntegrationLayer>::InitializeGeometry(const vecgeom::cxx::VPlacedVolume *world)
 {
   auto &cudaManager = vecgeom::cxx::CudaManager::Instance();
-  if(fCUDAStackLimit > 0)
-  {
-    std::cout <<  "CUDA Device stack limit: " <<  fCUDAStackLimit << "\n";
+  if (fCUDAStackLimit > 0) {
+    std::cout << "CUDA Device stack limit: " << fCUDAStackLimit << "\n";
     cudaDeviceSetLimit(cudaLimitStackSize, fCUDAStackLimit);
   }
   bool success = true;
@@ -219,8 +220,9 @@ template <typename IntegrationLayer>
 void AdePTTransport<IntegrationLayer>::Shower(int event, int /*threadId*/)
 {
   int tid = fIntegrationLayer.GetThreadID();
-  if (fDebugLevel > 0 && fBuffer.toDevice.size() == 0) {
-    std::cout << "[" << tid << "] AdePTTransport<IntegrationLayer>::Shower: No more particles in buffer. Exiting.\n";
+  if (fBuffer.toDevice.size() == 0) {
+    if (fDebugLevel > 0)
+      std::cout << "[" << tid << "] AdePTTransport<IntegrationLayer>::Shower: No more particles in buffer. Exiting.\n";
     return;
   }
 

include/AdePT/core/AdePTTransportInterface.hh

@@ -4,9 +4,13 @@
 #ifndef ADEPT_TRANSPORT_INTERFACE_H
 #define ADEPT_TRANSPORT_INTERFACE_H
 
+#include "G4VPhysicalVolume.hh"
+#include "VecGeom/navigation/NavigationState.h"
+
 #include <memory>
 #include <string>
 #include <vector>
+#include <unordered_map>
 
 class AdePTTransportInterface {
 public:
@@ -15,7 +19,7 @@ public:
   /// @brief Adds a track to the buffer
   virtual void AddTrack(int pdg, int parentID, double energy, double x, double y, double z, double dirx, double diry,
                         double dirz, double globalTime, double localTime, double properTime, int threadId,
-                        unsigned int eventId, unsigned int trackIndex) = 0;
+                        unsigned int eventId, unsigned int trackIndex, vecgeom::NavigationState &&state) = 0;
 
   /// @brief Set capacity of on-GPU track buffer.
   virtual void SetTrackCapacity(size_t capacity) = 0;
@@ -40,6 +44,8 @@ public:
   /// @brief Interface for transporting a buffer of tracks in AdePT.
   virtual void Shower(int event, int threadId) = 0;
   virtual void Cleanup()                       = 0;
+  /// @brief Get VecGeom to G4 volume map
+  virtual std::unordered_map<size_t, const G4VPhysicalVolume *> GetVecGeomG4Map() const = 0;
 };
 
 #endif

include/AdePT/core/TrackData.h

@@ -6,12 +6,15 @@
 
 #include <AdePT/base/MParray.h>
 
+#include "VecGeom/navigation/NavigationState.h"
+
 namespace adeptint {
 
 /// @brief Track data exchanged between Geant4 and AdePT
 /// @details This struct is initialised from an AdePT Track, either in GPU or CPU, copied to
 /// the destination, and used to reconstruct the track
 struct TrackData {
+  vecgeom::NavigationState navState;
   double position[3];
   double direction[3];
   double eKin{0};
@@ -23,12 +26,13 @@ struct TrackData {
 
   TrackData() = default;
   TrackData(int pdg_id, int parentID, double ene, double x, double y, double z, double dirx, double diry, double dirz,
-            double gTime, double lTime, double pTime)
-      : position{x, y, z}, direction{dirx, diry, dirz}, eKin{ene}, globalTime{gTime}, localTime{lTime},
+            double gTime, double lTime, double pTime, vecgeom::NavigationState state)
+      : navState{state}, position{x, y, z}, direction{dirx, diry, dirz}, eKin{ene}, globalTime{gTime}, localTime{lTime},
         properTime{pTime}, pdg{pdg_id}, parentID{parentID}
   {
   }
 
+  // fixme: add include navigation state in operators?
   friend bool operator==(TrackData const &a, TrackData const &b) { return !(a < b && b < a); }
   friend bool operator!=(TrackData const &a, TrackData const &b) { return !(a == b); }
   inline bool operator<(TrackData const &t) const

include/AdePT/integration/AdePTGeant4Integration.hh

@@ -30,7 +30,7 @@ class AdePTGeant4Integration {
 public:
   static constexpr G4int kAdePTTrackID =
       std::numeric_limits<G4int>::min() + 2; // TrackID to signify that the track came from AdePT
-  AdePTGeant4Integration()  = default;
+  AdePTGeant4Integration() = default;
   ~AdePTGeant4Integration();
 
   /// @brief Initializes VecGeom geometry
@@ -71,6 +71,8 @@ public:
 
   int GetThreadID() const { return G4Threading::G4GetThreadId(); }
 
+  std::unordered_map<size_t, const G4VPhysicalVolume *> GetVecGeomG4Map() const { return fglobal_vecgeom_to_g4_map; }
+
 private:
   /// @brief Reconstruct G4TouchableHistory from a VecGeom Navigation index
   void FillG4NavigationHistory(vecgeom::NavigationState aNavState, G4NavigationHistory *aG4NavigationHistory) const;

include/AdePT/integration/AdePTTrackingManager.hh

@@ -51,6 +51,11 @@ private:
   /// @brief Steps a track using the Generic G4TrackingManager until it enters a GPU region or stops
   void StepInHostRegion(G4Track *aTrack);
 
+  /// @brief Get the corresponding VecGeom NavigationState from the G4NavigationHistory
+  /// @param aG4NavigationHistory the given G4NavigationHistory
+  /// @return the corresponding vecgeom::NavigationState
+  const vecgeom::NavigationState GetVecGeomFromG4State(const G4Track *aG4Track);
+
   std::unique_ptr<G4HepEmTrackingManagerSpecialized> fHepEmTrackingManager;
   static inline int fNumThreads{0};
   std::set<G4Region const *> fGPURegions{};

include/AdePT/integration/G4HepEmTrackingManagerSpecialized.hh

@@ -20,6 +20,9 @@ public:
   ~G4HepEmTrackingManagerSpecialized();
 
   void SetGPURegions(const std::set<G4Region const *> &gpuRegions) { fGPURegions = gpuRegions; }
+  /// @brief Set whether AdePT should transport particles across the whole geometry
+  void SetTrackInAllRegions(bool trackInAllRegions) { fTrackInAllRegions = trackInAllRegions; }
+  bool GetTrackInAllRegions() const { return fTrackInAllRegions; }
 
   // Implement HandOverTrack that returns the track if it ends up in the GPU region
   void HandOverOneTrack(G4Track *aTrack) override;
@@ -29,7 +32,9 @@ public:
                               G4TrackVector &secondaries) const override;
 
 private:
-  std::set<G4Region const *> fGPURegions{};
+  bool fTrackInAllRegions = false;          ///< Whether the whole geometry is a GPU region
+  std::set<G4Region const *> fGPURegions{}; ///< List of GPU regions
+
   // G4Region const * fPreviousRegion = nullptr;
 };