delphes_EIC/SimpleAnalysis/KaonPIDModule.cc

0001 #include "KaonPIDModule.h"
0002
0003 #include "TClonesArray.h"
0004 #include "TRandom3.h"
0005 #include "Math/PdfFuncMathCore.h"
0006
0007 #include "AnalysisFunctions.cc"
0008
0009 #include <iostream>
0010
0011 KaonPIDModule::KaonPIDModule(ExRootTreeReader* data)
0012   : Module(data)
0013 {
0014
0015 }
0016
0017 KaonPIDModule::~KaonPIDModule()
0018 {
0019
0020 }
0021
0022
0023
0024 bool KaonPIDModule::execute(std::map<std::string, std::any>* DataStore)
0025 {
0026   auto data = getData();
0027
0028   // If event contains at least 1 jet
0029   // std::vector<Jet*> all_jets;
0030   // for (int ijet = 0; ijet < getJets()->GetEntries(); ijet++)
0031   //   {
0032   //     // Take first jet
0033   //     Jet *jet = (Jet*) getJets()->At(ijet);
0034   //     all_jets.push_back(jet);
0035   //   }
0036
0037   // std::vector<Jet*> fiducial_jets = SelectorFcn<Jet>(all_jets, [](Jet* j){ return (TMath::Abs(j->Eta) < 3.0 && j->PT > 5.0); });
0038   // std::vector<Jet*> charmJets = SelectorFcn<Jet>(fiducial_jets, [](Jet* j){ return (j->Flavor == 4); });
0039
0040   std::vector<Track*> all_kaons;
0041
0042
0043   if (DataStore->find("TracksForPID") != DataStore->end()) {
0044     for (auto track : std::any_cast<std::vector<Track*>>((*DataStore)["TracksForPID"])) {
0045       if (KaonPID(track, 0.90, 3.0))
0046     all_kaons.push_back(track);
0047     }
0048
0049
0050   } else {
0051     for (int itrk = 0; itrk < getEFlowTracks()->GetEntries(); itrk++)
0052       {
0053     Track* track = (Track*) getEFlowTracks()->At(itrk);
0054     if (KaonPID(track, 0.90, 3.0))
0055       all_kaons.push_back(track);
0056       }
0057
0058     std::vector<Track*> tracks_for_PID;
0059     for (int itrk = 0; itrk < getEFlowTracks()->GetEntries(); itrk++)
0060       {
0061     Track* track = (Track*) getEFlowTracks()->At(itrk);
0062     tracks_for_PID.push_back(track);
0063       }
0064     (*DataStore)["TracksForPID"] = tracks_for_PID;
0065   }
0066
0067   //auto reconstructed_kaons = all_kaons;
0068    std::vector<Track*> reconstructed_kaons = SelectorFcn<Track>(all_kaons, [](Track* t){ return (t->PT >= 0.1); });
0069
0070    std::vector<Track*> tracks_for_PID = std::any_cast<std::vector<Track*>>((*DataStore)["TracksForPID"]);
0071    for (auto kaon : reconstructed_kaons) {
0072      tracks_for_PID.erase(std::find(tracks_for_PID.begin(), tracks_for_PID.end(), kaon));
0073    }
0074    (*DataStore)["TracksForPID"] = tracks_for_PID;
0075
0076
0077   (*DataStore)["Kaons"] = reconstructed_kaons;
0078
0079
0080   return true;
0081 }
0082
0083
0084 bool KaonPIDModule::KaonPID(Track* track, float kIDprob, float separation)
0085 {
0086   bool TrackIsKaon = false;
0087
0088
0089   // Apply a basic parameterized kaon PID to tracks. If the track is really a
0090   // kaon from truth information, apply a flat ID probability. If it's a pion,
0091   // use the separation (in Gaussian sigma) to determine if it's mis-identified.
0092
0093   if (track->PT < 0.1)
0094     return TrackIsKaon;
0095
0096   int track_truth = track->PID;
0097
0098   if (TMath::Abs(track_truth) == 321) {
0099     // true charged kaon
0100     if (gRandom->Uniform(0, 1) <= kIDprob) {
0101       TrackIsKaon = true;
0102     } else {
0103       TrackIsKaon = false;
0104     }
0105
0106
0107
0108   } else if (TMath::Abs(track_truth) == 211) {
0109     // true charged pion
0110     if (gRandom->Uniform(0,1) <= ROOT::Math::gaussian_pdf(separation)) {
0111       TrackIsKaon = true;
0112     } else {
0113       TrackIsKaon = false;
0114     }
0115
0116   } else {
0117     // ignore ALL other species for now
0118     TrackIsKaon = false;
0119   }
0120
0121
0122   return TrackIsKaon;
0123 }
0124
0125
0126
0127