Adding tools for jet clustering for cluster and truth particles (#80)

* Small changes to cmake lists * Adding tools for jet reconstruction * Adding truth jet algorithms * Fixing CMake change * Addressing some comments from the PR * Switching to thread safe Gaudi functional * Minor code cleanup * More code cleanup * Adding an extra algorithm to filter truth particles * One last removal for filtering * Adding separate jet clustering helper class * Adding some checks of the configuration in ClusterJet * A few updates to namespace etc * Addressing some changes from MR * Removing some commented lines * Update RecCalorimeter/src/components/FilterTruthParticles.cpp Co-authored-by: Brieuc Francois <[email protected]> * Removing more unnecessary code * Adding tool descriptions * Addressing a few more MR comments * Fixing the tests in the CMake file * Changing the naming for truth particle filter * Fixing naming for truth jet test * Switching to using associations for truth jets --------- Co-authored-by: Brieuc Francois <[email protected]>
HEP-FCC · Jul 9, 2024 · 5f27a43 · 5f27a43
1 parent 0341395
commit 5f27a43
Show file tree

Hide file tree

Showing 13 changed files with 957 additions and 5 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -29,6 +29,7 @@ find_package(EDM4HEP 0.10.1) # implicit: Podio
 find_package(DD4hep)
 find_package(k4geo)
 find_package(ONNXRuntime)
+find_package(FastJet)
 # New versions of ONNRuntime package provide onnxruntime-Config.cmake
 # and use the name onnxruntime
 find_package(onnxruntime)
@@ -64,9 +65,6 @@ add_subdirectory(RecFCChhCalorimeter)
 add_subdirectory(RecFCCeeCalorimeter)
 
 
-
-
-
 install(EXPORT ${PROJECT_NAME}Targets
   NAMESPACE ${PROJECT_NAME}::
   FILE "${PROJECT_NAME}Targets.cmake"

diff --git a/RecCalorimeter/CMakeLists.txt b/RecCalorimeter/CMakeLists.txt
@@ -3,9 +3,10 @@
 ################################################################################
 
 
-file(GLOB _module_sources src/components/*.cpp)
+file(GLOB _module_sources src/components/*.cpp )
+
 gaudi_add_module(k4RecCalorimeterPlugins
-                 SOURCES ${_module_sources}
+                 SOURCES ${_module_sources} 
                  LINK k4FWCore::k4FWCore
                       k4FWCore::k4Interface
                       Gaudi::GaudiAlgLib
@@ -17,6 +18,7 @@ gaudi_add_module(k4RecCalorimeterPlugins
                       ROOT::Hist
                       k4geo::detectorSegmentations
                       k4geo::detectorCommon
+                      ${FASTJET_LIBRARIES}
                       )
 
 install(TARGETS k4RecCalorimeterPlugins
@@ -26,8 +28,24 @@ install(TARGETS k4RecCalorimeterPlugins
   COMPONENT dev)
 
 
+
+include_directories(${FASTJET_INCLUDE_DIRS} include/RecCalorimeter)
+link_directories(${FASTJET_LIBRARIES})
+
 include(CTest)
 
+add_test(NAME FCC_createJets
+  COMMAND k4run RecCalorimeter/tests/options/runJetReco.py
+         WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
+)
+set_test_env(FCC_createJets)
+
+add_test(NAME FCC_createTruthJets
+  COMMAND k4run RecCalorimeter/tests/options/runTruthJetReco.py
+         WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
+)
+set_test_env(FCC_createTruthJets)
+
 #install(DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/tests/options DESTINATION ${CMAKE_INSTALL_DATADIR}/${CMAKE_PROJECT_NAME}/Reconstruction/RecCalorimeter)
 #
 #gaudi_add_test(genJetClustering

diff --git a/RecCalorimeter/include/RecCalorimeter/ClusterJet.h b/RecCalorimeter/include/RecCalorimeter/ClusterJet.h
@@ -0,0 +1,77 @@
+#ifndef RECCALORIMETER_CLUSTERJET_H
+#define RECCALORIMETER_CLUSTERJET_H
+
+// std
+#include <map>
+
+// Gaudi
+#include "GaudiAlg/GaudiAlgorithm.h"
+
+#include "fastjet/JetDefinition.hh"
+
+// EDM4hep
+#include "edm4hep/ClusterCollection.h"
+
+
+namespace k4::recCalo{
+
+class ClusterInfo : public fastjet::PseudoJet::UserInfoBase{
+ public:
+ ClusterInfo(const int & index) : _index(index){}
+  int index() const{ return _index; }
+ protected:
+  int _index;
+};
+
+
+/** @class ClusterJet k4RecCalorimeter/RecCalorimeter/src/components/ClusterJet.h
+ *
+ *  Helper class for reconstructing jets, used to avoid code duplication in the classes that require specific input types.
+ *  It takes a vector of fastjet::PseudoJets as inputs, and returns the a vector of pseudojets after clustering.
+ *  
+ *  JetAlg: A string corresponding to the jet algorithm for clustering
+ *  JetRadius: The radius of the jets being clustered
+ *  MinPt: The pT threshold below which jets are ignored
+ *  isExclusiveClustering: 1 if jets should use an exclusive clustering, 0 otherwise
+ *  clusterArgs: Any other clustering arguments, such as the number of jets for exclusive clustering. 
+ *               This is not used everywhere, so check the code to make sure it is implemented for your use-case.
+ *
+ *  @author Jennifer Roloff
+ *  @date 2024-7
+ */
+
+
+
+class ClusterJet {
+public:
+  ClusterJet(std::string jetAlg, double jetRadius, int isExclusiveClustering = 0, double minPt = 0, int clusterArgs = 0);
+  StatusCode initialize();
+  ~ClusterJet(){if(m_clustSeq) delete m_clustSeq;}
+
+  std::vector<fastjet::PseudoJet> cluster(const   std::vector<fastjet::PseudoJet> clustersPJ);
+
+private:
+
+
+
+  std::map<std::string, fastjet::JetAlgorithm> m_jetAlgMap = {
+                        {"kt",                    fastjet::JetAlgorithm::kt_algorithm},
+                        {"cambridge",             fastjet::JetAlgorithm::cambridge_algorithm},
+                        {"antikt",                fastjet::JetAlgorithm::antikt_algorithm},
+                        {"genkt",                 fastjet::JetAlgorithm::genkt_algorithm},
+                        {"ee_kt",                 fastjet::JetAlgorithm::ee_kt_algorithm},
+                        {"ee_genkt",              fastjet::JetAlgorithm::ee_genkt_algorithm},
+  };
+
+  std::string m_jetAlg = "antikt";
+  double m_jetRadius = 0.4;
+  int m_isExclusiveClustering = 0; // Inclusive clustering by default
+  double m_minPt = 10; // Only relevant for inclusive clustering
+  int m_njets = 0; // Only relevant for exclusive clustering
+
+  fastjet::ClusterSequence* m_clustSeq;
+
+};
+
+}
+#endif /* RECCALORIMETER_CLUSTERJET_H */
diff --git a/RecCalorimeter/src/components/ClusterJet.cpp b/RecCalorimeter/src/components/ClusterJet.cpp
@@ -0,0 +1,47 @@
+#include "ClusterJet.h"
+
+// std
+#include <vector>
+#include <math.h>
+
+namespace k4::recCalo{
+
+ClusterJet::ClusterJet(std::string jetAlg, double jetRadius, int isExclusiveClustering, double minPt, int njets): m_jetAlg(jetAlg), m_jetRadius(jetRadius), m_isExclusiveClustering(isExclusiveClustering), m_minPt(minPt), m_njets(njets){
+  m_clustSeq = nullptr;
+}
+
+
+StatusCode ClusterJet::initialize(){
+  if (m_jetAlgMap.find(m_jetAlg) == m_jetAlgMap.end()) {
+    std::cout << "ERROR: " << " is not in the list of supported jet algorithms" << std::endl;;
+    return StatusCode::FAILURE;
+  }
+  if(m_isExclusiveClustering > 1){
+    std::cout << "ERROR: " << "Clustering of " << m_isExclusiveClustering << " is currently not supported" << std::endl;;
+    return StatusCode::FAILURE;
+  }
+  return StatusCode::SUCCESS;
+}
+
+std::vector<fastjet::PseudoJet> ClusterJet::cluster(const   std::vector<fastjet::PseudoJet> clustersPJ) {
+  std::vector <fastjet::PseudoJet> jets;
+
+  fastjet::JetDefinition* jetDef = new fastjet::JetDefinition( m_jetAlgMap.at(m_jetAlg), m_jetRadius);
+  if(m_clustSeq) delete m_clustSeq;
+  m_clustSeq = new fastjet::ClusterSequence(clustersPJ, *jetDef);
+
+  // Note: initialize has already checked if m_isExclusiveClustering has the right range
+  if(m_isExclusiveClustering == 0){
+    jets = fastjet::sorted_by_pt(m_clustSeq->inclusive_jets(m_minPt));
+  }
+  else if(m_isExclusiveClustering == 1){
+    jets = fastjet::sorted_by_pt(m_clustSeq->exclusive_jets(m_njets));
+  }
+
+  return jets;
+
+}
+
+
+}
+
diff --git a/RecCalorimeter/src/components/CreateCaloJet.cpp b/RecCalorimeter/src/components/CreateCaloJet.cpp
@@ -0,0 +1,78 @@
+#include "CreateCaloJet.h"
+
+// std
+#include <vector>
+#include <math.h>
+
+// EDM4hep
+#include "edm4hep/ClusterCollection.h"
+
+DECLARE_COMPONENT(CreateCaloJet)
+
+CreateCaloJet::CreateCaloJet(const std::string& name, ISvcLocator* svcLoc) : Transformer(name, svcLoc,
+                    KeyValue("InputCollection", "CorrectedCaloClusters"),
+                    KeyValue("OutputCollection", "Jets")) {
+  declareProperty("JetAlg", m_jetAlg, "Name of jet clustering algorithm");
+  declareProperty("JetRadius", m_jetRadius, "Jet clustering radius");
+  declareProperty("MinPt", m_minPt, "Minimum pT for saved jets");
+  declareProperty("isExclusiveClustering", m_isExclusive, "1 if exclusive, 0 if inclusive");
+
+}
+
+
+StatusCode CreateCaloJet::initialize() {
+  clusterer = new k4::recCalo::ClusterJet(m_jetAlg, m_jetRadius, m_isExclusive, m_minPt);
+
+
+  return clusterer->initialize();
+}
+
+edm4hep::ReconstructedParticleCollection  CreateCaloJet::operator()(const edm4hep::ClusterCollection& input) const{
+  std::vector<fastjet::PseudoJet> clustersPJ;
+  int i=0;
+
+  // Need to convert cluster position in space to the momentum
+  for(auto cluster: input){
+    const edm4hep::Vector3f position = cluster.getPosition();
+    float x = position.x;
+    float y = position.y;
+    float z = position.z;
+    double theta = acos(sqrt(z*z / (x*x + y*y + z*z)));
+    double eta = -log(tan(theta / 2.));
+    double phi = atan2(y, x);
+    double pT = cluster.getEnergy()* sqrt( (x*x + y*y) / (x*x + y*y + z*z));
+
+    // Take clusters to be massless
+    fastjet::PseudoJet clusterPJ;
+    clusterPJ.reset_momentum_PtYPhiM(pT, eta, phi, 0);
+    // Attach the cluster index to the pseudojet
+    clusterPJ.set_user_info(new k4::recCalo::ClusterInfo(i));
+    clustersPJ.push_back(clusterPJ);
+    i++;
+  }
+
+
+  std::vector<fastjet::PseudoJet> inclusiveJets = clusterer->cluster(clustersPJ);
+
+
+  edm4hep::ReconstructedParticleCollection edmJets = edm4hep::ReconstructedParticleCollection();
+  //Add a reconstructed particle for each jet
+  for(auto cjet : inclusiveJets){
+    edm4hep::MutableReconstructedParticle jet;
+    jet.setMomentum(edm4hep::Vector3f(cjet.px(), cjet.py(), cjet.pz()));
+    jet.setEnergy(cjet.e());
+    jet.setMass(cjet.m());
+
+    // Also add the clusters that were used to make the jets
+    std::vector<fastjet::PseudoJet> constits = cjet.constituents();
+    for(auto constit : constits){
+      int index = constit.user_info<k4::recCalo::ClusterInfo>().index();
+      jet.addToClusters((input)[index]);
+    }
+    edmJets.push_back(jet);
+  }
+
+  return edmJets;
+}
+
+
diff --git a/RecCalorimeter/src/components/CreateCaloJet.h b/RecCalorimeter/src/components/CreateCaloJet.h
@@ -0,0 +1,59 @@
+#ifndef RECCALORIMETER_CREATECALOJET_H
+#define RECCALORIMETER_CREATECALOJET_H
+
+
+// Gaudi
+#include "GaudiAlg/GaudiAlgorithm.h"
+#include "GaudiAlg/Transformer.h"
+#include "k4FWCore/BaseClass.h"
+
+// k4FWCore
+#include "k4FWCore/DataHandle.h"
+
+//EDM4Hep
+#include "edm4hep/ReconstructedParticleCollection.h"
+
+// EDM4hep
+#include "edm4hep/ClusterCollection.h"
+
+
+#include "ClusterJet.h"
+
+
+
+/** @class CreateCaloJet k4RecCalorimeter/RecFCCeeCalorimeter/src/components/CreateCaloJet.h
+ *
+ *  Tool for reconstructing jets from calorimeter clusters.
+ *  It takes as input a ClusterCollection, and it outputs a ReconstructedParticleCollection with the jets.
+ *  
+ *  JetAlg: A string corresponding to the jet algorithm for clustering
+ *  JetRadius: The radius of the jets being clustered
+ *  MinPt: The pT threshold below which jets are ignored
+ *  isExclusiveClustering: 1 if jets should use an exclusive clustering, 0 otherwise
+ *
+ *  @author Jennifer Roloff
+ *  @date 2024-7
+ */
+
+
+
+class CreateCaloJet : public Gaudi::Functional::Transformer <edm4hep::ReconstructedParticleCollection(const edm4hep::ClusterCollection&), BaseClass_t>  {
+public:
+  CreateCaloJet(const std::string& name, ISvcLocator* svcLoc);
+  edm4hep::ReconstructedParticleCollection operator()(const edm4hep::ClusterCollection& input) const override;
+
+  StatusCode initialize() override;
+
+private:
+
+
+  double m_minPt = 10;
+  double m_jetRadius = 0.4;
+  std::string m_jetAlg = "antikt";
+  int m_isExclusive = 0;
+
+  k4::recCalo::ClusterJet* clusterer;
+
+
+};
+#endif /* RECCALORIMETER_CREATECALOJET_H */
diff --git a/RecCalorimeter/src/components/CreateTruthJet.cpp b/RecCalorimeter/src/components/CreateTruthJet.cpp
@@ -0,0 +1,67 @@
+#include "CreateTruthJet.h"
+
+// std
+#include <vector>
+#include <math.h>
+
+DECLARE_COMPONENT(CreateTruthJet)
+
+CreateTruthJet::CreateTruthJet(const std::string& name, ISvcLocator* svcLoc) : Transformer(name, svcLoc,
+                    KeyValue("InputCollection", "MCParticles"),
+                    KeyValue("OutputCollection", "TruthJets")) {
+  declareProperty("JetAlg", m_jetAlg, "Name of jet clustering algorithm");
+  declareProperty("JetRadius", m_jetRadius, "Jet clustering radius");
+  declareProperty("MinPt", m_minPt, "Minimum pT for saved jets");
+  declareProperty("isExclusiveClustering", m_isExclusive, "1 if exclusive, 0 if inclusive");
+  declareProperty("outputAssociation", m_outputAssocCollection, "TruthJetParticleAssociation");
+}
+
+
+
+StatusCode CreateTruthJet::initialize() {
+  clusterer = new k4::recCalo::ClusterJet(m_jetAlg, m_jetRadius, m_isExclusive, m_minPt);
+  return clusterer->initialize();
+
+}
+
+edm4hep::ReconstructedParticleCollection  CreateTruthJet::operator()(const edm4hep::MCParticleCollection& input) const{
+  std::vector<fastjet::PseudoJet> clustersPJ;
+  auto& assoc    = *(m_outputAssocCollection.createAndPut());
+  int i=0;
+  for(auto particle: input){
+    fastjet::PseudoJet clusterPJ(particle.getMomentum().x, particle.getMomentum().y, particle.getMomentum().z, particle.getEnergy());
+    clusterPJ.set_user_info(new k4::recCalo::ClusterInfo(i));
+    clustersPJ.push_back(clusterPJ);
+    i++;
+  }
+
+  std::vector<fastjet::PseudoJet> inclusiveJets = clusterer->cluster(clustersPJ);
+
+
+  edm4hep::ReconstructedParticleCollection edmJets = edm4hep::ReconstructedParticleCollection();
+
+  for(auto cjet : inclusiveJets){
+    edm4hep::MutableReconstructedParticle jet;
+    jet.setMomentum(edm4hep::Vector3f(cjet.px(), cjet.py(), cjet.pz()));
+    jet.setEnergy(cjet.e());
+    jet.setMass(cjet.m());
+
+    std::vector<fastjet::PseudoJet> constits = cjet.constituents();
+
+    for(auto constit : constits){
+      int index = constit.user_info<k4::recCalo::ClusterInfo>().index();
+
+      edm4hep::MutableMCRecoParticleAssociation association;
+      association.setRec(jet);
+      association.setSim((input)[index]);
+      assoc.push_back(association);
+    }
+
+    edmJets.push_back(jet);
+  }
+
+  return edmJets;
+}
+
+
+