Call G4 user tracking actions when offloading

src/accel/LocalTransporter.cc

@@ -10,12 +10,15 @@
 #include <memory>
 #include <mutex>
 #include <string>
+#include <unordered_map>
 #include <CLHEP/Units/SystemOfUnits.h>
+#include <G4DynamicParticle.hh>
 #include <G4EventManager.hh>
 #include <G4MTRunManager.hh>
 #include <G4ParticleDefinition.hh>
 #include <G4ThreeVector.hh>
 #include <G4Track.hh>
+#include <G4UserTrackingAction.hh>
 
 #include "corecel/Config.hh"
 
@@ -48,6 +51,7 @@
 #include "celeritas/optical/OpticalCollector.hh"
 #include "celeritas/phys/PDGNumber.hh"
 #include "celeritas/phys/ParticleParams.hh"  // IWYU pragma: keep
+#include "celeritas/user/DetectorSteps.hh"
 
 #include "SetupOptions.hh"
 #include "SharedParams.hh"
@@ -141,6 +145,31 @@ void trace(StepperResult const& track_counts)
             }                                                       \
         }                                                           \
     } while (0)
+
+//---------------------------------------------------------------------------//
+/*!
+ * Apply GPU terminal state from a death record to the given G4Track.
+ *
+ * Called in Flush() before firing PostUserTrackingAction. The track must
+ * already be configured for the correct particle type via view().
+ * Celeritas positions are in cm (native), energies in MeV, time in s
+ * (native) -- convert to Geant4 CLHEP units (mm, MeV, ns) via standard
+ * quantity helpers.
+ */
+void apply_death_state(G4Track& track, TrackDeathRecord const& d)
+{
+    // Position: Celeritas native (cm) -> Geant4 (mm)
+    track.SetPosition(native_to_geant<lengthunits::ClhepLength>(d.final_pos));
+    // Momentum direction: dimensionless unit vector
+    track.SetMomentumDirection(
+        to_g4vector(static_array_cast<double>(d.final_dir)));
+    // Kinetic energy: MeV value -- CLHEP::MeV == 1, so value() is correct
+    const_cast<G4DynamicParticle*>(track.GetDynamicParticle())
+        ->SetKineticEnergy(d.final_energy.value());
+    // Time: Celeritas native (s) -> Geant4 (ns)
+    track.SetGlobalTime(native_to_geant<units::ClhepTime>(d.final_time));
+}
+
 }  // namespace
 
 //---------------------------------------------------------------------------//
@@ -237,10 +266,6 @@ void LocalTransporter::InitializeEvent(int id)
             step_->reseed(event_id_);
         }
     }
-    if (hit_processor_)
-    {
-        hit_processor_->track_reconstruction().init_event();
-    }
 }
 
 //---------------------------------------------------------------------------//
@@ -251,6 +276,13 @@ void LocalTransporter::Push(G4Track& g4track)
 {
     CELER_EXPECT(*this);
 
+    if (flushing_tracking_actions_)
+    {
+        // Ignore re-offload attempts from PreUserTrackingAction callbacks
+        // fired during Flush() for reconstructed tracks
+        return;
+    }
+
     ScopedProfiling profile_this{"push"};
 
     GeantTrackView gtv{g4track};
@@ -272,21 +304,19 @@ void LocalTransporter::Push(G4Track& g4track)
 
     Primary track;
 
+    track.energy = gtv.energy();
+    track.particle_id = particles_->find(gtv.particle().pdg());
+
     // Generate Celeritas-specific PrimaryID and capture user info
     if (hit_processor_)
     {
-        track.primary_id
-            = hit_processor_->track_reconstruction().acquire(g4track);
+        track.primary_id = hit_processor_->track_reconstruction().acquire(
+            g4track, track.particle_id);
     }
-
-    track.energy = gtv.energy();
-    track.particle_id = particles_->find(gtv.particle().pdg());
     track.position = static_array_cast<real_type>(native_value_from(gtv.pos()));
     track.direction = static_array_cast<real_type>(gtv.dir());
     track.time = static_cast<real_type>(native_value_from(gtv.time()));
     track.weight = gtv.weight();
-    track.primary_id = celeritas::id_cast<PrimaryId>(
-        track.primary_id.unchecked_get() + g4track.GetTrackID());
 
     CELER_VALIDATE(track.particle_id,
                    << "cannot offload '" << gtv.particle().name()
@@ -319,21 +349,17 @@ void LocalTransporter::Flush()
 
     ScopedProfiling profile_this("flush");
 
-    if (event_manager_ || !event_id_)
+    if (!event_manager_)
     {
-        if (CELER_UNLIKELY(!event_manager_))
-        {
-            // Save the event manager pointer, thereby marking that
-            // *subsequent* events need to have their IDs checked as well
-            event_manager_ = G4EventManager::GetEventManager();
-            CELER_ASSERT(event_manager_);
-        }
+        event_manager_ = G4EventManager::GetEventManager();
+        CELER_ASSERT(event_manager_);
+    }
 
+    {
         G4Event const* event = event_manager_->GetConstCurrentEvent();
         CELER_ASSERT(event);
         if (event_id_ != id_cast<UniqueEventId>(event->GetEventID()))
         {
-            // The event ID has changed: reseed it
             this->InitializeEvent(event->GetEventID());
         }
     }
@@ -415,6 +441,43 @@ void LocalTransporter::Flush()
                                    << " hits for event " << event_id_.get();
             run_accum_.hits += num_hits;
         }
+
+        // Fire Pre/PostUserTrackingAction for every offloaded track so
+        // MC-truth frameworks (e.g. DD4hep Geant4ParticleHandler) can
+        // register equivalence entries. Pre receives the original handover
+        // state; Post receives the GPU terminal state if available.
+        if (auto* ta = event_manager_->GetUserTrackingAction())
+        {
+            flushing_tracking_actions_ = true;
+            auto& recon = hit_processor_->track_reconstruction();
+
+            auto const& deaths = hit_processor_->last_deaths();
+            std::unordered_map<size_type, TrackDeathRecord const*> death_map;
+            for (auto const& d : deaths)
+            {
+                if (d.primary_id)
+                {
+                    death_map[d.primary_id.unchecked_get()] = &d;
+                }
+            }
+
+            for (size_type i = 0; i < recon.num_primaries(); ++i)
+            {
+                auto pid = id_cast<PrimaryId>(i);
+                G4Track& g4track
+                    = recon.view_initial(recon.particle_id(pid), pid);
+                ta->PreUserTrackingAction(&g4track);
+
+                if (auto it = death_map.find(i); it != death_map.end())
+                {
+                    apply_death_state(g4track, *it->second);
+                }
+
+                ta->PostUserTrackingAction(&g4track);
+            }
+            flushing_tracking_actions_ = false;
+        }
+        hit_processor_->clear_deaths();
         hit_processor_->track_reconstruction().clear();
     }
 }

