Adding ghost matching jet flavour defn

analyzers/dataframe/src/JetTaggingUtils.cc

@@ -63,6 +63,212 @@ ROOT::VecOps::RVec<int> get_flavour(ROOT::VecOps::RVec<fastjet::PseudoJet> in,
   return result;
 }
 
+
+get_ghostFlavour::get_ghostFlavour(const int & arg_algo, const float & arg_radius, const int & arg_exclusive, const float & arg_cut, const int & arg_sorted, const int & arg_recombination,
+				   const float & arg_add1, const float & arg_add2)
+{m_algo = arg_algo; m_radius = arg_radius; m_exclusive = arg_exclusive; m_cut = arg_cut; m_sorted = arg_sorted; m_recombination = arg_recombination; m_add1 = arg_add1; m_add2 = arg_add2;}
+
+ghostFlavour JetTaggingUtils::get_ghostFlavour::operator() (const ROOT::VecOps::RVec<edm4hep::MCParticleData> & Particle,
+							    const ROOT::VecOps::RVec<int> & ind,
+							    std::vector<fastjet::PseudoJet> & pseudoJets,
+							    const int & partonFlag) {
+
+
+  ghostFlavour result;
+
+  unsigned int index = pseudoJets.size();
+  ROOT::VecOps::RVec<int> pdg(pseudoJets.size(),0);
+  ROOT::VecOps::RVec<int> ghostStatus(pseudoJets.size(),0);
+  ROOT::VecOps::RVec<int> MCindex(pseudoJets.size(),-1);
+
+
+  ROOT::VecOps::RVec<int> partonStatus = {20, 30};
+
+  // In loop below ghosts are selected from MCParticle collection
+  for (size_t i = 0; i < Particle.size(); ++i) {
+    bool isGhost = false;
+
+    if(!partonFlag){
+      // Ghost partons as any partons that do not have partons as daughters
+      if (std::abs(Particle[i].PDG)<=5||Particle[i].PDG==21) {
+        isGhost = true;
+        auto daughters = MCParticle::get_list_of_particles_from_decay(i, Particle, ind);
+        for(auto& daughter_index : daughters){
+          if (std::abs(Particle[daughter_index].PDG)<=5||Particle[daughter_index].PDG==21) {
+              isGhost = false;
+              break;
+            }
+        }
+        if (isGhost) ghostStatus.push_back(1);
+      }
+    }
+    else{
+      // Ghost partons are selected via chosen Pythia8 status codes
+      if ((Particle[i].generatorStatus<partonStatus[1]) && (Particle[i].generatorStatus>partonStatus[0])) {
+        isGhost = true;
+        ghostStatus.push_back(1);
+      }
+    }
+
+
+    // Ghost hadrons are selected as b/c hadrons that do not have b/c hadrons as daughters
+    if ((std::abs(int((Particle[i].PDG/100))%10)==5)||(std::abs(int((Particle[i].PDG/1000))%10)==5)){
+      isGhost = true;
+      auto daughters = MCParticle::get_list_of_particles_from_decay(i, Particle, ind);
+      for(auto& daughter_index : daughters){
+          if((std::abs(int((Particle[daughter_index].PDG/100))%10)==5)||(std::abs(int((Particle[daughter_index].PDG/1000))%10)==5)){
+            isGhost = false;
+            break;
+          }
+      }
+      if (isGhost) ghostStatus.push_back(2);
+
+    }
+    else if ((std::abs(int((Particle[i].PDG/100))%10)==4)||(std::abs(int((Particle[i].PDG/1000))%10)==4)){
+      isGhost = true;
+      auto daughters = MCParticle::get_list_of_particles_from_decay(i, Particle, ind);
+      for(auto& daughter_index : daughters){
+          if((std::abs(int((Particle[daughter_index].PDG/100))%10)==4)||(std::abs(int((Particle[daughter_index].PDG/1000))%10)==4)){
+            isGhost = false;
+            break;
+          }
+      }
+      if (isGhost) ghostStatus.push_back(2);
+    }
+
+    // Ghosts 4-mom is scaled by 10^-18  
+    if (isGhost){
+      pdg.push_back(Particle[i].PDG);
+      MCindex.push_back(i);
+      // the double conversion here is verbose but for precision if I recall...
+      double px = Particle[i].momentum.x;//*pow(10, -1);
+      double py = Particle[i].momentum.y;
+      double pz = Particle[i].momentum.z;
+      double m = Particle[i].mass;
+      double  E = sqrt(px*px + py*py + pz*pz + m*m);
+      pseudoJets.emplace_back(px*pow(10, -18), py*pow(10, -18), pz*pow(10, -18), E*pow(10, -18));
+      pseudoJets.back().set_user_index(index);
+      ++index;
+    }
+  }
+
+  result.ghostStatus = ghostStatus;
+  result.MCindex = MCindex;
+
+
+  // Jet clustering algorithm is run according to user choice m_algo
+  JetClusteringUtils::FCCAnalysesJet FCCAnalysesGhostJets;
+
+  if (m_algo == 0) FCCAnalysesGhostJets = JetClustering::clustering_kt(m_radius, m_exclusive, m_cut, m_sorted, m_recombination)(pseudoJets);
+  else if (m_algo == 1) FCCAnalysesGhostJets = JetClustering::clustering_antikt(m_radius, m_exclusive, m_cut, m_sorted, m_recombination)(pseudoJets);
+  else if (m_algo == 2) FCCAnalysesGhostJets = JetClustering::clustering_cambridge(m_radius, m_exclusive, m_cut, m_sorted, m_recombination)(pseudoJets);
+  else if (m_algo == 3) FCCAnalysesGhostJets = JetClustering::clustering_ee_kt(m_exclusive, m_cut, m_sorted, m_recombination)(pseudoJets);
+  else if (m_algo == 4) FCCAnalysesGhostJets = JetClustering::clustering_ee_genkt(m_radius, m_exclusive, m_cut, m_sorted, m_recombination, m_add1)(pseudoJets);
+  else if (m_algo == 5) FCCAnalysesGhostJets = JetClustering::clustering_genkt(m_radius, m_exclusive, m_cut, m_sorted, m_recombination, m_add1)(pseudoJets);
+  else if (m_algo == 6) FCCAnalysesGhostJets = JetClustering::clustering_valencia(m_radius, m_exclusive, m_cut, m_sorted, m_recombination, m_add1, m_add2)(pseudoJets);
+  else if (m_algo == 7) FCCAnalysesGhostJets = JetClustering::clustering_jade(m_radius, m_exclusive, m_cut, m_sorted, m_recombination)(pseudoJets);
+  else return result;
+
+
+  result.jets = FCCAnalysesGhostJets;
+
+  // Jet constituents and pseudojets are read from resulting jet struct
+  auto ghostJets_ee_kt = JetClusteringUtils::get_pseudoJets(FCCAnalysesGhostJets);
+
+  auto jetconstituents = JetClusteringUtils::get_constituents(FCCAnalysesGhostJets);
+
+
+
+
+
+  // Flav vector is defined before jets are checked for clustered ghosts
+  std::vector<std::vector<int>> flav_vector;
+
+
+  std::vector<int> partonFlavs;
+  std::vector<int> hadronFlavs;
+  for (auto& consti_index : jetconstituents) {
+
+  int partonFlav = 0;
+  float partonMom2 = 0;
+  float partonMom2_b = 0;
+  float partonMom2_c = 0;
+  int hadronFlav = 0;
+  float hadronMom2_b = 0;
+  float hadronMom2_c = 0;
+
+    for (auto& i : consti_index) {
+
+      //Parton-flav loop 
+      if (ghostStatus[i]==1) {
+          if (std::abs(pdg[i])==5){
+              if (pseudoJets[i].modp2()>partonMom2_b){
+                  partonFlav = pdg[i];
+                  partonMom2_b = pseudoJets[i].modp2();
+              }
+          }
+          else if ((std::abs(pdg[i])==4) && (std::abs(partonFlav)<5)) {
+              if (pseudoJets[i].modp2()>partonMom2_c){
+                  partonFlav = pdg[i];
+                  partonMom2_c = pseudoJets[i].modp2();
+              }
+          }
+          else if (((std::abs(pdg[i])==3) || (std::abs(pdg[i])==2) || (std::abs(pdg[i])==1) || (std::abs(pdg[i])==21)) && ((std::abs(partonFlav)<4) || (std::abs(partonFlav)==21))) {
+              if (pseudoJets[i].modp2()>partonMom2){
+                  partonFlav = pdg[i];
+                  partonMom2 = pseudoJets[i].modp2();
+              }
+          }
+      }
+
+      // Hadron-flav loop    
+      if (ghostStatus[i]==2) {
+          if ((std::abs(int((pdg[i]/100))%10)==5)||(std::abs(int((pdg[i]/1000))%10)==5)){
+              if (pseudoJets[i].modp2()>hadronMom2_b){
+                  hadronFlav = ((pdg[i]<0)-(pdg[i]>0))*5;
+                  hadronMom2_b = pseudoJets[i].modp2();
+              }
+          }
+          else if (((std::abs(int((pdg[i]/100))%10)==4)||(std::abs(int((pdg[i]/1000))%10)==4)) && (std::abs(hadronFlav)<5)) {
+              if (pseudoJets[i].modp2()>hadronMom2_c){
+                  hadronFlav = ((pdg[i]>0)-(pdg[i]<0))*4;
+                  hadronMom2_c = pseudoJets[i].modp2();
+              }
+          }
+      }
+    }
+    partonFlavs.push_back(partonFlav);
+    hadronFlavs.push_back(hadronFlav);
+  }
+
+  flav_vector.push_back(partonFlavs);
+  flav_vector.push_back(hadronFlavs);
+
+  result.flavour = flav_vector;
+  return result;
+
+}
+
+std::vector<std::vector<int>> JetTaggingUtils::get_flavour(ghostFlavour ghostStruct){
+  return ghostStruct.flavour;
+}
+
+JetClusteringUtils::FCCAnalysesJet JetTaggingUtils::get_jets(ghostFlavour ghostStruct){
+  return ghostStruct.jets;
+}
+
+
+ROOT::VecOps::RVec<int> JetTaggingUtils::get_ghostStatus(ghostFlavour ghostStruct){
+  return ghostStruct.ghostStatus;
+}
+
+ROOT::VecOps::RVec<int> JetTaggingUtils::get_MCindex(ghostFlavour ghostStruct){
+  return ghostStruct.MCindex;
+}
+
+
+
+
 ROOT::VecOps::RVec<int>
 get_btag(ROOT::VecOps::RVec<int> in,
                           float efficiency, float mistag_c,

examples/FCCee/ghostFlavour/analysis.py

@@ -0,0 +1,170 @@
+import sys
+import ROOT
+
+print ("Load cxx analyzers ... ",)
+ROOT.gSystem.Load("libedm4hep")
+ROOT.gSystem.Load("libpodio")
+ROOT.gSystem.Load("libFCCAnalyses")
+ROOT.gErrorIgnoreLevel = ROOT.kFatal
+_edm  = ROOT.edm4hep.ReconstructedParticleData()
+_pod  = ROOT.podio.ObjectID()
+_fcc  = ROOT.dummyLoader
+
+print ('edm4hep  ',_edm)
+print ('podio    ',_pod)
+print ('fccana   ',_fcc)
+
+class analysis():
+
+    #__________________________________________________________
+    def __init__(self, inputlist, outname, ncpu):
+        self.outname = outname
+        if ".root" not in outname:
+            self.outname+=".root"
+
+        ROOT.ROOT.EnableImplicitMT(ncpu)
+
+        self.df = ROOT.RDataFrame("events", inputlist)
+        print (" done")
+    #__________________________________________________________
+    def run(self):
+        df2 = (self.df
+        #df2 = (self.df.Range(0, 10)
+               #alias for dealing with # in python
+               .Alias("Particle0", "Particle#0.index") 
+               .Alias("Particle1", "Particle#1.index") 
+               .Alias("Jet3","Jet#3.index")
+
+               .Define("MC_px",  "MCParticle::get_px(Particle)")
+               .Define("MC_py",  "MCParticle::get_py(Particle)")
+               .Define("MC_pz",  "MCParticle::get_pz(Particle)")
+               .Define("MC_p",  "MCParticle::get_p(Particle)")
+               .Define("MC_e",  "MCParticle::get_e(Particle)")
+               .Define("MC_m",  "MCParticle::get_mass(Particle)")
+               .Define("MC_theta",  "MCParticle::get_theta(Particle)")
+               .Define("MC_pdg", "MCParticle::get_pdg(Particle)")
+               .Define("MC_status", "MCParticle::get_genStatus(Particle)")
+
+
+               #define the RP px, py, pz and e
+               .Define("RP_px",          "ReconstructedParticle::get_px(ReconstructedParticles)")
+               .Define("RP_py",          "ReconstructedParticle::get_py(ReconstructedParticles)")
+               .Define("RP_pz",          "ReconstructedParticle::get_pz(ReconstructedParticles)")               
+               .Define("RP_m",           "ReconstructedParticle::get_mass(ReconstructedParticles)")
+
+
+
+
+               ################
+               #Jet Clustering#
+               ################
+
+               #build pseudo jets with the RP, using the interface that takes px,py,pz,m for better
+               #handling of rounding errors
+               .Define("pseudo_jets",    "JetClusteringUtils::set_pseudoJets_xyzm(RP_px, RP_py, RP_pz, RP_m)")
+
+               #EE-KT ALGORITHM
+               #run jet clustering with all MC particles. ee_kt_algorithm, exclusive clustering, exactly 2 jets, E-scheme
+               .Define("FCCAnalysesJets_ee_kt", "JetClustering::clustering_ee_kt(2, 2, 1, 0)(pseudo_jets)")
+
+               #get the jets out of the structure
+               .Define("jets_ee_kt",            "JetClusteringUtils::get_pseudoJets(FCCAnalysesJets_ee_kt)")
+
+               #get the jet constituents out of the structure
+               .Define("jetconstituents_ee_kt", "JetClusteringUtils::get_constituents(FCCAnalysesJets_ee_kt)")
+
+               #get some jet variables
+               .Define("jets_ee_kt_e",          "JetClusteringUtils::get_e(jets_ee_kt)")
+               .Define("jets_ee_kt_px",         "JetClusteringUtils::get_px(jets_ee_kt)")
+               .Define("jets_ee_kt_py",         "JetClusteringUtils::get_py(jets_ee_kt)")
+               .Define("jets_ee_kt_pz",         "JetClusteringUtils::get_pz(jets_ee_kt)")
+               .Define("jets_ee_kt_flavour",    "JetTaggingUtils::get_flavour(jets_ee_kt, Particle)")
+
+               ###############
+               #Ghost Flavour#
+               ###############
+
+               #pseudo jets defined above, along with Particle, Particle 1 collections are passed to get_ghostflavour
+               .Define("ghostFlavour",        "JetTaggingUtils::get_ghostFlavour(3, 0, 2, 2, 1, 0)(Particle, Particle1, pseudo_jets, 0)")
+
+               #the flavour vector can be obtained from the struct by using get_flavour
+               .Define("ghostFlavour_flav", "JetTaggingUtils::get_flavour(ghostFlavour)")
+
+               #an FCCAnalysesJet struct can be obtain by using get_jets 
+               .Define("ghostFlavour_jets", "JetTaggingUtils::get_jets(ghostFlavour)")
+               #the pseudojets can then be obtained (as above) using get_pseudojets
+               .Define("ghostJets", "JetClusteringUtils::get_pseudoJets(ghostFlavour_jets)")
+               #with the kinematics of the pseudojets accessible via their own functions
+               .Define("ghostJets_e",          "JetClusteringUtils::get_e(ghostJets)")
+               .Define("ghostJets_px",         "JetClusteringUtils::get_px(ghostJets)")
+               .Define("ghostJets_py",         "JetClusteringUtils::get_py(ghostJets)")
+               .Define("ghostJets_pz",         "JetClusteringUtils::get_pz(ghostJets)")
+
+               #the ghostStatus and MCindex vector can likewise be obtained
+               .Define("ghostFlavour_MCindex", "JetTaggingUtils::get_MCindex(ghostFlavour)")
+               .Define("ghostFlavour_ghostStatus", "JetTaggingUtils::get_ghostStatus(ghostFlavour)")
+
+
+
+
+        )
+
+
+
+
+        # select branches for output file
+        branchList = ROOT.vector('string')()
+        for branchName in [
+                "MC_px",
+                "MC_py",
+                "MC_pz",
+                "MC_p",
+                "MC_e",
+                "MC_theta",
+                "MC_pdg",
+                "MC_status",
+
+                "jets_ee_kt_e",
+                "jets_ee_kt_px",
+                "jets_ee_kt_py",
+                "jets_ee_kt_pz",
+                "jets_ee_kt_flavour",
+                "jetconstituents_ee_kt",
+
+                "ghostFlavour_flav",
+                "ghostFlavour_MCindex",
+                "ghostFlavour_ghostStatus",
+
+                "ghostJets_e",
+                "ghostJets_px",
+                "ghostJets_py",
+                "ghostJets_pz",
+
+
+                ]:
+            branchList.push_back(branchName)
+        df2.Snapshot("events", self.outname, branchList)
+
+# example call for standalone file
+# python examples/FCCee/top/hadronic/analysis.py /eos/experiment/fcc/ee/generation/DelphesEvents/fcc_tmp/p8_ee_tt_fullhad_ecm365/events_196309147.root 
+if __name__ == "__main__":
+
+    if len(sys.argv)==1:
+        print ("usage:")
+        print ("python ",sys.argv[0]," file.root")
+        sys.exit(3)
+    infile = sys.argv[1]
+    outDir = sys.argv[0].replace(sys.argv[0].split('/')[0],'outputs/FCCee').replace('analysis.py','')+'/'
+    import os
+    os.system("mkdir -p {}".format(outDir))
+    outfile = outDir+infile.split('/')[-1]
+    ncpus = 2
+    analysis = analysis(infile, outfile, ncpus)
+    analysis.run()
+
+    tf = ROOT.TFile(infile)
+    entries = tf.events.GetEntries()
+    p = ROOT.TParameter(int)( "eventsProcessed", entries)
+    outf=ROOT.TFile(outfile,"UPDATE")
+    p.Write()
+