Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
TrackletAna.cxx
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file TrackletAna.cxx
1 #include <TFile.h>
2 #include <TMath.h>
3 #include <TObjString.h>
4 #include <TRandom3.h>
5 #include <TTree.h>
6 #include <TTreeIndex.h>
7 
8 #include <fstream>
9 #include <iostream>
10 #include <sstream>
11 #include <string>
12 #include <vector>
13 
14 #include "GenHadron.h"
15 #include "Tracklet.h"
16 #include "Vertex.h"
17 #include "pdgidfunc.h"
18 
19 int main(int argc, char *argv[])
20 {
21  // Usage: ./TrackletAna [isdata] [evt-vtx map file] [infile] [outfile] [NevtToRun] [skip] [dRCut]
22  bool IsData = (TString(argv[1]).Atoi() == 1) ? true : false;
23  TString EvtVtx_map_filename = TString(argv[2]);
24  TString infilename = TString(argv[3]);
25  TString outfilename = TString(argv[4]);
26  int NevtToRun_ = TString(argv[5]).Atoi();
27  float dRCut = TString(argv[6]).Atof();
28 
29  // int iniEvt = skip;
30 
31  cout << "[Run Info] Event-vertex map file = " << EvtVtx_map_filename << endl
32  << " Input file = " << infilename << endl
33  << " Output file = " << outfilename << endl
34  << " Number of events to run = " << NevtToRun_ << endl
35  << " dRCut = " << dRCut << endl
36  << "-----------" << endl;
37 
38  cout << "[Run Info] NevtToRun: " << NevtToRun_ << ", dRCut: " << dRCut << endl;
39 
40  TrackletData tkldata = {};
41 
42  std::map<int, vector<float>> EvtVtx_map = EvtVtx_map_tklcluster(EvtVtx_map_filename.Data());
43 
44  TFile *f = new TFile(infilename, "READ");
45  TTree *t = (TTree *)f->Get("EventTree");
46  t->BuildIndex("event"); // Reference: https://root-forum.cern.ch/t/sort-ttree-entries/13138
47  TTreeIndex *index = (TTreeIndex *)t->GetTreeIndex();
48  int event, NTruthVtx;
49  float TruthPV_trig_x, TruthPV_trig_y, TruthPV_trig_z, centrality_mbd, centrality_mbdquantity;
50  vector<int> *ClusLayer = 0;
51  vector<float> *ClusX = 0, *ClusY = 0, *ClusZ = 0, *ClusPhiSize = 0, *ClusZSize = 0;
52  vector<int> *UniqueAncG4P_PID = 0;
53  vector<float> *UniqueAncG4P_Pt = 0, *UniqueAncG4P_Eta = 0, *UniqueAncG4P_Phi = 0, *UniqueAncG4P_E = 0;
54  t->SetBranchAddress("event", &event);
55  if (!IsData)
56  {
57  t->SetBranchAddress("NTruthVtx", &NTruthVtx);
58  t->SetBranchAddress("TruthPV_trig_x", &TruthPV_trig_x);
59  t->SetBranchAddress("TruthPV_trig_y", &TruthPV_trig_y);
60  t->SetBranchAddress("TruthPV_trig_z", &TruthPV_trig_z);
61  t->SetBranchAddress("UniqueAncG4P_PID", &UniqueAncG4P_PID);
62  t->SetBranchAddress("UniqueAncG4P_Pt", &UniqueAncG4P_Pt);
63  t->SetBranchAddress("UniqueAncG4P_Eta", &UniqueAncG4P_Eta);
64  t->SetBranchAddress("UniqueAncG4P_Phi", &UniqueAncG4P_Phi);
65  t->SetBranchAddress("UniqueAncG4P_E", &UniqueAncG4P_E);
66  }
67  t->SetBranchAddress("centrality_mbd", &centrality_mbd);
68  t->SetBranchAddress("centrality_mbdquantity", &centrality_mbdquantity);
69  t->SetBranchAddress("ClusLayer", &ClusLayer);
70  t->SetBranchAddress("ClusX", &ClusX);
71  t->SetBranchAddress("ClusY", &ClusY);
72  t->SetBranchAddress("ClusZ", &ClusZ);
73  t->SetBranchAddress("ClusPhiSize", &ClusPhiSize);
74  t->SetBranchAddress("ClusZSize", &ClusZSize);
75 
76  TFile *outfile = new TFile(outfilename, "RECREATE");
77  TTree *minitree = new TTree("minitree", "Minitree of Reconstructed Tracklets");
78  SetMinitree(minitree, tkldata);
79 
80  for (int i = 0; i < NevtToRun_; i++)
81  {
82  Long64_t local = t->LoadTree(index->GetIndex()[i]);
83  t->GetEntry(local);
84 
85  // cout << "event = " << event << " local = " << local << endl;
86  vector<float> PV = EvtVtx_map[event];
87  tkldata.event = event;
88  tkldata.PV_x = PV[0];
89  tkldata.PV_y = PV[1];
90  tkldata.PV_z = PV[2];
91  if (!IsData)
92  {
93  tkldata.TruthPV_x = TruthPV_trig_x;
94  tkldata.TruthPV_y = TruthPV_trig_y;
95  tkldata.TruthPV_z = TruthPV_trig_z;
96  }
97  // Centrality
98  tkldata.Centrality_mbd = centrality_mbd;
99  tkldata.Centrality_mbdquantity = centrality_mbdquantity;
100 
101  cout << "[Event info] Event = " << event << endl;
102 
103  // A selection to select events -> mimic 0-pileup environment
104  tkldata.pu0_sel = (NTruthVtx == 1) ? true : false;
105  tkldata.trig = true;
106  tkldata.process = 0; // Exclude single diffractive events
107 
108  // Prepare clusters
109  for (size_t ihit = 0; ihit < ClusLayer->size(); ihit++)
110  {
111  if (ClusLayer->at(ihit) < 3 || ClusLayer->at(ihit) > 6)
112  {
113  cout << "[WARNING] Unknown layer: " << ClusLayer->at(ihit) << endl; // Should not happen
114  continue;
115  }
116 
117  int layer = (ClusLayer->at(ihit) == 3 || ClusLayer->at(ihit) == 4) ? 0 : 1;
118  Hit *hit = new Hit(ClusX->at(ihit), ClusY->at(ihit), ClusZ->at(ihit), PV[0], PV[1], PV[2], 0);
119  tkldata.layers[layer].push_back(hit);
120  tkldata.clusphi.push_back(hit->Phi());
121  tkldata.cluseta.push_back(hit->Eta());
122  tkldata.clusphisize.push_back(ClusPhiSize->at(ihit));
123  tkldata.cluszsize.push_back(ClusZSize->at(ihit));
124  }
125 
126  tkldata.NClusLayer1 = tkldata.layers[0].size();
127 
128  // Tracklet reconstruction: proto-tracklets -> reco-tracklets -> gen-hadron matching
129  ProtoTracklets(tkldata, dRCut);
130  RecoTracklets(tkldata);
131  if (!IsData)
132  {
133  // Generated charged hadrons
134  for (size_t ihad = 0; ihad < UniqueAncG4P_PID->size(); ihad++)
135  {
136  if (is_chargedHadron(UniqueAncG4P_PID->at(ihad)) == false)
137  continue;
138 
139  GenHadron *genhadron = new GenHadron(UniqueAncG4P_Pt->at(ihad), UniqueAncG4P_Eta->at(ihad), UniqueAncG4P_Phi->at(ihad), UniqueAncG4P_E->at(ihad));
140  tkldata.GenHadrons.push_back(genhadron);
141  tkldata.GenHadron_Pt.push_back(UniqueAncG4P_Pt->at(ihad));
142  tkldata.GenHadron_eta.push_back(UniqueAncG4P_Eta->at(ihad));
143  tkldata.GenHadron_phi.push_back(UniqueAncG4P_Phi->at(ihad));
144  tkldata.GenHadron_E.push_back(UniqueAncG4P_E->at(ihad));
145  }
146  tkldata.NGenHadron = tkldata.GenHadrons.size();
147 
148  GenMatch_Recotkl(tkldata);
149 
150  cout << "NCluster layer 1 = " << tkldata.NClusLayer1 << "; NRecotkl_Raw = " << tkldata.NRecotkl_Raw << "; NRecotkl_GenMatched = " << tkldata.NRecotkl_GenMatched << endl;
151  }
152  minitree->Fill();
153  ResetVec(tkldata);
154  cout << "----------" << endl;
155  }
156 
157  outfile->cd();
158  minitree->Write("", TObject::kOverwrite);
159  outfile->Close();
160 
161  f->Close();
162 
163  return 0;
164 }