algorithms/reco/JetReconstruction.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 Derek Anderson, Zhongling Ji, Dmitry Kalinkin, John Lajoie
0003
0004 // class definition
0005 #include "JetReconstruction.h"
0006
0007 // for error handling
0008 #include <JANA/JException.h>
0009 #include <edm4hep/MCParticleCollection.h> // IWYU pragma: keep
0010 #include <edm4hep/Vector3f.h>
0011 #include <edm4hep/utils/vector_utils.h>
0012 #include <fastjet/ClusterSequenceArea.hh>
0013 #include <fastjet/GhostedAreaSpec.hh>
0014 // for fastjet objects
0015 #include <fastjet/PseudoJet.hh>
0016 #include <fastjet/contrib/Centauro.hh>
0017 #include <fmt/core.h>
0018 #include <gsl/pointers>
0019 #include <stdexcept>
0020 #include <vector>
0021
0022 #include "algorithms/reco/JetReconstructionConfig.h"
0023
0024 using namespace fastjet;
0025
0026 namespace eicrecon {
0027
0028 template <typename InputT> void JetReconstruction<InputT>::init() {
0029
0030   this->trace("Initialized");
0031
0032   // if specified algorithm, recomb. scheme, or area type
0033   // are not defined, then issue error and throw exception
0034   try {
0035     m_mapJetAlgo.at(m_cfg.jetAlgo);
0036   } catch (std::out_of_range& out) {
0037     this->error(" Unknown jet algorithm \"{}\" specified!", m_cfg.jetAlgo);
0038     throw JException(out.what());
0039   }
0040
0041   try {
0042     m_mapRecombScheme.at(m_cfg.recombScheme);
0043   } catch (std::out_of_range& out) {
0044     this->error(" Unknown recombination scheme \"{}\" specified!", m_cfg.recombScheme);
0045     throw JException(out.what());
0046   }
0047
0048   try {
0049     m_mapAreaType.at(m_cfg.areaType);
0050   } catch (std::out_of_range& out) {
0051     this->error(" Unknown area type \"{}\" specified!", m_cfg.areaType);
0052     throw JException(out.what());
0053   }
0054
0055   // Choose jet definition based on no. of parameters
0056   switch (m_mapJetAlgo[m_cfg.jetAlgo]) {
0057
0058   // contributed algorithms
0059   case JetAlgorithm::plugin_algorithm:
0060
0061     // expand to other algorithms as required
0062     if (m_cfg.jetContribAlgo == "Centauro") {
0063       m_jet_plugin = std::make_unique<contrib::CentauroPlugin>(m_cfg.rJet);
0064       m_jet_def    = std::make_unique<JetDefinition>(m_jet_plugin.get());
0065     } else {
0066       this->error(" Unknown contributed FastJet algorithm \"{}\" specified!", m_cfg.jetContribAlgo);
0067       throw JException("Invalid contributed FastJet algorithm");
0068     }
0069     break;
0070
0071   // 0 parameter algorithms
0072   case JetAlgorithm::ee_kt_algorithm:
0073     m_jet_def = std::make_unique<JetDefinition>(m_mapJetAlgo[m_cfg.jetAlgo],
0074                                                 m_mapRecombScheme[m_cfg.recombScheme]);
0075     break;
0076
0077   // 2 parameter algorithms
0078   case JetAlgorithm::genkt_algorithm:
0079     [[fallthrough]];
0080
0081   case JetAlgorithm::ee_genkt_algorithm:
0082     m_jet_def = std::make_unique<JetDefinition>(m_mapJetAlgo[m_cfg.jetAlgo], m_cfg.rJet, m_cfg.pJet,
0083                                                 m_mapRecombScheme[m_cfg.recombScheme]);
0084     break;
0085
0086   // all others have only 1 parameter
0087   default:
0088     m_jet_def = std::make_unique<JetDefinition>(m_mapJetAlgo[m_cfg.jetAlgo], m_cfg.rJet,
0089                                                 m_mapRecombScheme[m_cfg.recombScheme]);
0090     break;
0091
0092   } // end switch (jet algorithm)
0093
0094   // Define jet area
0095   m_area_def = std::make_unique<AreaDefinition>(
0096       m_mapAreaType[m_cfg.areaType],
0097       GhostedAreaSpec(m_cfg.ghostMaxRap, m_cfg.numGhostRepeat, m_cfg.ghostArea));
0098
0099 } // end 'init()'
0100
0101 template <typename InputT>
0102 void JetReconstruction<InputT>::process(
0103     const typename JetReconstructionAlgorithm<InputT>::Input& input,
0104     const typename JetReconstructionAlgorithm<InputT>::Output& output) const {
0105   // Grab input collections
0106   const auto [input_collection] = input;
0107   auto [jet_collection]         = output;
0108
0109   // extract input momenta and collect into pseudojets
0110   std::vector<PseudoJet> particles;
0111   for (unsigned iInput = 0; const auto& input : *input_collection) {
0112
0113     // get 4-vector
0114     const auto& momentum = input.getMomentum();
0115     const auto& energy   = input.getEnergy();
0116     const auto pt        = edm4hep::utils::magnitudeTransverse(momentum);
0117
0118     // Only cluster particles within the given pt Range
0119     if ((pt > m_cfg.minCstPt) && (pt < m_cfg.maxCstPt)) {
0120       particles.emplace_back(momentum.x, momentum.y, momentum.z, energy);
0121       particles.back().set_user_index(iInput);
0122     }
0123     ++iInput;
0124   }
0125
0126   // Skip empty
0127   if (particles.empty()) {
0128     this->trace("  Empty particle list.");
0129     return;
0130   }
0131   this->trace("  Number of particles: {}", particles.size());
0132
0133   // Run the clustering, extract the jets
0134   fastjet::ClusterSequenceArea m_clus_seq(particles, *m_jet_def, *m_area_def);
0135   std::vector<PseudoJet> jets = sorted_by_pt(m_clus_seq.inclusive_jets(m_cfg.minJetPt));
0136
0137   // Print out some infos
0138   this->trace("  Clustering with : {}", m_jet_def->description());
0139
0140   // loop over jets
0141   for (unsigned i = 0; i < jets.size(); i++) {
0142
0143     this->trace("  jet {}: pt = {}, y = {}, phi = {}", i, jets[i].pt(), jets[i].rap(),
0144                 jets[i].phi());
0145
0146     // create jet to store in output collection
0147     edm4eic::MutableReconstructedParticle jet_output = jet_collection->create();
0148     jet_output.setMomentum(edm4hep::Vector3f(jets[i].px(), jets[i].py(), jets[i].pz()));
0149     jet_output.setEnergy(jets[i].e());
0150     jet_output.setMass(jets[i].m());
0151
0152     // link constituents to jet kinematic info
0153     std::vector<PseudoJet> csts = jets[i].constituents();
0154     for (unsigned j = 0; j < csts.size(); j++) {
0155       jet_output.addToParticles(input_collection->at(csts[j].user_index()));
0156     } // for constituent j
0157   }   // for jet i
0158
0159   // return the jets
0160   return;
0161 } // end 'process(const T&)'
0162
0163 template class JetReconstruction<edm4eic::ReconstructedParticle>;
0164
0165 } // end namespace eicrecon