EICMCAnalysis/jets/fill_jet_tree.C

0001 #include "fastjet/ClusterSequence.hh"
0002
0003 #include <vector>
0004 #include <iostream>
0005 #include <fstream>
0006 #include <cstdio>
0007 #include <cstdlib>
0008 #include <cassert>
0009
0010 #include "TTree.h"
0011 #include "TFile.h"
0012
0013 #include "eicsmear/erhic/EventBase.h"
0014 #include "eicsmear/erhic/EventPythia.h"
0015 #include "eicsmear/erhic/Particle.h"
0016 #include "eicsmear/erhic/ParticleMC.h"
0017 #include "eicsmear/smear/EventSmear.h"
0018 #include "eicsmear/smear/ParticleMCS.h"
0019
0020 using namespace fastjet;
0021 using namespace std;
0022
0023
0024 /** Find jet in given vector of jets that's closest in eta-phi to reference jet
0025  */
0026 PseudoJet* find_matching_jet( const PseudoJet* refjet, vector<PseudoJet>* vjets )
0027 {
0028   /* PidCandidate with eta, phi closest to reference */
0029   PseudoJet* best_candidate = NULL;
0030
0031   float eta_ref = refjet->eta();
0032   float phi_ref = refjet->phi();
0033
0034   /* find jet with smallest delta_R */
0035   float min_delta_R = 100;
0036   vector<PseudoJet>::iterator min_delta_R_iter = vjets->end();
0037
0038   vector<PseudoJet>::iterator ijet;
0039
0040   for ( ijet = vjets->begin(); ijet != vjets->end(); ijet++ )
0041     {
0042       float eta = ijet->eta();
0043       float phi = ijet->phi();
0044
0045       float delta_R = sqrt( pow( eta - eta_ref, 2 ) + pow( phi - phi_ref, 2 ) );
0046
0047       if ( delta_R < min_delta_R )
0048         {
0049           min_delta_R_iter = ijet;            ;
0050           min_delta_R = delta_R;
0051         }
0052     }
0053
0054   /* set best_candidate to PidCandiate with smallest delta_R within reasonable range*/
0055   if ( min_delta_R_iter != vjets->end() && min_delta_R < 0.5 )
0056     best_candidate = &(*min_delta_R_iter);
0057
0058   return best_candidate;
0059 }
0060
0061
0062 /** Main function
0063  */
0064 int main(int argc, char* argv[]) {
0065
0066   if ( argc != 6 )
0067     {
0068       cerr << "Wrong number of arguments! Exit." << endl;
0069       return -1;
0070     }
0071
0072   const char* truthFileName = argv[1];
0073   const char* smearFileName = argv[2];
0074   const char* outFileName = argv[3];
0075
0076   /* R value for fastjet algorithm */
0077   const float fastjetR = atof(argv[4]);
0078
0079   /* minimum pt for jets that are kept for analysis */
0080   const float ptmin = atof(argv[5]);
0081
0082   cout << "truthFileName = " << truthFileName << endl;
0083   cout << "smearFileName = " << smearFileName << endl;
0084   cout << "outFileName = " << outFileName << endl;
0085   cout << "fastjetR      = " << fastjetR << endl;
0086   cout << "ptmin         = " << ptmin << endl;
0087
0088   /* Open files and retrieve trees */
0089   TFile *file_mc_truth = new TFile(truthFileName, "OPEN");
0090   TFile *file_mc_smeared = new TFile(smearFileName, "OPEN");
0091
0092   TTree* tree_truth = (TTree*)file_mc_truth->Get("EICTree");
0093   TTree *tree_smeared = (TTree*)file_mc_smeared->Get("Smeared");
0094
0095   tree_truth->AddFriend(tree_smeared);
0096
0097   /* Setup Input Event Buffer */
0098   erhic::EventPythia* event_truth(NULL);
0099   Smear::Event* event_smear(NULL);
0100
0101   tree_truth->SetBranchAddress("event",&event_truth);
0102   tree_truth->SetBranchAddress("eventS", &event_smear);
0103
0104   /* Set more buffers */
0105   Double_t event_x = NAN;
0106   Double_t event_q2 = NAN;
0107
0108   tree_truth->SetBranchAddress("trueX",&event_x);
0109   tree_truth->SetBranchAddress("trueQ2",&event_q2);
0110
0111   /* Open Output File */
0112   TFile *ofile = TFile::Open(outFileName,"recreate");
0113   assert(ofile);
0114
0115   /* Create Jet Tree */
0116   Int_t   _event_id    = -999;
0117   Int_t   _event_njets = -999;
0118   Double_t _event_truth_x     = NAN;
0119   Double_t _event_truth_q2    = NAN;
0120
0121   Int_t   _jet_truth_id          = -999;
0122   Int_t   _jet_truth_ncomp       = -999;
0123   Int_t   _jet_truth_ncharged    = -999;
0124   Double_t _jet_truth_e           = NAN;
0125   Double_t _jet_truth_et          = NAN;
0126   Double_t _jet_truth_eta         = NAN;
0127   Double_t _jet_truth_phi         = NAN;
0128   Double_t _jet_truth_minv        = NAN;
0129   Double_t _jet_truth_eem         = NAN;
0130   Double_t _jet_truth_rvtx        = NAN;
0131   Double_t _jet_truth_rmean       = NAN;
0132   Double_t _jet_truth_rstdev      = NAN;
0133
0134   Int_t   _jet_smear_id          = -999;
0135   Int_t   _jet_smear_ncomp       = -999;
0136   Int_t   _jet_smear_ncharged    = -999;
0137   Double_t _jet_smear_e           = NAN;
0138   Double_t _jet_smear_et          = NAN;
0139   Double_t _jet_smear_eta         = NAN;
0140   Double_t _jet_smear_phi         = NAN;
0141   Double_t _jet_smear_minv        = NAN;
0142   Double_t _jet_smear_eem         = NAN;
0143   Double_t _jet_smear_rvtx        = NAN;
0144   Double_t _jet_smear_rmean       = NAN;
0145   Double_t _jet_smear_rstdev      = NAN;
0146
0147   /* Create jet tree */
0148   TTree *mTree = new TTree("jets","Jet Tree");
0149   mTree->Branch("event_id", &_event_id);
0150   mTree->Branch("event_njets", &_event_njets);
0151   mTree->Branch("event_truth_x", &_event_truth_x);
0152   mTree->Branch("event_truth_q2", &_event_truth_q2);
0153
0154   mTree->Branch("jet_truth_id", &_jet_truth_id);
0155   mTree->Branch("jet_truth_ncomp", &_jet_truth_ncomp);
0156   mTree->Branch("jet_truth_ncharged", &_jet_truth_ncharged);
0157   mTree->Branch("jet_truth_e", &_jet_truth_e);
0158   mTree->Branch("jet_truth_et", &_jet_truth_et);
0159   mTree->Branch("jet_truth_eta", &_jet_truth_eta);
0160   mTree->Branch("jet_truth_phi", &_jet_truth_phi);
0161   mTree->Branch("jet_truth_minv", &_jet_truth_minv);
0162   mTree->Branch("jet_truth_eem", &_jet_truth_eem);
0163   mTree->Branch("jet_truth_rvtx", &_jet_truth_rvtx);
0164   mTree->Branch("jet_truth_rmean", &_jet_truth_rmean);
0165   mTree->Branch("jet_truth_rstdev", &_jet_truth_rstdev);
0166
0167   mTree->Branch("jet_smear_id", &_jet_smear_id);
0168   mTree->Branch("jet_smear_ncomp", &_jet_smear_ncomp);
0169   mTree->Branch("jet_smear_ncharged", &_jet_smear_ncharged);
0170   mTree->Branch("jet_smear_e", &_jet_smear_e);
0171   mTree->Branch("jet_smear_et", &_jet_smear_et);
0172   mTree->Branch("jet_smear_eta", &_jet_smear_eta);
0173   mTree->Branch("jet_smear_phi", &_jet_smear_phi);
0174   mTree->Branch("jet_smear_minv", &_jet_smear_minv);
0175   mTree->Branch("jet_smear_eem", &_jet_smear_eem);
0176   mTree->Branch("jet_smear_rvtx", &_jet_smear_rvtx);
0177   mTree->Branch("jet_smear_rmean", &_jet_smear_rmean);
0178   mTree->Branch("jet_smear_rstdev", &_jet_smear_rstdev);
0179
0180   /* Loop Over Events in Simu Trees */
0181   for(Int_t iEvent=0; iEvent<tree_truth->GetEntries(); iEvent++)
0182     {
0183       /* Read Next Event */
0184       if(tree_truth->GetEntry(iEvent) <=0) break;
0185
0186       if(iEvent%10000 == 0) cout << "Event " << iEvent << endl;
0187
0188       /* reset tree variables */
0189       _event_id = iEvent;
0190       _event_njets = 0;
0191       _event_truth_x = event_x;
0192       _event_truth_q2 = event_q2;
0193
0194       /* Create laboratory frame jet vectors */
0195       vector<PseudoJet> jetcomponent_truth;
0196       vector<PseudoJet> jetcomponent_smear;
0197
0198       /* Loop over truth Particles, fill jet component vector */
0199       for(Int_t j=0; j<event_truth->GetNTracks(); j++)
0200     {
0201       const Particle* inParticle = event_truth->GetTrack(j);
0202
0203       if ( !inParticle)
0204         continue;
0205
0206       /* Select Particles for Jets */
0207       if(j>10 && inParticle->GetStatus() == 1 && inParticle->GetParentIndex() != 3)
0208         {
0209           if( abs(inParticle->GetEta()) <= 5 && inParticle->GetE() >= 0.250 )
0210         {
0211           /* Truth particle: Get all information directly from particle */
0212           Double_t px = inParticle->GetPx();
0213           Double_t py = inParticle->GetPy();
0214           Double_t pz = inParticle->GetPz();
0215           Double_t E = inParticle->GetE();
0216
0217           fastjet::PseudoJet pPt(px,py,pz,E);
0218           pPt.set_user_index(inParticle->GetIndex());
0219           jetcomponent_truth.push_back(pPt);
0220         }
0221         }
0222     }
0223
0224       /* Loop over SMEARED jet components: Energy in Calorimeter (like 'tower jets'). Fill jet component vector. */
0225       for(Int_t js=0; js<event_smear->GetNTracks(); js++)
0226     {
0227       const Smear::ParticleMCS* inParticle = event_smear->GetTrack(js);
0228
0229       if ( !inParticle)
0230         continue;
0231
0232       /* Select Particles for Jets */
0233       if(js>10 && inParticle->GetStatus() == 1 && inParticle->GetParentIndex() != 3)
0234         {
0235           if( inParticle->GetE() >= 0.250 )
0236         {
0237           /* Calorimeter: Get energy from smeared particle, and ... */
0238           Double_t E = inParticle->GetE();
0239
0240 //        if ( E == 0 )
0241 //          cout << "E == 0 found! PID = " << event_truth->GetTrack(js)->GetPdgCode() << " , E_true = " <<  event_truth->GetTrack(js)->GetE() << " , Eta_true = " <<  event_truth->GetTrack(js)->GetEta() << endl;
0242
0243           /* ... get theta, phi from truth particle */
0244           Double_t phi = event_truth->GetTrack(js)->GetPhi();
0245           Double_t eta = event_truth->GetTrack(js)->GetEta();
0246
0247           /* Recalculate px, py, pz based on calorimeter enrgy and truth direction */
0248           Double_t pT = E / cosh( eta );
0249           Double_t px = pT * cos( phi );
0250           Double_t py = pT * sin( phi );
0251           Double_t pz = pT * sinh( eta );
0252
0253           fastjet::PseudoJet pPt(px,py,pz,E);
0254           pPt.set_user_index(event_truth->GetTrack(js)->GetIndex());
0255           jetcomponent_smear.push_back(pPt);
0256         }
0257         }
0258     }
0259
0260       /* Set Jet Definitions */
0261       JetDefinition jet_def_antikt(antikt_algorithm,fastjetR);
0262
0263       /* Run Clustering and Extract the Jets */
0264
0265       /* Lab Frame Cluster */
0266       ClusterSequence cluster_truth_antikt(jetcomponent_truth, jet_def_antikt);
0267       ClusterSequence cluster_smear_antikt(jetcomponent_smear, jet_def_antikt);
0268
0269       /* Lab Frame Jets*/
0270       vector<PseudoJet> jets_truth_antikt = sorted_by_pt(cluster_truth_antikt.inclusive_jets(ptmin));
0271       vector<PseudoJet> jets_smear_antikt = sorted_by_pt(cluster_smear_antikt.inclusive_jets(ptmin));
0272
0273       /* loop over SMEARED jets */
0274       _event_njets =  jets_smear_antikt.size();
0275       for ( unsigned ijet = 0; ijet < _event_njets; ijet++ )
0276     {
0277       PseudoJet* jetMatch = find_matching_jet( &(jets_smear_antikt.at(ijet)), &jets_truth_antikt );
0278
0279       /* Set SMEARED jet variables */
0280       _jet_smear_id          = ijet;
0281       _jet_smear_ncomp       = -999;
0282       _jet_smear_ncharged    = -999;
0283       _jet_smear_e           = jets_smear_antikt.at(ijet).E();
0284       _jet_smear_et          = jets_smear_antikt.at(ijet).Et();
0285       _jet_smear_eta         = jets_smear_antikt.at(ijet).eta();
0286       _jet_smear_phi         = jets_smear_antikt.at(ijet).phi_std();
0287       _jet_smear_minv        = jets_smear_antikt.at(ijet).m();
0288       _jet_smear_eem         = NAN;
0289       _jet_smear_rvtx        = NAN;
0290       _jet_smear_rmean       = NAN;
0291       _jet_smear_rstdev      = NAN;
0292
0293       /* Set TRUTH jet variables */
0294       _jet_truth_id          = -999;
0295       _jet_truth_ncomp       = -999;
0296       _jet_truth_ncharged    = -999;
0297       _jet_truth_e           = NAN;
0298       _jet_truth_et          = NAN;
0299       _jet_truth_eta         = NAN;
0300       _jet_truth_phi         = NAN;
0301       _jet_truth_minv        = NAN;
0302       _jet_truth_eem         = NAN;
0303       _jet_truth_rvtx        = NAN;
0304       _jet_truth_rmean       = NAN;
0305       _jet_truth_rstdev      = NAN;
0306
0307       if ( jetMatch )
0308         {
0309           _jet_truth_id          = ijet;
0310           _jet_truth_ncomp       = -999;
0311           _jet_truth_ncharged    = -999;
0312           _jet_truth_e           = jetMatch->E();
0313           _jet_truth_et          = jetMatch->Et();
0314           _jet_truth_eta         = jetMatch->eta();
0315           _jet_truth_phi         = jetMatch->phi_std();
0316           _jet_truth_minv        = jetMatch->m();
0317           _jet_truth_eem         = NAN;
0318           _jet_truth_rvtx        = NAN;
0319           _jet_truth_rmean       = NAN;
0320           _jet_truth_rstdev      = NAN;
0321         }
0322
0323       /* Fill tree */
0324       mTree->Fill();
0325     }
0326     }
0327
0328   /* Write output tree and close file */
0329   mTree->Write();
0330   ofile->Close();
0331
0332   return 0;
0333 }