refactor: Refactor wiring for number of hits in simulation (#3001)

Currently we recalculate the number of hits per particles in different places. Here I try to unify this by outputting the number of hits from the simulation and attaching it to the final particle of the simulation.

Examples/Algorithms/Geant4/src/ParticleTrackingAction.cpp

@@ -119,12 +119,19 @@ ActsExamples::SimParticle ActsExamples::ParticleTrackingAction::convert(
   G4ThreeVector pDirection = aTrack.GetMomentumDirection();
   G4double p = convertEnergy * aTrack.GetKineticEnergy();
 
+  std::uint32_t numberOfHits = 0;
+  if (auto it = eventStore().particleHitCount.find(particleId);
+      it != eventStore().particleHitCount.end()) {
+    numberOfHits = it->second;
+  }
+
   // Now create the Particle
   ActsExamples::SimParticle aParticle(particleId, Acts::PdgParticle(pdg),
                                       charge, mass);
   aParticle.setPosition4(pPosition[0], pPosition[1], pPosition[2], pTime);
   aParticle.setDirection(pDirection[0], pDirection[1], pDirection[2]);
   aParticle.setAbsoluteMomentum(p);
+  aParticle.setNumberOfHits(numberOfHits);
   return aParticle;
 }
 

Examples/Io/Root/include/ActsExamples/Io/Root/RootParticleWriter.hpp

@@ -41,8 +41,6 @@ class RootParticleWriter final : public WriterT<SimParticleContainer> {
     /// Optional. If given, the the energy loss and traversed material is
     /// computed and written.
     std::string inputFinalParticles;
-    /// Optional. If given, the number of measurements is computed and written.
-    std::string inputSimHits;
     /// Path to the output file.
     std::string filePath;
     /// Output file access mode.
@@ -79,7 +77,6 @@ class RootParticleWriter final : public WriterT<SimParticleContainer> {
 
   ReadDataHandle<SimParticleContainer> m_inputFinalParticles{
       this, "InputFinalParticles"};
-  ReadDataHandle<SimHitContainer> m_inputSimHits{this, "InputSimHits"};
 
   std::mutex m_writeMutex;
 

Examples/Io/Root/src/RootParticleWriter.cpp

@@ -38,7 +38,6 @@ ActsExamples::RootParticleWriter::RootParticleWriter(
     throw std::invalid_argument("Missing tree name");
   }
 
-  m_inputSimHits.maybeInitialize(m_cfg.inputSimHits);
   m_inputFinalParticles.maybeInitialize(m_cfg.inputFinalParticles);
 
   // open root file and create the tree
@@ -80,8 +79,6 @@ ActsExamples::RootParticleWriter::RootParticleWriter(
     m_outputTree->Branch("e_loss", &m_eLoss);
     m_outputTree->Branch("total_x0", &m_pathInX0);
     m_outputTree->Branch("total_l0", &m_pathInL0);
-  }
-  if (m_inputSimHits.isInitialized()) {
     m_outputTree->Branch("number_of_hits", &m_numberOfHits);
   }
 }
@@ -105,18 +102,16 @@ ActsExamples::ProcessCode ActsExamples::RootParticleWriter::finalize() {
 
 ActsExamples::ProcessCode ActsExamples::RootParticleWriter::writeT(
     const AlgorithmContext& ctx, const SimParticleContainer& particles) {
+  const SimParticleContainer* finalParticles = nullptr;
+  if (m_inputFinalParticles.isInitialized()) {
+    finalParticles = &m_inputFinalParticles(ctx);
+  }
+
   // ensure exclusive access to tree/file while writing
   std::lock_guard<std::mutex> lock(m_writeMutex);
 
   auto nan = std::numeric_limits<float>::quiet_NaN();
 
-  std::unordered_map<ActsFatras::Barcode, std::uint32_t> hitsPerParticle;
-  if (m_inputSimHits.isInitialized()) {
-    for (const auto& simHit : m_inputSimHits(ctx)) {
-      ++hitsPerParticle[simHit.particleId()];
-    }
-  }
-
   m_eventId = ctx.eventNumber;
   for (const auto& particle : particles) {
     m_particleId.push_back(particle.particleId().value());
@@ -156,11 +151,11 @@ ActsExamples::ProcessCode ActsExamples::RootParticleWriter::writeT(
     m_generation.push_back(particle.particleId().generation());
     m_subParticle.push_back(particle.particleId().subParticle());
 
-    if (m_inputFinalParticles.isInitialized()) {
-      const auto& finalParticles = m_inputFinalParticles(ctx);
+    bool wroteFinalParticle = false;
+    if (finalParticles != nullptr) {
       // get the final particle
-      auto it = finalParticles.find(particle);
-      if (it == finalParticles.end()) {
+      auto it = finalParticles->find(particle);
+      if (it == finalParticles->end()) {
         ACTS_ERROR("Could not find final particle for "
                    << particle.particleId() << " in event " << ctx.eventNumber);
       } else {
@@ -175,25 +170,15 @@ ActsExamples::ProcessCode ActsExamples::RootParticleWriter::writeT(
         // get the path in L0
         m_pathInL0.push_back(Acts::clampValue<float>(finalParticle.pathInL0() /
                                                      Acts::UnitConstants::mm));
+        // get the number of hits
+        m_numberOfHits.push_back(finalParticle.numberOfHits());
+        wroteFinalParticle = true;
       }
-    } else {
+    }
+    if (!wroteFinalParticle) {
       m_eLoss.push_back(nan);
       m_pathInX0.push_back(nan);
       m_pathInL0.push_back(nan);
-    }
-
-    if (m_inputSimHits.isInitialized()) {
-      // get the sim hits
-      auto it = hitsPerParticle.find(particle.particleId());
-      if (it == hitsPerParticle.end()) {
-        ACTS_DEBUG("Could not find sim hits for "
-                   << particle.particleId() << " in event " << ctx.eventNumber);
-        m_numberOfHits.push_back(0);
-      } else {
-        // get the number of hits
-        m_numberOfHits.push_back(it->second);
-      }
-    } else {
       m_numberOfHits.push_back(-1);
     }
   }

Examples/Python/python/acts/examples/simulation.py

@@ -551,7 +551,6 @@ def addSimWriters(
                 level=customLogLevel(),
                 inputParticles=particlesInitial,
                 inputFinalParticles=particlesFinal,
-                inputSimHits=simHits,
                 filePath=str(outputDirRoot / "particles_simulation.root"),
             )
         )

Examples/Python/src/Output.cpp

@@ -175,8 +175,7 @@ void addOutput(Context& ctx) {
 
   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootParticleWriter, mex,
                              "RootParticleWriter", inputParticles,
-                             inputFinalParticles, inputSimHits, filePath,
-                             fileMode, treeName);
+                             inputFinalParticles, filePath, fileMode, treeName);
 
   ACTS_PYTHON_DECLARE_WRITER(ActsExamples::TrackFinderPerformanceWriter, mex,
                              "TrackFinderPerformanceWriter", inputProtoTracks,

Fatras/include/ActsFatras/EventData/Particle.hpp

@@ -284,6 +284,17 @@ class Particle {
         fourPosition(), direction(), qOverP(), std::nullopt, hypothesis());
   }
 
+  /// Set the number of hits.
+  ///
+  /// @param nHits number of hits
+  constexpr Particle &setNumberOfHits(std::uint32_t nHits) {
+    m_numberOfHits = nHits;
+    return *this;
+  }
+
+  /// Number of hits.
+  constexpr std::uint32_t numberOfHits() const { return m_numberOfHits; }
+
  private:
   // identity, i.e. things that do not change over the particle lifetime.
   /// Particle identifier within the event.
@@ -299,12 +310,14 @@ class Particle {
   Vector3 m_direction = Vector3::UnitZ();
   Scalar m_absMomentum = Scalar(0);
   Vector4 m_position4 = Vector4::Zero();
-  // proper time in the particle rest frame
+  /// proper time in the particle rest frame
   Scalar m_properTime = Scalar(0);
   // accumulated material
   Scalar m_pathInX0 = Scalar(0);
   Scalar m_pathInL0 = Scalar(0);
-  // reference surface
+  /// number of hits
+  std::uint32_t m_numberOfHits = 0;
+  /// reference surface
   const Acts::Surface *m_referenceSurface{nullptr};
 };
 

Fatras/include/ActsFatras/Kernel/detail/SimulationActor.hpp

@@ -198,7 +198,7 @@ struct SimulationActor {
         }
       }
     }
-    const Particle &after = result.particle;
+    Particle &after = result.particle;
 
     // store results of this interaction step, including potential hits
     if (selectHitSurface(surface)) {
@@ -207,6 +207,8 @@ struct SimulationActor {
           // the interaction could potentially modify the particle position
           Hit::Scalar(0.5) * (before.fourPosition() + after.fourPosition()),
           before.fourMomentum(), after.fourMomentum(), result.hits.size());
+
+      after.setNumberOfHits(result.hits.size());
     }
 
     if (after.absoluteMomentum() == 0.0) {