Switching to using associations for truth jets

HEP-FCC · Jul 8, 2024 · 16f774d · 16f774d
1 parent e5a419e
commit 16f774d
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/RecCalorimeter/src/components/CreateTruthJet.cpp b/RecCalorimeter/src/components/CreateTruthJet.cpp
@@ -4,8 +4,6 @@
 #include <vector>
 #include <math.h>
 
-
-
 DECLARE_COMPONENT(CreateTruthJet)
 
 CreateTruthJet::CreateTruthJet(const std::string& name, ISvcLocator* svcLoc) : Transformer(name, svcLoc,
@@ -15,6 +13,7 @@ CreateTruthJet::CreateTruthJet(const std::string& name, ISvcLocator* svcLoc) : T
   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");
 }
 
 
@@ -27,6 +26,7 @@ StatusCode CreateTruthJet::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());
@@ -47,17 +47,14 @@ edm4hep::ReconstructedParticleCollection  CreateTruthJet::operator()(const edm4h
     jet.setMass(cjet.m());
 
     std::vector<fastjet::PseudoJet> constits = cjet.constituents();
+
     for(auto constit : constits){
-      edm4hep::MutableReconstructedParticle jetInput;
-      jetInput.setMomentum(edm4hep::Vector3f(constit.px(), constit.py(), constit.pz()));
-      jetInput.setEnergy(constit.e());
-      jetInput.setMass(constit.m());
-
       int index = constit.user_info<k4::recCalo::ClusterInfo>().index();
-      // This function is not available in all versions, so I am leaving it commented for now
-      jetInput.setPDG((input)[index].getPDG());
 
-      jet.addToParticles(jetInput);
+      edm4hep::MutableMCRecoParticleAssociation association;
+      association.setRec(jet);
+      association.setSim((input)[index]);
+      assoc.push_back(association);
     }
 
     edmJets.push_back(jet);

diff --git a/RecCalorimeter/src/components/CreateTruthJet.h b/RecCalorimeter/src/components/CreateTruthJet.h
@@ -11,6 +11,8 @@
 
 #include "edm4hep/ReconstructedParticleCollection.h"
 #include "edm4hep/MCParticleCollection.h"
+#include "edm4hep/MCRecoParticleAssociation.h"
+#include "edm4hep/MCRecoParticleAssociationCollection.h"
 
 #include "ClusterJet.h"
 
@@ -24,6 +26,7 @@
  *  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
+ *  outputAssociation: Name of the output association collection to link jets to their constituents
  *
 
  *  @author Jennifer Roloff
@@ -40,8 +43,11 @@ class CreateTruthJet : public Gaudi::Functional::Transformer <edm4hep::Reconstru
 
 private:
 
+
+  mutable DataHandle<edm4hep::MCRecoParticleAssociationCollection> m_outputAssocCollection{"MCRecoParticleAssociation",
+                                                                                           Gaudi::DataHandle::Writer, this};
 
-  double m_minPt = 10;
+  double m_minPt = 1;
   double m_jetRadius = 0.4;
   std::string m_jetAlg = "antikt";
   int m_isExclusive = 0;