feat: Introduce TrackToTruthJetAlgorithm

Examples/Algorithms/Jets/CMakeLists.txt

@@ -1,11 +1,16 @@
-acts_add_library(ExamplesJets SHARED src/TruthJetAlgorithm.cpp)
+acts_add_library(
+    ExamplesJets
+    SHARED
+    src/TruthJetAlgorithm.cpp
+    src/TrackToTruthJetAlgorithm.cpp
+)
 target_include_directories(
     ActsExamplesJets
     PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
 )
 target_link_libraries(
     ActsExamplesJets
-    PUBLIC Acts::PluginFastJet Acts::ExamplesFramework
+    PUBLIC Acts::Core Acts::PluginFastJet Acts::ExamplesFramework
 )
 
 acts_compile_headers(ExamplesJets GLOB "include/**/*.hpp")

Examples/Algorithms/Jets/include/ActsExamples/Jets/TrackToTruthJetAlgorithm.hpp

@@ -0,0 +1,57 @@
+// This file is part of the ACTS project.
+//
+// Copyright (C) 2016 CERN for the benefit of the ACTS project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#pragma once
+
+#include "Acts/Definitions/Algebra.hpp"
+#include "Acts/Plugins/FastJet/Jets.hpp"
+#include "Acts/Utilities/Logger.hpp"
+#include "ActsExamples/EventData/SimParticle.hpp"
+#include "ActsExamples/EventData/Track.hpp"
+#include "ActsExamples/Framework/DataHandle.hpp"
+#include "ActsExamples/Framework/IAlgorithm.hpp"
+#include "ActsExamples/Framework/ProcessCode.hpp"
+
+#include <string>
+
+#include <fastjet/ClusterSequence.hh>
+
+namespace ActsExamples {
+struct AlgorithmContext;
+
+/// Print all particles.
+class TrackToTruthJetAlgorithm : public IAlgorithm {
+ public:
+  struct Config {
+    /// Input tracks collection.
+    std::string inputTracks;
+    /// Input jets collection.
+    std::string inputJets;
+    /// Output track jets collection.
+    std::string outputTrackJets;
+    /// Maximum delta R for track to jet matching.
+    double maxDeltaR = 0.4;
+  };
+
+  TrackToTruthJetAlgorithm(const Config& cfg, Acts::Logging::Level lvl);
+
+  ProcessCode execute(const AlgorithmContext& ctx) const override;
+  ProcessCode finalize() override;
+
+  const Config& config() const { return m_cfg; }
+
+ private:
+  Config m_cfg;
+  ReadDataHandle<ConstTrackContainer> m_inputTracks{this, "inputTracks"};
+  ReadDataHandle<Acts::FastJet::TrackJetContainer> m_inputJets{this,
+                                                               "inputJets"};
+  WriteDataHandle<Acts::FastJet::TrackJetContainer> m_outputTrackJets{
+      this, "outputTrackJets"};
+};
+
+}  // namespace ActsExamples

Examples/Algorithms/Jets/include/ActsExamples/Jets/TruthJetAlgorithm.hpp

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "Acts/Plugins/FastJet/Jets.hpp"
 #include "Acts/Utilities/Logger.hpp"
 #include "ActsExamples/EventData/SimParticle.hpp"
 #include "ActsExamples/Framework/DataHandle.hpp"
@@ -20,6 +21,11 @@ namespace fastjet {
 class PseudoJet;
 }
 
+namespace ActsFastJet {
+class TruthJetBuilder;
+class JetProperties;
+}  // namespace ActsFastJet
+
 namespace ActsExamples {
 struct AlgorithmContext;
 
@@ -45,8 +51,8 @@ class TruthJetAlgorithm final : public IAlgorithm {
   Config m_cfg;
   ReadDataHandle<SimParticleContainer> m_inputTruthParticles{
       this, "inputTruthParticles"};
-  WriteDataHandle<std::vector<fastjet::PseudoJet>> m_outputJets{this,
-                                                                "outputJets"};
+  WriteDataHandle<Acts::FastJet::TrackJetContainer> m_outputJets{this,
+                                                                 "outputJets"};
 };
 
 }  // namespace ActsExamples

Examples/Algorithms/Jets/src/TrackToTruthJetAlgorithm.cpp

