JetsAndHF/eFinder/find_ScatElectron.C

0001 #include <fstream>
0002 #include <map>
0003 #include "TROOT.h"
0004 #include "TClass.h"
0005 #include "TFile.h"
0006 #include "TLine.h"
0007 #include "TTree.h"
0008 #include "TBrowser.h"
0009 #include "TH2.h"
0010 #include "TProfile.h"
0011 #include "TRandom.h"
0012 #include "TMath.h"
0013 #include <vector>
0014 #include "TVector3.h"
0015 #include "TLorentzVector.h"
0016 #include "TTreeReader.h"
0017 #include "TTreeReaderValue.h"
0018 #include "TTreeReaderArray.h"
0019 #include "TLatex.h"
0020 #include <string>
0021
0022 using namespace std;
0023
0024 void find_ScatElectron(TString listname = "file.list")
0025 {
0026   TChain *chain = new TChain("events");
0027   TTree* meta = NULL;
0028
0029   int nfiles = 0;
0030   char filename[512];
0031   ifstream *inputstream = new ifstream;
0032   inputstream->open(listname.Data());
0033   if(!inputstream)
0034     {
0035       printf("[e] Cannot open file list: %s\n", listname.Data());
0036     }
0037   while(inputstream->good())
0038     {
0039       inputstream->getline(filename, 512);
0040       if(inputstream->good())
0041     {
0042       TFile *ftmp = TFile::Open(filename, "read");
0043       if(!ftmp||!(ftmp->IsOpen())||!(ftmp->GetNkeys()))
0044         {
0045           printf("[e] Could you open file: %s\n", filename);
0046         }
0047       else
0048         {
0049           cout<<"[i] Add "<<nfiles<<"th file: "<<filename<<endl;
0050           chain->Add(filename);
0051           if(nfiles==0) meta  = (TTree*)ftmp->Get("podio_metadata");
0052           nfiles++;
0053         }
0054     }
0055     }
0056   inputstream->close();
0057   printf("[i] Read in %d files with %lld events in total\n", nfiles, chain->GetEntries());
0058
0059   TTreeReader treereader(chain);
0060
0061   // scattered electron
0062   TTreeReaderArray<int> mc_ScatElec = {treereader, "MCScatteredElectrons_objIdx.index"};
0063
0064   // Reconstructed particles
0065   TTreeReaderArray<float> rcMomPx = {treereader, "ReconstructedParticles.momentum.x"};
0066   TTreeReaderArray<float> rcMomPy = {treereader, "ReconstructedParticles.momentum.y"};
0067   TTreeReaderArray<float> rcMomPz = {treereader, "ReconstructedParticles.momentum.z"};
0068   TTreeReaderArray<float> rcCharge = {treereader, "ReconstructedParticles.charge"};
0069   TTreeReaderArray<int> rcPdg = {treereader, "ReconstructedParticles.PDG"};
0070
0071   TTreeReaderArray<int> rc_index = {treereader, "_ReconstructedParticleAssociations_rec.index"};
0072   TTreeReaderArray<int> mc_index = {treereader, "_ReconstructedParticleAssociations_sim.index"};
0073
0074   TTreeReaderArray<unsigned int> clusters_begin = {treereader, "ReconstructedParticles.clusters_begin"};
0075   TTreeReaderArray<unsigned int> clusters_end = {treereader, "ReconstructedParticles.clusters_end"};
0076   TTreeReaderArray<int> clusters_index = {treereader, "_ReconstructedParticles_clusters.index"};
0077   TTreeReaderArray<unsigned int> clusters_collectionID = {treereader, "_ReconstructedParticles_clusters.collectionID"};
0078
0079   TTreeReaderArray<unsigned int> tracks_begin = {treereader, "ReconstructedParticles.tracks_begin"};
0080   TTreeReaderArray<unsigned int> tracks_end = {treereader, "ReconstructedParticles.tracks_end"};
0081   TTreeReaderArray<int> tracks_index = {treereader, "_ReconstructedParticles_tracks.index"};
0082   TTreeReaderArray<unsigned int> tracks_collectionID = {treereader, "_ReconstructedParticles_tracks.collectionID"};
0083
0084   // cluster collections
0085   map<string, int> nameToIndex;
0086   const int nClusterCollections = 3;
0087   string cluster_names[nClusterCollections] = {"EcalBarrelScFiClusters", "EcalEndcapNClusters", "EcalEndcapPClusters"};
0088   TTreeReaderArray<float>* clusterE[nClusterCollections], *clusterX[nClusterCollections], *clusterY[nClusterCollections], *clusterZ[nClusterCollections];
0089   for(int i=0; i<nClusterCollections; i++)
0090     {
0091       nameToIndex[cluster_names[i]] = i;
0092       clusterE[i] = new TTreeReaderArray<float>(treereader, Form("%s.energy", cluster_names[i].c_str()));
0093       clusterX[i] = new TTreeReaderArray<float>(treereader, Form("%s.position.x", cluster_names[i].c_str()));
0094       clusterY[i] = new TTreeReaderArray<float>(treereader, Form("%s.position.y", cluster_names[i].c_str()));
0095       clusterZ[i] = new TTreeReaderArray<float>(treereader, Form("%s.position.z", cluster_names[i].c_str()));
0096     }
0097
0098   // Get the relevant leafs from podio_metadata tree
0099   TLeaf* idLeaf = meta->FindLeaf("events___CollectionTypeInfo.collectionID");
0100   TLeafC* nameLeafBase = (TLeafC*)meta->FindLeaf("events___CollectionTypeInfo.name");
0101   meta->GetEntry(0);
0102   Int_t nCollections = idLeaf->GetLen();
0103   const char* tmpFile = "podio_metadata.txt";
0104   {
0105     TRedirectOutputGuard guard(tmpFile, "w");
0106     nameLeafBase->PrintValue(nCollections);
0107   }
0108   ifstream infile(tmpFile);
0109   string out;
0110   getline(infile, out);
0111   infile.close();
0112
0113   // parse the out string to save into a map for lookup
0114   string prefix = "events___CollectionTypeInfo.name =";
0115   if (out.rfind(prefix, 0) == 0)
0116     {
0117       out = out.substr(prefix.size());
0118     }
0119   vector<string> names;
0120   stringstream ss(out);
0121   string item;
0122
0123   // split by comma
0124   while (getline(ss, item, ','))
0125     {
0126       // trim leading space
0127       if (!item.empty() && item[0] == ' ')
0128     item.erase(0, 1);
0129       names.push_back(item);
0130     }
0131
0132   if (names.size() != nCollections)
0133     {
0134       cerr << "Mismatch: names =" << names.size() << " vs IDs =" << nCollections << endl;
0135       return;
0136     }
0137
0138   // build map
0139   unordered_map<UInt_t, string> idToName;
0140   for (size_t i = 0; i < names.size(); ++i)
0141     {
0142       UInt_t cid = static_cast<UInt_t>(idLeaf->GetValue(i));
0143       idToName[cid] = names[i];
0144     }
0145
0146   // Cuts
0147   const double mIsoR = 0.4;
0148   const double mEoP_min = 0.8;
0149   const double mEoP_max = 1.2;
0150   const double mIsoE = 0.9;
0151
0152   // find scattered electron
0153   int nevents = 0;
0154   while(treereader.Next())
0155     {
0156       //if(nevents>11) return;
0157       if(nevents%100==0) printf("\n[i] New event %d\n",nevents);
0158       printf("\n[i] New event %d\n",nevents);
0159       nevents++;
0160
0161       // loop over reconstructed particles
0162       int scat_elec_index = -1;
0163       double scat_elec_pt = -1;
0164       for(int rc_index=0; rc_index<rcMomPx.GetSize(); rc_index++)
0165     {
0166       //cout << rcMomPx[rc_index] << "  " << chain->FindLeaf("ReconstructedParticles.momentum.x")->GetValue(rc_index) << endl;
0167       int nclusters = clusters_end[rc_index] - clusters_begin[rc_index];
0168       int ntracks = tracks_end[rc_index] - tracks_begin[rc_index];
0169
0170       // at least one cluster and one track
0171       if(nclusters == 0 || ntracks == 0) continue;
0172
0173       // Negative charge
0174       if(rcCharge[rc_index] >= 0) continue;
0175
0176       // Calculate associate cluster energy
0177       double rcpart_sum_cluster_E = 0;
0178       double rcpart_isolation_E = 0;
0179       double lead_eta = -999;
0180       double lead_phi = -999;
0181       double lead_E = 0;
0182       for(int ic = clusters_begin[rc_index]; ic < clusters_end[rc_index]; ic++)
0183         {
0184           int cluster_index = clusters_index[ic];
0185           int collectionID = clusters_collectionID[ic];
0186           string collection_name = idToName[collectionID];
0187           if(!nameToIndex.count(collection_name))
0188         {
0189           cout << "[e] Error: unknown collection name: " << collection_name.c_str() << endl;
0190           cout << "[e] Please update the map" << endl;
0191           return;
0192         }
0193           int collection_index = nameToIndex[collection_name];
0194           double energy = clusterE[collection_index]->At(cluster_index);
0195           rcpart_sum_cluster_E += energy;
0196           if(energy>lead_E)
0197         {
0198           lead_E = energy;
0199           TVector3 pos(clusterX[collection_index]->At(cluster_index),
0200                    clusterY[collection_index]->At(cluster_index),
0201                    clusterZ[collection_index]->At(cluster_index));
0202           lead_eta = pos.Eta();
0203           lead_phi = pos.Phi();
0204         }
0205         }
0206
0207       // loop again to get isolation energy
0208       for(int j=0; j<rcMomPx.GetSize(); j++)
0209         {
0210           for(int ic = clusters_begin[j]; ic < clusters_end[j]; ic++)
0211         {
0212           int cluster_index = clusters_index[ic];
0213           int collectionID = clusters_collectionID[ic];
0214           string collection_name = idToName[collectionID];
0215           if(!nameToIndex.count(collection_name))
0216             {
0217               cout << "[e] Error: unknown collection name: " << collection_name.c_str() << endl;
0218               cout << "[e] Please update the map" << endl;
0219               return;
0220             }
0221           int collection_index = nameToIndex[collection_name];
0222           TVector3 pos(clusterX[collection_index]->At(cluster_index),
0223                    clusterY[collection_index]->At(cluster_index),
0224                    clusterZ[collection_index]->At(cluster_index));
0225
0226           double d_eta = pos.Eta() - lead_eta;
0227           double d_phi = pos.Phi() - lead_phi;
0228           if (d_phi > TMath::Pi()) d_phi-=2*TMath::Pi();
0229           if (d_phi < -TMath::Pi()) d_phi+=2*TMath::Pi();
0230
0231           double dR = std::sqrt(std::pow(d_eta, 2) + std::pow(d_phi, 2));
0232           // Check if the cluster is within the isolation cone
0233           if (dR < mIsoR)
0234             {
0235               rcpart_isolation_E += clusterE[collection_index]->At(cluster_index);
0236             }
0237         }
0238         }
0239
0240       //printf("[i] rcpart_sum_cluster_E = %f, rcpart_isolation_E = %f\n", rcpart_sum_cluster_E,  rcpart_isolation_E);
0241
0242       // Apply scattered electron ID cuts
0243       TVector3 part_mom(rcMomPx[rc_index], rcMomPy[rc_index], rcMomPz[rc_index]);
0244       double recon_EoP = rcpart_sum_cluster_E / part_mom.Mag();
0245       double recon_isoE = rcpart_sum_cluster_E / rcpart_isolation_E;
0246       if(recon_EoP < mEoP_min || recon_EoP > mEoP_max) continue;
0247       if(recon_isoE < mIsoE) continue;
0248
0249       // in case of multiple candidate, select the one with largest pt
0250       if(part_mom.Pt()>scat_elec_pt)
0251         {
0252           scat_elec_pt = part_mom.Pt();
0253           scat_elec_index = rc_index;
0254         }
0255     }
0256
0257       // Sanity check
0258       if(scat_elec_index>=0)
0259     {
0260       // find corresponding MC particle
0261       int scat_elec_mc_index = -1;
0262       for(unsigned int i=0; i<rc_index.GetSize(); i++)
0263         {
0264           if(rc_index[i] == scat_elec_index)
0265         {
0266           scat_elec_mc_index = mc_index[i];
0267           break;
0268         }
0269         }
0270
0271       // true MC electron index
0272       int scat_elec_mc_truth = mc_ScatElec[0];
0273
0274       printf("[i] Found a scattered electron candidate: index = %d, pt = %f, mc index = %d, truth = %d\n", scat_elec_index, scat_elec_pt, scat_elec_mc_index, scat_elec_mc_truth);
0275     }
0276     }
0277 }