delphes_EIC/SimpleAnalysis/ElectronPIDModule.cc

0001 #include "ElectronPIDModule.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 #include <algorithm>
0011
0012 ElectronPIDModule::ElectronPIDModule(ExRootTreeReader* data)
0013   : Module(data)
0014 {
0015
0016 }
0017
0018 ElectronPIDModule::~ElectronPIDModule()
0019 {
0020
0021 }
0022
0023
0024
0025 bool ElectronPIDModule::execute(std::map<std::string, std::any>* DataStore)
0026 {
0027   auto data = getData();
0028
0029   // If event contains at least 1 jet
0030   // std::vector<Jet*> all_jets;
0031   // for (int ijet = 0; ijet < getJets()->GetEntries(); ijet++)
0032   //   {
0033   //     // Take first jet
0034   //     Jet *jet = (Jet*) getJets()->At(ijet);
0035   //     all_jets.push_back(jet);
0036   //   }
0037
0038   // std::vector<Jet*> fiducial_jets = SelectorFcn<Jet>(all_jets, [](Jet* j){ return (TMath::Abs(j->Eta) < 3.0 && j->PT > 5.0); });
0039   // std::vector<Jet*> charmJets = SelectorFcn<Jet>(fiducial_jets, [](Jet* j){ return (j->Flavor == 4); });
0040
0041   std::vector<Track*> all_electrons;
0042
0043   // If the DataStore contains already a list of tracks to be used for PID assignment,
0044   // use that. If not, use the getEFlowTracks() method to get the general list of all tracks.
0045   // TracksForPID is special - it contains tracks NOT already used by another PID algorithm,
0046   // to avoid using the same track (pion) twice in two categories.
0047
0048   float e_efficiency = 0.90;
0049   float epi_separation = 2.4; //sigma
0050   //float e_efficiency = 1.00;
0051   //float epi_separation = 100.0; //sigma
0052
0053   if (DataStore->find("TracksForPID") != DataStore->end()) {
0054     for (auto track : std::any_cast<std::vector<Track*>>((*DataStore)["TracksForPID"])) {
0055       if (ElectronPID(track, e_efficiency, epi_separation))
0056     all_electrons.push_back(track);
0057     }
0058
0059
0060   } else {
0061     for (int itrk = 0; itrk < getEFlowTracks()->GetEntries(); itrk++)
0062       {
0063     Track* track = (Track*) getEFlowTracks()->At(itrk);
0064     if (ElectronPID(track, e_efficiency, epi_separation))
0065       all_electrons.push_back(track);
0066       }
0067
0068     std::vector<Track*> tracks_for_PID;
0069     for (int itrk = 0; itrk < getEFlowTracks()->GetEntries(); itrk++)
0070       {
0071     Track* track = (Track*) getEFlowTracks()->At(itrk);
0072     tracks_for_PID.push_back(track);
0073       }
0074     (*DataStore)["TracksForPID"] = tracks_for_PID;
0075   }
0076
0077   //auto reconstructed_electrons = all_electrons;
0078   std::vector<Track*> reconstructed_electrons = SelectorFcn<Track>(all_electrons,
0079                                    [](Track* t){ return (t->PT >= 0.1); });
0080
0081   std::vector<Track*> tracks_for_PID = std::any_cast<std::vector<Track*>>((*DataStore)["TracksForPID"]);
0082   for (auto electron : reconstructed_electrons) {
0083     tracks_for_PID.erase(std::find(tracks_for_PID.begin(), tracks_for_PID.end(), electron));
0084   }
0085   (*DataStore)["TracksForPID"] = tracks_for_PID;
0086
0087
0088   // store the electrons
0089   (*DataStore)["Electrons"] = reconstructed_electrons;
0090
0091
0092   return true;
0093 }
0094
0095
0096 bool ElectronPIDModule::ElectronPID(Track* track, float eIDprob, float separation)
0097 {
0098   bool TrackIsElectron = false;
0099
0100
0101   // Apply a basic parameterized electron PID to tracks. If the track is really a
0102   // electron from truth information, apply a flat ID probability. If it's a pion,
0103   // use the separation (in Gaussian sigma) to determine if it's mis-identified.
0104
0105   if (track->PT < 0.1)
0106     return TrackIsElectron;
0107
0108   int track_truth = track->PID;
0109
0110   if (TMath::Abs(track_truth) == 11) {
0111     // true charged electron
0112     if (gRandom->Uniform(0, 1) <= eIDprob) {
0113       TrackIsElectron = true;
0114     } else {
0115       TrackIsElectron = false;
0116     }
0117
0118
0119
0120   } else if (TMath::Abs(track_truth) == 211) {
0121     // true charged pion
0122     if (gRandom->Uniform(0,1) <= ROOT::Math::gaussian_pdf(separation)) {
0123       TrackIsElectron = true;
0124     } else {
0125       TrackIsElectron = false;
0126     }
0127
0128   } else {
0129     // ignore ALL other species for now
0130     TrackIsElectron = false;
0131   }
0132
0133
0134   return TrackIsElectron;
0135 }
0136
0137
0138
0139