@@ -0,0 +1,124 @@
+// This file is part of the ACTS project.
+//
+// Copyright (C) 2016 CERN for the benefit of the ACTS project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+#include "ActsExamples/Jets/TrackToTruthJetAlgorithm.hpp"
+
+#include "Acts/Utilities/Logger.hpp"
+#include "ActsExamples/Framework/AlgorithmContext.hpp"
+
+#include <ostream>
+#include <stdexcept>
+
+namespace ActsExamples {
+
+TrackToTruthJetAlgorithm::TrackToTruthJetAlgorithm(const Config& cfg,
+                                                   Acts::Logging::Level lvl)
+    : IAlgorithm("TrackToTruthJetAlgorithm", lvl), m_cfg(cfg) {
+  if (m_cfg.inputTracks.empty()) {
+    throw std::invalid_argument("Input tracks are not configured");
+  }
+  if (m_cfg.inputJets.empty()) {
+    throw std::invalid_argument("Input tracks are not configured");
+  }
+  m_inputTracks.initialize(m_cfg.inputTracks);
+  m_inputJets.initialize(m_cfg.inputJets);
+  m_outputTrackJets.initialize(m_cfg.outputTrackJets);
+}
+
+ProcessCode ActsExamples::TrackToTruthJetAlgorithm::execute(
+    const ActsExamples::AlgorithmContext& ctx) const {
+  ACTS_DEBUG("Executing track to truth jet matching algorithm");
+
+  const auto& tracks = m_inputTracks(ctx);
+  const auto& truthJets = m_inputJets(ctx);
+  // Take a copy that we will modify
+  Acts::FastJet::TrackJetContainer jets = truthJets;
+
+  ACTS_DEBUG("TrackToTruthJetAlg - Number of tracks: " << tracks.size());
+  ACTS_DEBUG("TrackToTruthJetAlg - Number of truth jets: " << truthJets.size());
+  ACTS_DEBUG("TrackToTruthJetAlg - Number of jets: " << jets.size());
+
+  for (const auto& track : tracks) {
+    double minDeltaR = m_cfg.maxDeltaR;
+
+    ACTS_VERBOSE("Track index: " << track.index()
+                                 << ", momentum: " << track.momentum().x()
+                                 << ", " << track.momentum().y() << ", "
+                                 << track.momentum().z());
+
+    Acts::FastJet::TruthJetBuilder* matchedJet = nullptr;
+
+    // Loop over the jets to find the closest one
+    std::size_t i = 0;
+    for (auto& jet : jets) {
+      Acts::FastJet::JetProperties jetProps(jet);
+
+      double jetPx = jetProps.getFourMomentum().x();
+      double jetPy = jetProps.getFourMomentum().y();
+      double jetPz = jetProps.getFourMomentum().z();
+      double trackPx = track.momentum().x();
+      double trackPy = track.momentum().y();
+      double trackPz = track.momentum().z();
+      double pjet = std::sqrt(jetPx * jetPx + jetPy * jetPy + jetPz * jetPz);
+      double ptrack =
+          std::sqrt(trackPx * trackPx + trackPy * trackPy + trackPz * trackPz);
+
+      // Calculate eta and phi for the jet and track
+      // Note: eta = arctanh(pz/|p|), phi = atan2(py, px)
+      // where theta is the polar angle of the momentum vector
+      // and phi is the azimuthal angle in the xy-plane.
+      // Here we use the four-momentum to calculate eta and phi.
+
+      double jetEta = std::atanh(jetPz / pjet);
+      double jetPhi = std::atan2(jetPy, jetPx);
+      double trackEta = std::atanh(trackPz / ptrack);
+      double trackPhi = std::atan2(trackPy, trackPx);
+
+      // Calculate delta R
+      double deltaEta = jetEta - trackEta;
+      double deltaPhi = jetPhi - trackPhi;
+      double deltaR = std::sqrt(deltaEta * deltaEta + deltaPhi * deltaPhi);
+
+      // Acts::Vector3 jetMom = jet.getFourMomentum().head<3>();
+      // double deltaR = Acts::VectorHelpers::deltaR(jetMom, track.momentum());
+
+      if (deltaR < minDeltaR) {
+        minDeltaR = deltaR;
+        matchedJet = &jet;
+      }
+
+      ACTS_DEBUG("Track " << track.index() << " delta R to jet " << i << ": "
+                          << deltaR
+                          << ", jet px: " << jetProps.getFourMomentum().x());
+      if (deltaR < m_cfg.maxDeltaR) {
+        ACTS_DEBUG("Track " << track.index() << " matches jet " << i
+                            << " with delta R: " << deltaR);
+      }
+
+      i++;
+    }  // loop over jets
+
+    if (matchedJet != nullptr) {
+      Acts::FastJet::JetProperties matchedJetProps(*matchedJet);
+      matchedJetProps.addTrack(track.index());
+    }
+
+  }  // loop over tracks
+
+  // Write the matched jets to the output
+  m_outputTrackJets(ctx, std::move(jets));
+
+  return ProcessCode::SUCCESS;
+}
+
+ProcessCode ActsExamples::TrackToTruthJetAlgorithm::finalize() {
+  ACTS_INFO("Finalizing track to truth jet matching algorithm");
+  return ProcessCode::SUCCESS;
+}
+
+};  // namespace ActsExamples

