feat: Allow particle selection with measurment count (acts-project#2570)

Enhances the particle selector to allow to select particles based on the number of measurements. The measurement related inputs are only required if the number of measurements is actually constrained.

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);
   }
 
   {