src/accel/LocalTransporter.hh

@@ -142,6 +142,10 @@ class LocalTransporter final : public TrackOffloadInterface
 
     // Shared across threads to write flushed particles
     SPOffloadWriter dump_primaries_;
+
+    // True while firing Pre/PostUserTrackingAction in Flush to prevent
+    // re-offloading of reconstructed tracks
+    bool flushing_tracking_actions_{false};
 };
 
 //---------------------------------------------------------------------------//

src/accel/TrackingManager.cc

@@ -122,7 +122,16 @@ void TrackingManager::PreparePhysicsTable(G4ParticleDefinition const& part)
 
 //---------------------------------------------------------------------------//
 /*!
- * Offload the incoming track to Celeritas.
+ * Offload the track to Celeritas for transport.
+ *
+ * A \c G4VTrackingManager is responsible for firing user action callbacks,
+ * just as \c G4TrackingManager::ProcessOneTrack does for standard tracks.
+ * For offloaded tracks, both \c PreUserTrackingAction and
+ * \c PostUserTrackingAction are fired together in
+ * \c LocalTransporter::Flush() after GPU transport completes, using the
+ * original handover state for Pre and the GPU terminal state for Post.
+ * This ensures MC-truth frameworks that read \c m_currTrack in Post see
+ * the state set by Pre, preserving the expected Pre/Post pairing semantics.
  *
  * This will \em not be called in the "master" thread of an MT run.
  */

src/accel/TrackingManager.hh

@@ -11,6 +11,8 @@
 #    error "Tracking manager offload requires Geant4 11.0 or higher"
 #endif
 
+#include <G4ParticleDefinition.hh>
+#include <G4Track.hh>
 #include <G4VTrackingManager.hh>
 
 namespace celeritas
@@ -20,6 +22,13 @@ namespace celeritas
 class SharedParams;
 class TrackOffloadInterface;
 
+//---------------------------------------------------------------------------//
+//! Whether a track's particle type is offloaded to Celeritas.
+inline bool IsTrackOffloadedToCeleritas(G4Track const* track)
+{
+    return track->GetParticleDefinition()->GetTrackingManager() != nullptr;
+}
+
 //---------------------------------------------------------------------------//
 /*!
  * Offload to Celeritas via the per-particle Geant4 "tracking manager".

src/celeritas/ext/GeantSd.cc

@@ -75,6 +75,7 @@ GeantSd::GeantSd(ParticleParams const& par,
 
     // Convert setup options to step data
     selection_.particle = setup.track;
+    selection_.primary_id = setup.track;
     selection_.weight = setup.track;
     selection_.energy_deposition = setup.energy_deposition;
     selection_.step_length = setup.step_length;
@@ -143,6 +144,7 @@ auto GeantSd::filters() const -> Filters
     }
 
     result.nonzero_energy_deposition = nonzero_energy_deposition_;
+    result.track_death = true;
 
     return result;
 }