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<TrackContainer>>
0069       outputJetContainer{};
0070
0071   Acts::ScopedTimer globalTimer("TruthJetAlgorithm", logger(),
0072                                 Acts::Logging::DEBUG);
0073
0074   const SimParticleContainer& truthParticlesRaw = m_inputTruthParticles(ctx);
0075   std::vector<const SimParticle*> truthParticles;
0076   truthParticles.reserve(truthParticlesRaw.size());
0077   std::ranges::transform(truthParticlesRaw, std::back_inserter(truthParticles),
0078                          [](const auto& particle) { return &particle; });
0079   ACTS_DEBUG("Number of truth particles: " << truthParticles.size());
0080
0081   const fastjet::JetDefinition defaultJetDefinition = fastjet::JetDefinition(
0082       fastjet::antikt_algorithm, m_cfg.jetClusteringRadius);
0083
0084   // Get the 4-momentum information from the simulated truth particles
0085   // and create fastjet::PseudoJet objects
0086   std::vector<fastjet::PseudoJet> inputPseudoJets;
0087   {
0088     Acts::ScopedTimer timer("Input particle building", logger(),
0089                             Acts::Logging::DEBUG);
0090
0091     for (unsigned int i = 0; i < truthParticles.size(); ++i) {
0092       const auto* particle = truthParticles.at(i);
0093
0094       detail::PrimaryVertexIdGetter primaryVertexIdGetter;
0095
0096       ACTS_VERBOSE("Primary vertex ID: "
0097                    << primaryVertexIdGetter(*particle).vertexPrimary()
0098                    << ", PDG: " << static_cast<int>(particle->pdg()) << ", pT: "
0099                    << particle->transverseMomentum() / Acts::UnitConstants::GeV
0100                    << " GeV");
0101
0102       if (m_cfg.clusterHSParticlesOnly &&
0103           primaryVertexIdGetter(*particle).vertexPrimary() != 1) {
0104         continue;
0105       }
0106
0107       fastjet::PseudoJet pseudoJet(
0108           particle->momentum().x(), particle->momentum().y(),
0109           particle->momentum().z(), particle->energy());
0110
0111       pseudoJet.set_user_index(i);
0112       inputPseudoJets.push_back(pseudoJet);
0113     }
0114   }
0115
0116   ACTS_DEBUG("Number of input pseudo jets: " << inputPseudoJets.size());
0117
0118   std::vector<fastjet::PseudoJet> jets;
0119   fastjet::ClusterSequence clusterSeq;
0120   {
0121     Acts::ScopedTimer timer("Jet clustering", logger(), Acts::Logging::DEBUG);
0122     // Run the jet clustering
0123     clusterSeq =
0124         fastjet::ClusterSequence(inputPseudoJets, defaultJetDefinition);
0125     // Get the jets above a certain pt threshold
0126     jets = sorted_by_pt(clusterSeq.inclusive_jets(m_cfg.jetPtMin));
0127     // Apply eta range cut if specified
0128     double minEta = m_cfg.jetEtaRange.first;
0129     double maxEta = m_cfg.jetEtaRange.second;
0130
0131     std::erase_if(jets, [minEta, maxEta](const auto& jet) {
0132       return jet.eta() < minEta || jet.eta() > maxEta;
0133     });
0134     ACTS_DEBUG("Number of clustered jets: " << jets.size());
0135   }
0136
0137   // Find hadrons for jet labeling with sim particles
0138   std::vector<std::pair<ActsPlugins::FastJet::JetLabel, const SimParticle*>>
0139       hadrons;
0140
0141   if (m_cfg.doJetLabeling) {
0142     Acts::ScopedTimer timer("hadron finding", logger(), Acts::Logging::DEBUG);
0143     ACTS_DEBUG("Jet labeling is enabled. Finding hadrons for jet labeling.");
0144
0145     // A lazy view over the simulated particles for hadron finding
0146     auto hadronView =
0147         truthParticles | std::views::filter([this](const auto* particle) {
0148           if (m_cfg.jetLabelingHSHadronsOnly) {
0149             detail::PrimaryVertexIdGetter primaryVertexIdGetter;
0150             if (primaryVertexIdGetter(*particle).vertexPrimary() != 1) {
0151               return false;
0152             }
0153           }
0154
0155           Acts::PdgParticle pdgId{particle->pdg()};
0156           if (!Acts::ParticleIdHelper::isHadron(pdgId)) {
0157             return false;
0158           }
0159
0160           // Apply a pt cut on B or C hadrons
0161           auto label =
0162               jetLabelFromHadronType(Acts::ParticleIdHelper::hadronType(pdgId));
0163           using enum ActsPlugins::FastJet::JetLabel;
0164
0165           if (label == BJet || label == CJet) {
0166             if (particle->transverseMomentum() < m_cfg.jetLabelingHadronPtMin) {
0167               return false;
0168             }
0169           }
0170           return true;
0171         }) |
0172         std::views::transform([](const auto* particle) {
0173           Acts::PdgParticle pdgId{particle->pdg()};
0174           auto type = Acts::ParticleIdHelper::hadronType(pdgId);
0175           auto label = jetLabelFromHadronType(type);
0176           return std::make_pair(label, particle);
0177         }) |
0178         std::views::filter([](const auto& hadron) {
0179           return hadron.first > ActsPlugins::FastJet::JetLabel::Unknown;
0180         });
0181
0182     std::ranges::copy(hadronView, std::back_inserter(hadrons));
0183
0184     // Deduplicate hadrons based on their pdg id
0185     std::ranges::sort(hadrons, [](const auto& a, const auto& b) {
0186       return a.second->pdg() < b.second->pdg();
0187     });
0188
0189     auto unique = std::ranges::unique(hadrons);
0190     hadrons.erase(unique.begin(), unique.end());
0191   }
0192
0193   // Jet classification
0194
0195   auto classifyJet = [&](const fastjet::PseudoJet& jet) {
0196     auto hadronsInJetView =
0197         hadrons | std::views::filter([&jet, this](const auto& hadron) {
0198           const auto& momentum = hadron.second->momentum();
0199           Acts::Vector3 hadronJetMom{momentum[0], momentum[1], momentum[2]};
0200           Acts::Vector3 jetMom{jet.px(), jet.py(), jet.pz()};
0201           return Acts::VectorHelpers::deltaR(jetMom, hadronJetMom) <
0202                  m_cfg.jetLabelingDeltaR;
0203         }) |
0204         std::views::transform([](const auto& hadron) {
0205           return std::pair{
0206               hadron.second,
0207               jetLabelFromHadronType(Acts::ParticleIdHelper::hadronType(
0208                   Acts::PdgParticle{hadron.second->pdg()}))};
0209         });
0210
0211     std::vector<std::pair<const SimParticle*, ActsPlugins::FastJet::JetLabel>>
0212         hadronsInJet;
0213     std::ranges::copy(hadronsInJetView, std::back_inserter(hadronsInJet));
0214
0215     ACTS_VERBOSE("-> hadrons in jet: " << hadronsInJet.size());
0216     for (const auto& hadron : hadronsInJet) {
0217       ACTS_VERBOSE(
0218           "  - " << hadron.first->pdg() << " "
0219                  << Acts::findName(Acts::PdgParticle{hadron.first->pdg()})
0220                         .value_or("UNKNOWN")
0221                  << " label=" << hadron.second);
0222     }
0223
0224     auto maxHadronIt = std::ranges::max_element(
0225         hadronsInJet, [](const auto& a, const auto& b) { return a < b; },
0226         [](const auto& a) {
0227           const auto& [hadron, label] = a;
0228           return label;
0229         });
0230
0231     if (maxHadronIt == hadronsInJet.end()) {
0232       // Now hadronic "jet"
0233       return ActsPlugins::FastJet::JetLabel::Unknown;
0234     }
0235
0236     const auto& [maxHadron, maxHadronLabel] = *maxHadronIt;
0237
0238     ACTS_VERBOSE("-> max hadron type="
0239                  << Acts::findName(Acts::PdgParticle{maxHadron->pdg()})
0240                         .value_or("UNKNOWN")
0241                  << " label=" << maxHadronLabel);
0242
0243     return maxHadronLabel;
0244   };  // jet classification
0245
0246   boost::container::flat_map<ActsPlugins::FastJet::JetLabel, std::size_t>
0247       jetLabelCounts;
0248
0249   {
0250     Acts::AveragingScopedTimer timer("Jet classification", logger(),
0251                                      Acts::Logging::DEBUG);
0252
0253     for (unsigned int i = 0; i < jets.size(); ++i) {
0254       const auto& jet = jets.at(i);
0255
0256       // If jet labeling is enabled, classify the jet based on its hadronic
0257       // content
0258       ActsPlugins::FastJet::JetLabel jetLabel =
0259           ActsPlugins::FastJet::JetLabel::Unknown;
0260       if (m_cfg.doJetLabeling) {
0261         ACTS_DEBUG("Classifying jet " << i);
0262         auto sample = timer.sample();
0263         jetLabel = classifyJet(jet);
0264       }
0265
0266       // Initialize truth jet for storing in the output container
0267       Acts::Vector4 jetFourMom{jet.px(), jet.py(), jet.pz(), jet.e()};
0268       ActsPlugins::FastJet::TruthJet<TrackContainer> truthJet(jetFourMom,
0269                                                               jetLabel);
0270
0271       std::vector<fastjet::PseudoJet> jetConstituents = jet.constituents();
0272       std::vector<int> constituentIndices;
0273       constituentIndices.reserve(jetConstituents.size());
0274
0275       for (unsigned int j = 0; j < jetConstituents.size(); ++j) {
0276         constituentIndices.push_back(jetConstituents[j].user_index());
0277       }
0278
0279       truthJet.setConstituentIndices(constituentIndices);
0280       outputJetContainer.push_back(truthJet);
0281
0282       jetLabelCounts[jetLabel] += 1;
0283
0284       ACTS_VERBOSE("-> jet label: " << jetLabel);
0285       ACTS_VERBOSE("-> jet constituents: ");
0286
0287       if (logger().doPrint(Acts::Logging::VERBOSE)) {
0288         for (const auto& constituent : constituentIndices) {
0289           const auto& particle = truthParticles.at(constituent);
0290           ACTS_VERBOSE("- " << particle);
0291         }
0292       }
0293     }
0294   }
0295
0296   ACTS_DEBUG("-> jet label counts: ");
0297   for (const auto& [label, count] : jetLabelCounts) {
0298     ACTS_DEBUG("  - " << label << ": " << count);
0299   }
0300
0301   m_outputJets(ctx, std::move(outputJetContainer));
0302
0303   return ProcessCode::SUCCESS;
0304 }
0305
0306 ProcessCode ActsExamples::TruthJetAlgorithm::finalize() {
0307   ACTS_INFO("Finalizing truth jet clustering");
0308   return ProcessCode::SUCCESS;
0309 }
0310
0311 }  // namespace ActsExamples