feat: Allow particle selection with measurment count

Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/ParticleSelector.cpp

@@ -10,6 +10,8 @@
 
 #include "Acts/Definitions/Common.hpp"
 #include "Acts/Utilities/VectorHelpers.hpp"
+#include "ActsExamples/EventData/GeometryContainers.hpp"
+#include "ActsExamples/EventData/IndexSourceLink.hpp"
 #include "ActsExamples/EventData/SimParticle.hpp"
 #include "ActsExamples/Framework/AlgorithmContext.hpp"
 #include "ActsFatras/EventData/Particle.hpp"
@@ -50,14 +52,38 @@ ActsExamples::ParticleSelector::ParticleSelector(const Config& config,
   ACTS_DEBUG("remove charged particles " << m_cfg.removeCharged);
   ACTS_DEBUG("remove neutral particles " << m_cfg.removeNeutral);
   ACTS_DEBUG("remove secondary particles " << m_cfg.removeSecondaries);
+
+  // We only initialize this if we actually select on this
+  if (m_cfg.measurementsMin > 0 or
+      m_cfg.measurementsMax < std::numeric_limits<std::size_t>::max()) {
+    m_inputMap.initialize(m_cfg.inputMeasurementParticlesMap);
+    ACTS_DEBUG("selection particle number of measurements ["
+               << m_cfg.measurementsMin << "," << m_cfg.measurementsMax << ")");
+  }
 }
 
 ActsExamples::ProcessCode ActsExamples::ParticleSelector::execute(
     const AlgorithmContext& ctx) const {
+  using ParticlesMeasurmentMap =
+      boost::container::flat_multimap<ActsFatras::Barcode, Index>;
+
+  // prepare input/ output types
+  const auto& inputParticles = m_inputParticles(ctx);
+
+  // Make global particles measurement map if necessary
+  std::optional<ParticlesMeasurmentMap> particlesMeasMap;
+  if (m_inputMap.isInitialized()) {
+    particlesMeasMap = invertIndexMultimap(m_inputMap(ctx));
+  }
+
+  std::size_t nInvalidCharge = 0;
+  std::size_t nInvalidMeasurementCount = 0;
+
   // helper functions to select tracks
-  auto within = [](double x, double min, double max) {
+  auto within = [](auto x, auto min, auto max) {
     return (min <= x) and (x < max);
   };
+
   auto isValidParticle = [&](const ActsFatras::Particle& p) {
     const auto eta = Acts::VectorHelpers::eta(p.direction());
     const auto phi = Acts::VectorHelpers::phi(p.direction());
@@ -67,7 +93,26 @@ ActsExamples::ProcessCode ActsExamples::ParticleSelector::execute(
     const bool validCharged = (p.charge() != 0) and not m_cfg.removeCharged;
     const bool validCharge = validNeutral or validCharged;
     const bool validSecondary = not m_cfg.removeSecondaries or !p.isSecondary();
-    return validCharge and validSecondary and
+
+    nInvalidCharge += static_cast<std::size_t>(not validCharge);
+
+    // default valid measurement count to true and only change if we have loaded
+    // the measurement particles map
+    bool validMeasurementCount = true;
+    if (particlesMeasMap) {
+      auto [b, e] = particlesMeasMap->equal_range(p.particleId());
+      validMeasurementCount =
+          within(static_cast<std::size_t>(std::distance(b, e)),
+                 m_cfg.measurementsMin, m_cfg.measurementsMax);
+
+      ACTS_VERBOSE("Found " << std::distance(b, e) << " measurements for "
+                            << p.particleId());
+    }
+
+    nInvalidMeasurementCount +=
+        static_cast<std::size_t>(not validMeasurementCount);
+
+    return validCharge and validSecondary and validMeasurementCount and
            within(p.transverseMomentum(), m_cfg.ptMin, m_cfg.ptMax) and
            within(std::abs(eta), m_cfg.absEtaMin, m_cfg.absEtaMax) and
            within(eta, m_cfg.etaMin, m_cfg.etaMax) and
@@ -79,9 +124,6 @@ ActsExamples::ProcessCode ActsExamples::ParticleSelector::execute(
            within(p.mass(), m_cfg.mMin, m_cfg.mMax);
   };
 
-  // prepare input/ output types
-  const auto& inputParticles = m_inputParticles(ctx);
-
   SimParticleContainer outputParticles;
   outputParticles.reserve(inputParticles.size());
 
@@ -97,6 +139,9 @@ ActsExamples::ProcessCode ActsExamples::ParticleSelector::execute(
   ACTS_DEBUG("event " << ctx.eventNumber << " selected "
                       << outputParticles.size() << " from "
                       << inputParticles.size() << " particles");
+  ACTS_DEBUG("filtered out because of charge: " << nInvalidCharge);
+  ACTS_DEBUG("filtered out because of measurement count: "
+             << nInvalidMeasurementCount);
 
   m_outputParticles(ctx, std::move(outputParticles));
   return ProcessCode::SUCCESS;

Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/ParticleSelector.hpp

@@ -12,7 +12,10 @@
 
 #pragma once
 
+#include "Acts/Geometry/GeometryIdentifier.hpp"
 #include "Acts/Utilities/Logger.hpp"
+#include "ActsExamples/EventData/Index.hpp"
+#include "ActsExamples/EventData/Measurement.hpp"
 #include "ActsExamples/EventData/SimParticle.hpp"
 #include "ActsExamples/Framework/DataHandle.hpp"
 #include "ActsExamples/Framework/IAlgorithm.hpp"
@@ -30,6 +33,8 @@ class ParticleSelector final : public IAlgorithm {
   struct Config {
     /// The input particles collection.
     std::string inputParticles;
+    /// Input measurement particles map (Optional)
+    std::string inputMeasurementParticlesMap;
     /// The output particles collection.
     std::string outputParticles;
     // Minimum/maximum distance from the origin in the transverse plane.
@@ -54,6 +59,9 @@ class ParticleSelector final : public IAlgorithm {
     // Rest mass cuts
     double mMin = 0;
     double mMax = std::numeric_limits<double>::infinity();
+    /// Measurement number cuts
+    std::size_t measurementsMin = 0;
+    std::size_t measurementsMax = std::numeric_limits<std::size_t>::max();
     /// Remove charged particles.
     bool removeCharged = false;
     /// Remove neutral particles.
@@ -74,6 +82,8 @@ class ParticleSelector final : public IAlgorithm {
   Config m_cfg;
 
   ReadDataHandle<SimParticleContainer> m_inputParticles{this, "InputParticles"};
+  ReadDataHandle<IndexMultimap<ActsFatras::Barcode>> m_inputMap{
+      this, "InputMeasurementParticlesMap"};
 
   WriteDataHandle<SimParticleContainer> m_outputParticles{this,
                                                           "OutputParticles"};

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

@@ -40,11 +40,12 @@
         "absEta",  # (min,max)
         "pt",  # (min,max)
         "m",  # (min,max)
+        "measurements",  # (min,max)
         "removeCharged",  # bool
         "removeNeutral",  # bool
         "removeSecondaries",  # bool
     ],
-    defaults=[(None, None)] * 8 + [None] * 3,
+    defaults=[(None, None)] * 9 + [None] * 3,
 )
 
 
@@ -333,6 +334,7 @@ def addParticleSelection(
     config: ParticleSelectorConfig,
     inputParticles: str,
     outputParticles: str,
+    inputMeasurementParticlesMap: str = "",
     logLevel: Optional[acts.logging.Level] = None,
 ) -> None:
     """
@@ -370,13 +372,16 @@ def addParticleSelection(
                 ptMax=config.pt[1],
                 mMin=config.m[0],
                 mMax=config.m[1],
+                measurementsMin=config.measurements[0],
+                measurementsMax=config.measurements[1],
                 removeCharged=config.removeCharged,
                 removeNeutral=config.removeNeutral,
                 removeSecondaries=config.removeSecondaries,
             ),
             level=customLogLevel(),
             inputParticles=inputParticles,
             outputParticles=outputParticles,
+            inputMeasurementParticlesMap=inputMeasurementParticlesMap,
         )
     )
 

Examples/Python/src/TruthTracking.cpp

@@ -106,6 +106,7 @@ void addTruthTracking(Context& ctx) {
 
     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
     ACTS_PYTHON_MEMBER(inputParticles);
+    ACTS_PYTHON_MEMBER(inputMeasurementParticlesMap);
     ACTS_PYTHON_MEMBER(outputParticles);
     ACTS_PYTHON_MEMBER(rhoMin);
     ACTS_PYTHON_MEMBER(rhoMax);
@@ -123,6 +124,8 @@ void addTruthTracking(Context& ctx) {
     ACTS_PYTHON_MEMBER(mMax);
     ACTS_PYTHON_MEMBER(ptMin);
     ACTS_PYTHON_MEMBER(ptMax);
+    ACTS_PYTHON_MEMBER(measurementsMin);
+    ACTS_PYTHON_MEMBER(measurementsMax);
     ACTS_PYTHON_MEMBER(removeCharged);
     ACTS_PYTHON_MEMBER(removeNeutral);
     ACTS_PYTHON_MEMBER(removeSecondaries);
@@ -136,6 +139,8 @@ void addTruthTracking(Context& ctx) {
     pythonRangeProperty(c, "absEta", &Config::absEtaMin, &Config::absEtaMax);
     pythonRangeProperty(c, "m", &Config::mMin, &Config::mMax);
     pythonRangeProperty(c, "pt", &Config::ptMin, &Config::ptMax);
+    pythonRangeProperty(c, "measurements", &Config::measurementsMin,
+                        &Config::measurementsMax);
   }
 
   {