snippets/ElectronID/InclusiveSkim.C

0001 // Data model headers
0002 #include "edm4eic/ReconstructedParticleCollection.h"
0003 #include "edm4hep/MCParticleCollection.h"
0004 #include "edm4hep/utils/vector_utils.h"
0005 #include "edm4hep/utils/kinematics.h"
0006 #include "edm4eic/ClusterCollection.h"
0007 #include "edm4eic/MCRecoParticleAssociationCollection.h"
0008 #include "podio/Frame.h"
0009 #include "podio/ROOTFrameReader.h"
0010
0011 // ROOT headers
0012 #include "TTree.h"
0013 #include "TFile.h"
0014 #include "TString.h"
0015 #include "TVector3.h"
0016 #include "TLorentzVector.h"
0017
0018 // Analysis headers
0019 #include "InclusiveSkim.h"
0020 #include "ElectronID.cc"
0021
0022 void InclusiveSkim() {
0023
0024     double Ee = 10.;
0025     double Eh = 100.;
0026
0027     vector<std::string> inFiles = {"pythia8NCDIS_10x100_minQ2=10_beamEffects_xAngle=-0.025_hiDiv_1.0001.eicrecon.tree.edm4eic.root"};
0028
0029     auto reader = podio::ROOTFrameReader();
0030     reader.openFiles(inFiles);
0031
0032     ElectronID* eFinder = new ElectronID(Ee, Eh);
0033
0034     TString outFileName = Form("inclusive_skim_%.0fx%.0fGeV.root", Ee, Eh);
0035     CreateOutputTree(outFileName);
0036
0037
0038     for(size_t ev = 0; ev < reader.getEntries("events"); ev++) {
0039
0040         if(ev%1000==0) cout << "Event " << ev << "/" << reader.getEntries("events") << endl;
0041
0042         ResetVariables();
0043
0044         const auto event = podio::Frame(reader.readNextEntry("events"));
0045         eFinder->SetEvent(&event);
0046
0047         edm4hep::MCParticleCollection e_mc = eFinder->GetMCElectron();
0048
0049         // Skip events without scattered MC electron
0050         if(e_mc.size() == 0) continue;
0051
0052         // Find scattered electrons
0053         auto e_candidates = eFinder->FindScatteredElectron();
0054         edm4eic::ReconstructedParticle e_rec;
0055         // If there are multiple candidates,
0056         // select one with highest pT
0057         if(e_candidates.size() > 0) {
0058             positive_eID = 1;
0059             e_rec = eFinder->SelectHighestPT(e_candidates);
0060         }
0061
0062         // Get momentum vector elements from MC electron
0063         mc_p = edm4hep::utils::magnitude(e_mc[0].getMomentum());
0064         mc_eta = edm4hep::utils::eta(e_mc[0].getMomentum());
0065         mc_phi = edm4hep::utils::angleAzimuthal(e_mc[0].getMomentum());
0066
0067         // Calculate kinematic variables using MC electron
0068         TLorentzVector kprime;
0069         kprime.SetXYZM(e_mc[0].getMomentum().x, e_mc[0].getMomentum().y, e_mc[0].getMomentum().z, Me);
0070         CalculateElectronKinematics(Ee, Eh, kprime, mc_xB, mc_Q2, mc_W, mc_y, mc_nu);
0071
0072         // If electron was identified, get vector elements
0073         // and kinematic variables from reconstructed particle
0074         if(positive_eID) {
0075             track_p = edm4hep::utils::magnitude(e_rec.getMomentum());
0076             track_eta = edm4hep::utils::eta(e_rec.getMomentum());
0077             track_phi = edm4hep::utils::angleAzimuthal(e_rec.getMomentum());
0078
0079             kprime.SetXYZM(e_rec.getMomentum().x, e_rec.getMomentum().y, e_rec.getMomentum().z, Me);
0080             CalculateElectronKinematics(Ee, Eh, kprime, e_track_xB, e_track_Q2, e_track_W, e_track_y, e_track_nu);
0081
0082             // Calculate kinematic variables with calorimeter energy
0083             TVector3 e3v = kprime.Vect();
0084             double track_clust_E = eFinder->GetCalorimeterEnergy(e_rec);
0085             e3v.SetMag(track_clust_E);
0086             kprime.SetVectM(e3v, Me);
0087             CalculateElectronKinematics(Ee, Eh, kprime, e_clust_xB, e_clust_Q2, e_clust_W, e_clust_y, e_clust_nu);
0088         }
0089
0090
0091         outTree->Fill();
0092
0093     }
0094
0095     outFile->cd();
0096     outTree->Write(outTree->GetName(), 2);
0097
0098 }
0099
0100
0101 void CreateOutputTree(TString outFileName) {
0102
0103     outFile = new TFile(outFileName, "RECREATE");
0104     outTree = new TTree("T", "Reconstructed and generated information from EICRecon");
0105
0106     outTree->Branch("positive_eID",     &positive_eID);
0107
0108     outTree->Branch("mc_p",         &mc_p);
0109     outTree->Branch("mc_eta",       &mc_eta);
0110     outTree->Branch("mc_phi",       &mc_phi);
0111
0112     outTree->Branch("track_p",      &track_p);
0113     outTree->Branch("track_eta",        &track_eta);
0114     outTree->Branch("track_phi",        &track_phi);
0115
0116     outTree->Branch("mc_xB",        &mc_xB);
0117     outTree->Branch("mc_Q2",        &mc_Q2);
0118     outTree->Branch("mc_W",         &mc_W);
0119     outTree->Branch("mc_y",         &mc_y);
0120     outTree->Branch("mc_nu",        &mc_nu);
0121
0122     outTree->Branch("e_track_xB",       &e_track_xB);
0123     outTree->Branch("e_track_Q2",       &e_track_Q2);
0124     outTree->Branch("e_track_W",        &e_track_W);
0125     outTree->Branch("e_track_y",        &e_track_y);
0126     outTree->Branch("e_track_nu",       &e_track_nu);
0127
0128     outTree->Branch("e_clust_xB",       &e_clust_xB);
0129     outTree->Branch("e_clust_Q2",       &e_clust_Q2);
0130     outTree->Branch("e_clust_W",        &e_clust_W);
0131     outTree->Branch("e_clust_y",        &e_clust_y);
0132     outTree->Branch("e_clust_nu",       &e_clust_nu);
0133
0134 }
0135
0136 void ResetVariables() {
0137
0138     positive_eID = 0;
0139
0140     mc_p = -999;
0141     mc_eta = -999;
0142     mc_phi = -999;
0143
0144     track_p = -999;
0145     track_eta = -999;
0146     track_phi = -999;
0147
0148     mc_xB = -999;
0149     mc_Q2 = -999;
0150     mc_W = -999;
0151     mc_y = -999;
0152     mc_nu = -999;
0153
0154     e_track_xB = -999;
0155     e_track_Q2 = -999;
0156     e_track_W = -999;
0157     e_track_y = -999;
0158     e_track_nu = -999;
0159
0160     e_clust_xB = -999;
0161     e_clust_Q2 = -999;
0162     e_clust_W = -999;
0163     e_clust_y = -999;
0164     e_clust_nu = -999;
0165
0166
0167 }
0168
0169 void CalculateElectronKinematics(double fEe, double fEh, TLorentzVector kf, float& xB, float& Q2, float& W, float& y, float& nu) {
0170
0171         TLorentzVector ki; ki.SetXYZM(0., 0., -fEe, Me);
0172         TLorentzVector P = GetHadronBeam(fEh);
0173         TLorentzVector q = ki - kf;
0174         Q2 = -(q.Dot(q));
0175         nu = (q.Dot(P))/Mp;
0176         xB = Q2/(2.*Mp*nu);
0177         y  = (q.Dot(P))/(ki.Dot(P));
0178         W  = sqrt(Mp*Mp + (2.*Mp*nu) - Q2);
0179 }
0180
0181
0182 TLorentzVector GetHadronBeam(double fEh) {
0183
0184     TLorentzVector hadron_beam;
0185     hadron_beam.SetX(fEh*sin(crossing_angle));
0186     hadron_beam.SetY(0.);
0187     hadron_beam.SetZ(fEh*cos(crossing_angle));
0188     hadron_beam.SetE(std::hypot(fEh, Mp));
0189     return hadron_beam;
0190
0191 }