Examples/Algorithms/Jets/src/TruthJetAlgorithm.cpp

@@ -35,6 +35,8 @@ TruthJetAlgorithm::TruthJetAlgorithm(const Config& cfg,
 
 ProcessCode ActsExamples::TruthJetAlgorithm::execute(
     const ActsExamples::AlgorithmContext& ctx) const {
+  Acts::FastJet::TrackJetContainer outputJets;
+
   const auto& truthParticles = m_inputTruthParticles(ctx);
 
   ACTS_DEBUG("Number of truth particles: " << truthParticles.size());
@@ -56,16 +58,46 @@ ProcessCode ActsExamples::TruthJetAlgorithm::execute(
     inputPseudoJets.push_back(pseudoJet);
     particleIndex++;
   }
-  ACTS_DEBUG("Number of input pseudo jets: " << inputPseudoJets.size());
+  ACTS_DEBUG("Number of input pseudo jets from truth particles: "
+             << inputPseudoJets.size());
+
   // Run the jet clustering
   fastjet::ClusterSequence clusterSeq(inputPseudoJets, defaultJetDefinition);
+
   // Get the jets above a certain pt threshold
   std::vector<fastjet::PseudoJet> jets =
       sorted_by_pt(clusterSeq.inclusive_jets(m_cfg.jetPtMin));
-  ACTS_DEBUG("Number of clustered jets: " << jets.size());
-  // Store the jets in the output data handle
-  m_outputJets(ctx, std::move(jets));
+  ACTS_DEBUG("Number of clustered truth jets: " << jets.size());
+
+  // Prepare jets for the storage - conversion of jets to custom track jet class
+  // (and later add here the jet classification)
+
+  for (unsigned int i = 0; i < jets.size(); i++) {
+    // Get information on the jet constituents
+    std::vector<fastjet::PseudoJet> jetConstituents = jets[i].constituents();
+    std::vector<int> constituentIndices;
+    constituentIndices.reserve(jetConstituents.size());
+
+    Acts::Vector4 jetFourMomentum(jets[i].px(), jets[i].py(), jets[i].pz(),
+                                  jets[i].e());
+
+    // Initialize the (track) jet with 4-momentum
+    Acts::FastJet::TruthJetBuilder storedJet(jetFourMomentum);
+
+    // Add the jet constituents to the (track)jet
+    for (unsigned int j = 0; j < jetConstituents.size(); j++) {
+      // Get the index of the constituent in the original input pseudo jets
+      constituentIndices.push_back(jetConstituents[j].user_index());
+    }
+
+    Acts::FastJet::JetProperties jetProps(storedJet);
+    jetProps.setConstituents(constituentIndices);
+
+    outputJets.push_back(storedJet);
+    ACTS_DEBUG("Stored jet properties: " << jetProps);
+  }
 
+  m_outputJets(ctx, std::move(outputJets));
   return ProcessCode::SUCCESS;
 }
 

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

@@ -53,10 +53,11 @@
     defaults=[(None, None)] * 10 + [None] * 4,
 )
 
