Jets/src/TruthJetAlgorithm.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #include "ActsExamples/Jets/TruthJetAlgorithm.hpp"
0010
0011 #include "Acts/Definitions/ParticleData.hpp"
0012 #include "Acts/Definitions/PdgParticle.hpp"
0013 #include "Acts/Definitions/Units.hpp"
0014 #include "Acts/Utilities/Logger.hpp"
0015 #include "Acts/Utilities/ScopedTimer.hpp"
0016 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0017 #include "ActsFatras/EventData/ProcessType.hpp"
0018
0019 #include <algorithm>
0020 #include <ranges>
0021 #include <stdexcept>
0022
0023 #include <boost/container/flat_map.hpp>
0024 #include <fastjet/ClusterSequence.hh>
0025 #include <fastjet/JetDefinition.hh>
0026 #include <fastjet/PseudoJet.hh>
0027
0028 namespace ActsExamples {
0029
0030 TruthJetAlgorithm::TruthJetAlgorithm(const Config& cfg,
0031                                      Acts::Logging::Level lvl)
0032     : IAlgorithm("TruthJetAlgorithm", lvl), m_cfg(cfg) {
0033   if (m_cfg.inputTruthParticles.empty()) {
0034     throw std::invalid_argument("Input particles collection is not configured");
0035   }
0036   m_inputTruthParticles.initialize(m_cfg.inputTruthParticles);
0037   m_outputJets.initialize(m_cfg.outputJets);
0038 }
0039
0040 namespace {
0041 ActsPlugins::FastJet::JetLabel jetLabelFromHadronType(
0042     Acts::HadronType hadronType) {
0043   using enum Acts::HadronType;
0044   switch (hadronType) {
0045     case BBbarMeson:
0046     case BottomMeson:
0047     case BottomBaryon:
0048       return ActsPlugins::FastJet::JetLabel::BJet;
0049     case CCbarMeson:
0050     case CharmedMeson:
0051     case CharmedBaryon:
0052       return ActsPlugins::FastJet::JetLabel::CJet;
0053     case StrangeMeson:
0054     case StrangeBaryon:
0055     case LightMeson:
0056     case LightBaryon:
0057       return ActsPlugins::FastJet::JetLabel::LightJet;
0058     default:
0059       return ActsPlugins::FastJet::JetLabel::Unknown;
0060   }
0061 }
0062
0063 }  // namespace
0064
0065 ProcessCode ActsExamples::TruthJetAlgorithm::execute(
0066     const ActsExamples::AlgorithmContext& ctx) const {
0067   // Initialize the output container
0068   std::vector<ActsPlugins::FastJet::TruthJet> outputJetContainer{};
0069
0070   Acts::ScopedTimer globalTimer("TruthJetAlgorithm", logger(),
0071                                 Acts::Logging::DEBUG);
0072
0073   const SimParticleContainer& truthParticlesRaw = m_inputTruthParticles(ctx);
0074   std::vector<const SimParticle*> truthParticles;
0075   truthParticles.reserve(truthParticlesRaw.size());
0076   std::ranges::transform(truthParticlesRaw, std::back_inserter(truthParticles),
0077                          [](const auto& particle) { return &particle; });
0078   ACTS_DEBUG("Number of truth particles: " << truthParticles.size());
0079
0080   const fastjet::JetDefinition defaultJetDefinition = fastjet::JetDefinition(
0081       fastjet::antikt_algorithm, m_cfg.jetClusteringRadius);
0082
0083   // Get the 4-momentum information from the simulated truth particles
0084   // and create fastjet::PseudoJet objects
0085   std::vector<fastjet::PseudoJet> inputPseudoJets;
0086   {
0087     Acts::ScopedTimer timer("Input particle building", logger(),
0088                             Acts::Logging::DEBUG);
0089
0090     for (unsigned int i = 0; i < truthParticles.size(); ++i) {
0091       const auto* particle = truthParticles.at(i);
0092
0093       detail::PrimaryVertexIdGetter primaryVertexIdGetter;
0094
0095       ACTS_VERBOSE("Primary vertex ID: "
0096                    << primaryVertexIdGetter(*particle).vertexPrimary()
0097                    << ", PDG: " << static_cast<int>(particle->pdg()) << ", pT: "
0098                    << particle->transverseMomentum() / Acts::UnitConstants::GeV
0099                    << " GeV");
0100
0101       if (m_cfg.clusterHSParticlesOnly &&
0102           primaryVertexIdGetter(*particle).vertexPrimary() != 1) {
0103         continue;
0104       }
0105
0106       fastjet::PseudoJet pseudoJet(
0107           particle->momentum().x(), particle->momentum().y(),
0108           particle->momentum().z(), particle->energy());
0109
0110       pseudoJet.set_user_index(i);
0111       inputPseudoJets.push_back(pseudoJet);
0112     }
0113   }
0114
0115   ACTS_DEBUG("Number of input pseudo jets: " << inputPseudoJets.size());
0116
0117   std::vector<fastjet::PseudoJet> jets;
0118   fastjet::ClusterSequence clusterSeq;
0119   {
0120     Acts::ScopedTimer timer("Jet clustering", logger(), Acts::Logging::DEBUG);
0121     // Run the jet clustering
0122     clusterSeq =
0123         fastjet::ClusterSequence(inputPseudoJets, defaultJetDefinition);
0124     // Get the jets above a certain pt threshold
0125     jets = sorted_by_pt(clusterSeq.inclusive_jets(m_cfg.jetPtMin));
0126     // Apply eta range cut if specified
0127     double minEta = m_cfg.jetEtaRange.first;
0128     double maxEta = m_cfg.jetEtaRange.second;
0129
0130     std::erase_if(jets, [minEta, maxEta](const auto& jet) {
0131       return jet.eta() < minEta || jet.eta() > maxEta;
0132     });
0133     ACTS_DEBUG("Number of clustered jets: " << jets.size());
0134   }
0135
0136   // Find hadrons for jet labeling with sim particles
0137   std::vector<std::pair<ActsPlugins::FastJet::JetLabel, const SimParticle*>>
0138       hadrons;
0139
0140   if (m_cfg.doJetLabeling) {
0141     Acts::ScopedTimer timer("hadron finding", logger(), Acts::Logging::DEBUG);
0142     ACTS_DEBUG("Jet labeling is enabled. Finding hadrons for jet labeling.");
0143
0144     // A lazy view over the simulated particles for hadron finding
0145     auto hadronView =
0146         truthParticles | std::views::filter([this](const auto* particle) {
0147           if (m_cfg.jetLabelingHSHadronsOnly) {
0148             detail::PrimaryVertexIdGetter primaryVertexIdGetter;
0149             if (primaryVertexIdGetter(*particle).vertexPrimary() != 1) {
0150               return false;
0151             }
0152           }
0153
0154           Acts::PdgParticle pdgId{particle->pdg()};
0155           if (!Acts::ParticleIdHelper::isHadron(pdgId)) {
0156             return false;
0157           }
0158
0159           // Apply a pt cut on B or C hadrons
0160           auto label =
0161               jetLabelFromHadronType(Acts::ParticleIdHelper::hadronType(pdgId));
0162           using enum ActsPlugins::FastJet::JetLabel;
0163
0164           if (label == BJet || label == CJet) {
0165             if (particle->transverseMomentum() < m_cfg.jetLabelingHadronPtMin) {
0166               return false;
0167             }
0168           }
0169           return true;
0170         }) |
0171         std::views::transform([](const auto* particle) {
0172           Acts::PdgParticle pdgId{particle->pdg()};
0173           auto type = Acts::ParticleIdHelper::hadronType(pdgId);
0174           auto label = jetLabelFromHadronType(type);
0175           return std::make_pair(label, particle);
0176         }) |
0177         std::views::filter([](const auto& hadron) {
0178           return hadron.first > ActsPlugins::FastJet::JetLabel::Unknown;
0179         });
0180
0181     std::ranges::copy(hadronView, std::back_inserter(hadrons));
0182
0183     // Deduplicate hadrons based on their pdg id
0184     std::ranges::sort(hadrons, [](const auto& a, const auto& b) {
0185       return a.second->pdg() < b.second->pdg();
0186     });
0187
0188     auto unique = std::ranges::unique(hadrons);
0189     hadrons.erase(unique.begin(), unique.end());
0190   }
0191
0192   // Jet classification
0193
0194   auto classifyJet = [&](const fastjet::PseudoJet& jet) {
0195     auto hadronsInJetView =
0196         hadrons | std::views::filter([&jet, this](const auto& hadron) {
0197           const auto& momentum = hadron.second->momentum();
0198           Acts::Vector3 hadronJetMom{momentum[0], momentum[1], momentum[2]};
0199           Acts::Vector3 jetMom{jet.px(), jet.py(), jet.pz()};
0200           return Acts::VectorHelpers::deltaR(jetMom, hadronJetMom) <
0201                  m_cfg.jetLabelingDeltaR;
0202         }) |
0203         std::views::transform([](const auto& hadron) {
0204           return std::pair{
0205               hadron.second,
0206               jetLabelFromHadronType(Acts::ParticleIdHelper::hadronType(
0207                   Acts::PdgParticle{hadron.second->pdg()}))};
0208         });
0209
0210     std::vector<std::pair<const SimParticle*, ActsPlugins::FastJet::JetLabel>>
0211         hadronsInJet;
0212     std::ranges::copy(hadronsInJetView, std::back_inserter(hadronsInJet));
0213
0214     ACTS_VERBOSE("-> hadrons in jet: " << hadronsInJet.size());
0215     for (const auto& hadron : hadronsInJet) {
0216       ACTS_VERBOSE(
0217           "  - " << hadron.first->pdg() << " "
0218                  << Acts::findName(Acts::PdgParticle{hadron.first->pdg()})
0219                         .value_or("UNKNOWN")
0220                  << " label=" << hadron.second);
0221     }
0222
0223     auto maxHadronIt = std::ranges::max_element(
0224         hadronsInJet, [](const auto& a, const auto& b) { return a < b; },
0225         [](const auto& a) {
0226           const auto& [hadron, label] = a;
0227           return label;
0228         });
0229
0230     if (maxHadronIt == hadronsInJet.end()) {
0231       // Now hadronic "jet"
0232       return ActsPlugins::FastJet::JetLabel::Unknown;
0233     }
0234
0235     const auto& [maxHadron, maxHadronLabel] = *maxHadronIt;
0236
0237     ACTS_VERBOSE("-> max hadron type="
0238                  << Acts::findName(Acts::PdgParticle{maxHadron->pdg()})
0239                         .value_or("UNKNOWN")
0240                  << " label=" << maxHadronLabel);
0241
0242     return maxHadronLabel;
0243   };  // jet classification
0244
0245   boost::container::flat_map<ActsPlugins::FastJet::JetLabel, std::size_t>
0246       jetLabelCounts;
0247
0248   {
0249     Acts::AveragingScopedTimer timer("Jet classification", logger(),
0250                                      Acts::Logging::DEBUG);
0251
0252     for (unsigned int i = 0; i < jets.size(); ++i) {
0253       const auto& jet = jets.at(i);
0254
0255       // If jet labeling is enabled, classify the jet based on its hadronic
0256       // content
0257       ActsPlugins::FastJet::JetLabel jetLabel =
0258           ActsPlugins::FastJet::JetLabel::Unknown;
0259       if (m_cfg.doJetLabeling) {
0260         ACTS_DEBUG("Classifying jet " << i);
0261         auto sample = timer.sample();
0262         jetLabel = classifyJet(jet);
0263       }
0264
0265       // Initialize truth jet for storing in the output container
0266       Acts::Vector4 jetFourMom{jet.px(), jet.py(), jet.pz(), jet.e()};
0267       ActsPlugins::FastJet::TruthJet truthJet(jetFourMom, jetLabel);
0268
0269       std::vector<fastjet::PseudoJet> jetConstituents = jet.constituents();
0270       std::vector<int> constituentIndices;
0271       constituentIndices.reserve(jetConstituents.size());
0272
0273       for (unsigned int j = 0; j < jetConstituents.size(); ++j) {
0274         constituentIndices.push_back(jetConstituents[j].user_index());
0275       }
0276
0277       truthJet.setConstituentIndices(constituentIndices);
0278       outputJetContainer.push_back(truthJet);
0279
0280       jetLabelCounts[jetLabel] += 1;
0281
0282       ACTS_VERBOSE("-> jet label: " << jetLabel);
0283       ACTS_VERBOSE("-> jet constituents: ");
0284
0285       if (logger().doPrint(Acts::Logging::VERBOSE)) {
0286         for (const auto& constituent : constituentIndices) {
0287           const auto& particle = truthParticles.at(constituent);
0288           ACTS_VERBOSE("- " << particle);
0289         }
0290       }
0291     }
0292   }
0293
0294   ACTS_DEBUG("-> jet label counts: ");
0295   for (const auto& [label, count] : jetLabelCounts) {
0296     ACTS_DEBUG("  - " << label << ": " << count);
0297   }
0298
0299   m_outputJets(ctx, std::move(outputJetContainer));
0300
0301   return ProcessCode::SUCCESS;
0302 }
0303
0304 ProcessCode ActsExamples::TruthJetAlgorithm::finalize() {
0305   ACTS_INFO("Finalizing truth jet clustering");
0306   return ProcessCode::SUCCESS;
0307 }
0308
0309 }  // namespace ActsExamples