-TruthJetConfig = namedtuple(
-    "TruthJetConfig",
-    ["inputTruthParticles", "outputJets", "jetPtMin"],
-    defaults=[None, None],
+
+TrackToTruthJetConfig = namedtuple(
+    "TrackToTruthJetConfig",
+    ["inputTracks", "inputJets", "outputTrackJets", "maxDeltaR"],
+    defaults=[None, None, None, None],
 )
 
 
@@ -89,11 +90,12 @@ def _getParticleSelectionKWargs(config: ParticleSelectorConfig) -> dict:
     }
 
 
-def _getTruthJetKWargs(config: TruthJetConfig) -> dict:
+def _getTrackToTruthJetKWargs(config: TrackToTruthJetConfig) -> dict:
     return {
-        "inputTruthParticles": config.inputTruthParticles,
-        "outputJets": config.outputJets,
-        "jetPtMin": config.jetPtMin,
+        "inputTracks": config.inputTracks,
+        "inputJets": config.inputJets,
+        "outputTrackJets": config.outputTrackJets,
+        "maxDeltaR": config.maxDeltaR,
     }
 
 
@@ -852,17 +854,17 @@ def addDigiParticleSelection(
     )
 
 
-def addTruthJetAlg(
+def addTrackToTruthJetAlg(
     s: acts.examples.Sequencer,
-    config: TruthJetConfig,
+    config: TrackToTruthJetConfig,
     loglevel: Optional[acts.logging.Level] = None,
 ) -> None:
-    from acts.examples import TruthJetAlgorithm
+    from acts.examples import TrackToTruthJetAlgorithm
 
     customLogLevel = acts.examples.defaultLogging(s, loglevel)
-    truthJetAlg = acts.examples.TruthJetAlgorithm(
-        **acts.examples.defaultKWArgs(**_getTruthJetKWargs(config)),
+    trackToTruthJetAlg = acts.examples.TrackToTruthJetAlgorithm(
+        **acts.examples.defaultKWArgs(**_getTrackToTruthJetKWargs(config)),
         level=customLogLevel(),
     )
 
-    s.addAlgorithm(truthJetAlg)
+    s.addAlgorithm(trackToTruthJetAlg)

Examples/Python/src/ModuleEntry.cpp

@@ -40,6 +40,7 @@ void addTruthTracking(Context& ctx);
 void addTrackFitting(Context& ctx);
 void addTrackFinding(Context& ctx);
 void addTruthJet(Context& ctx);
+void addTrackToTruthJet(Context& ctx);
 void addVertexing(Context& ctx);
 void addAmbiguityResolution(Context& ctx);
 void addUtilities(Context& ctx);
@@ -79,6 +80,7 @@ void addModuleEntry(Context& ctx) {
   addTrackFitting(ctx);
   addTrackFinding(ctx);
   addTruthJet(ctx);
+  addTrackToTruthJet(ctx);
   addVertexing(ctx);
   addAmbiguityResolution(ctx);
   addUtilities(ctx);

Examples/Python/src/TruthJet.cpp

@@ -7,6 +7,7 @@
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
 #include "Acts/Utilities/Logger.hpp"
+#include "ActsExamples/Jets/TrackToTruthJetAlgorithm.hpp"
 #include "ActsExamples/Jets/TruthJetAlgorithm.hpp"
 #include "ActsPython/Utilities/Helpers.hpp"
 #include "ActsPython/Utilities/Macros.hpp"
@@ -29,4 +30,13 @@ void addTruthJet(Context& ctx) {
   ACTS_PYTHON_DECLARE_ALGORITHM(TruthJetAlgorithm, mex, "TruthJetAlgorithm",
                                 inputTruthParticles, outputJets, jetPtMin);
 }  // addTruthJet
+
+void addTrackToTruthJet(Context& ctx) {
+  auto mex = ctx.get("examples");
+
+  ACTS_PYTHON_DECLARE_ALGORITHM(TrackToTruthJetAlgorithm, mex,
+                                "TrackToTruthJetAlgorithm", inputTracks,
+                                inputJets, outputTrackJets, maxDeltaR);
+}  // addTrackToTruthJet
+
 }  // namespace ActsPython

Examples/Python/src/TruthJetStub.cpp

@@ -11,4 +11,5 @@
 
 namespace ActsPython {
 void addTruthJet(Context& /*ctx*/) {}
+void addTrackToTruthJet(Context& /*ctx*/) {}
 }  // namespace ActsPython