Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
TruthCaloTree.C
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file TruthCaloTree.C
1 // Truth Calorimeter Calibration TreeMaker
2 #include <fun4all/Fun4AllBase.h>
3 #include <iostream>
4 #include "TruthCaloTree.h"
5 
6 // General F4A includes
7 #include <fun4all/Fun4AllServer.h>
8 #include <fun4all/Fun4AllReturnCodes.h>
9 #include <phool/getClass.h>
10 #include <phool/PHCompositeNode.h>
11 #include <g4main/PHG4TruthInfoContainer.h>
12 #include <g4main/PHG4Particle.h>
13 #include <g4main/PHG4Hit.h>
14 #include <g4main/PHG4VtxPoint.h>
15 
16 // Calorimeter includes
17 #include <calobase/RawTowerv2.h>
18 #include <calobase/RawTowerGeom.h>
19 #include <calobase/RawTowerContainer.h>
20 #include <calobase/RawCluster.h>
21 #include <calobase/RawClusterContainer.h>
22 
23 // ROOT includes
24 #include "TLorentzVector.h"
25 #include "TMath.h"
26 #include "TROOT.h"
27 #include "TH2.h"
28 
29 TruthCaloTree::TruthCaloTree(const std::string &name) : SubsysReco("TRUTH_CALO_TREE")
30 {
31  _foutname = name;
32  _hcal_sim_name = "TOWER_SIM_HCALOUT";
33  _hcalIN_sim_name = "TOWER_SIM_HCALIN";
34  _EMcal_sim_name = "TOWER_SIM_CEMC";
35 
36 }
37 
38 int TruthCaloTree::reset_tree_vars() {
39 
40  _b_event = -99;
41  _b_tower_sim_n = -99;
42  _b_EMcal_sim_n = -99;
43  _b_hcalIN_sim_n = -99;
44  _nBH = -99;
45 
46  for (int i = 0; i <nTowers; i++){
47 
48  _b_tower_sim_E[i] = -99;
49  _b_tower_sim_eta[i] = -99;
50  _b_tower_sim_phi[i] = -99;
51  _b_tower_sim_ieta[i] = -99;
52  _b_tower_sim_iphi[i] = -99;
53 
54  _b_EMcal_sim_E[i] = -99;
55  _b_EMcal_sim_eta[i] = -99;
56  _b_EMcal_sim_phi[i] = -99;
57  _b_EMcal_sim_iphi[i] = -99;
58  _b_EMcal_sim_ieta[i] = -99;
59 
60  _b_hcalIN_sim_E[i] = -99;
61  _b_hcalIN_sim_eta[i] = -99;
62  _b_hcalIN_sim_phi[i] = -99;
63  _b_hcalIN_sim_iphi[i] = -99;
64  _b_hcalIN_sim_ieta[i] = -99;
65  }
66 
67  _b_n_truth = -99;
68  n_child = -99;
69  n_vertex = -99;
70 
71  for (int i = 0; i < nTruth; i++) {
72  _b_truthenergy[i] = -99;
73  _b_trutheta[i] = -99;
74  _b_truthphi[i] = -99;
75  _b_truthpt[i] = -99;
76  _b_truthp[i] = -99;
77  _b_truthpid[i] = -99;
78  _b_truth_trackid[i] = -99;
79  vertex_id[i] = -99;
80  vertex_x[i] = -99;
81  vertex_y[i] = -99;
82  vertex_z[i] = -99;
83  _BH_e[i] = -99;
84  _BH_px[i] = -99;
85  _BH_py[i] = -99;
86  _BH_pz[i] = -99;
87  _BH_track_id[i] = -99;
88  }
89 
90  for (int i = 0; i < nTruth; i++) {
91  child_vertex_id[i] = -99;
92  child_parent_id[i] = -99;
93  child_pid[i] = -99;
94  child_px[i] = -99;
95  child_py[i] = -99;
96  child_pz[i] = -99;
97  child_energy[i] = -99;
98  }
99 
100  return 1;
101 }
102 
103 int TruthCaloTree::GetNodes(PHCompositeNode *topNode) {
104 
105  blackhole = findNode::getClass<PHG4HitContainer>(topNode, "G4HIT_BH_1");
106  if (!blackhole) std::cout << "No blackhole" << std::endl;
107 
108  if (_debug) std::cout<<"GettingNodes..."<<std::endl;
109  _geomOH = findNode::getClass<RawTowerGeomContainer>(topNode, "TOWERGEOM_HCALOUT");
110  if(!_geomOH) std::cout<<"No TOWERGeOM_HCALOUT"<<std::endl;
111 
112  _geomIH = findNode::getClass<RawTowerGeomContainer>(topNode, "TOWERGEOM_HCALIN");
113  if(!_geomIH) std::cout<<"No TOWERGeOM_HCALIN"<<std::endl;
114 
115  _geomEM = findNode::getClass<RawTowerGeomContainer>(topNode, "TOWERGEOM_CEMC");
116  if(!_geomEM) std::cout<<"No TOWERGeOM_CEMC"<<std::endl;
117 
118  _towersSimOH = findNode::getClass<RawTowerContainer>(topNode, _hcal_sim_name);
119  if (!_towersSimOH) std::cout<<"No TOWER_SIM_HCALOUT Node"<<std::endl;
120 
121  _towersSimIH = findNode::getClass<RawTowerContainer>(topNode, _hcalIN_sim_name);
122  if (!_towersSimIH) std::cout<<"No TOWER_SIM_HCALIN Node"<<std::endl;
123 
124  _towersSimEM = findNode::getClass<RawTowerContainer>(topNode, _EMcal_sim_name);
125  if (!_towersSimEM) std::cout<<"No TOWER_SIM_CEMC Node"<<std::endl;
126 
127  truthinfo = findNode::getClass<PHG4TruthInfoContainer>(topNode, "G4TruthInfo");
128  if (!truthinfo) std::cout << "PHG4TruthInfoContainer node is missing, can't collect G4 truth particles"<< std::endl;
129 
130  return 1;
131 }
132 
133 int TruthCaloTree::Init(PHCompositeNode *topNode) {
134  _ievent = 0;
135  _b_event = -1;
136 
137  if (_debug) std::cout<<"Initiating..."<<std::endl;
138 
139  reset_tree_vars();
140 
141  _file = new TFile(_foutname.c_str(), "RECREATE");
142 
143  _tree = new TTree("T", "keep on giving tree");
144 
145  _tree->Branch("n_truth",&_b_n_truth,"n_truth/I");
146  _tree->Branch("truthenergy",_b_truthenergy,"truthenergy[n_truth]/F");
147  _tree->Branch("trutheta",_b_trutheta,"trutheta[n_truth]/F");
148  _tree->Branch("truthphi",_b_truthphi,"truthphi[n_truth]/F");
149  _tree->Branch("truthpt",_b_truthpt,"truthpt[n_truth]/F");
150  _tree->Branch("truthp",_b_truthp,"truthp[n_truth]/F");
151  _tree->Branch("truthpid",_b_truthpid,"truthpid[n_truth]/I");
152  _tree->Branch("truth_track_id",_b_truth_trackid,"truth_track_id[n_truth]/I");
153 
154  _tree->Branch("EMcal_sim_n", &_b_EMcal_sim_n, "EMcal_sim_n/I");
155  _tree->Branch("EMcal_sim_E", _b_EMcal_sim_E, "EMcal_sim_E[EMcal_sim_n]/F");
156  _tree->Branch("EMcal_sim_eta", _b_EMcal_sim_eta, "EMcal_sim_eta[EMcal_sim_n]/F");
157  _tree->Branch("EMcal_sim_phi", _b_EMcal_sim_phi, "EMcal_sim_phi[EMcal_sim_n]/F");
158  _tree->Branch("EMcal_sim_iphi", _b_EMcal_sim_iphi, "EMcal_sim_iphi[EMcal_sim_n]/I");
159  _tree->Branch("EMcal_sim_ieta", _b_EMcal_sim_ieta, "EMcal_sim_ieta[EMcal_sim_n]/I");
160 
161  _tree->Branch("hcalIN_sim_n", &_b_hcalIN_sim_n, "hcalIN_sim_n/I");
162  _tree->Branch("hcalIN_sim_E", _b_hcalIN_sim_E, "hcalIN_sim_E[hcalIN_sim_n]/F");
163  _tree->Branch("hcalIN_sim_eta", _b_hcalIN_sim_eta, "hcalIN_sim_eta[hcalIN_sim_n]/F");
164  _tree->Branch("hcalIN_sim_phi", _b_hcalIN_sim_phi, "hcalIN_sim_phi[hcalIN_sim_n]/F");
165  _tree->Branch("hcalIN_sim_iphi", _b_hcalIN_sim_iphi, "hcalIN_sim_iphi[hcalIN_sim_n]/I");
166  _tree->Branch("hcalIN_sim_ieta", _b_hcalIN_sim_ieta, "hcalIN_sim_ieta[hcalIN_sim_n]/I");
167 
168  _tree->Branch("tower_sim_n",&_b_tower_sim_n, "tower_sim_n/I");
169  _tree->Branch("tower_sim_E",_b_tower_sim_E, "tower_sim_E[tower_sim_n]/F");
170  _tree->Branch("tower_sim_eta",_b_tower_sim_eta, "tower_sim_eta[tower_sim_n]/F");
171  _tree->Branch("tower_sim_phi",_b_tower_sim_phi, "tower_sim_phi[tower_sim_n]/F");
172  _tree->Branch("tower_sim_ieta",_b_tower_sim_ieta, "tower_sim_ieta[tower_sim_n]/I");
173  _tree->Branch("tower_sim_iphi",_b_tower_sim_iphi, "tower_sim_iphi[tower_sim_n]/I");
174 
175  _tree->Branch("n_vertex",&n_vertex,"n_vertex/I");
176  _tree->Branch("vertex_id",vertex_id,"vertex_id[n_vertex]/I");
177  _tree->Branch("vertex_x",vertex_x,"vertex_x[n_vertex]/F");
178  _tree->Branch("vertex_y",vertex_y,"vertex_y[n_vertex]/F");
179  _tree->Branch("vertex_z",vertex_z,"vertex_z[n_vertex]/F");
180 
181  _tree->Branch("n_child",&n_child,"n_child/I");
182  _tree->Branch("child_vertex_id",child_vertex_id,"child_vertex_id[n_child]/I");
183  _tree->Branch("child_parent_id",child_parent_id,"child_parent_id[n_child]/I");
184  _tree->Branch("child_pid",child_pid,"child_pid[n_child]/I");
185  _tree->Branch("child_px",child_px,"child_px[n_child]/F");
186  _tree->Branch("child_py",child_py,"child_py[n_child]/F");
187  _tree->Branch("child_pz",child_pz,"child_pz[n_child]/F");
188  _tree->Branch("child_energy",child_energy,"child_energy[n_child]/F");
189 
190  _tree->Branch("BH_n",&_nBH,"BH_n/I");
191  _tree->Branch("BH_px",_BH_px,"BH_px[BH_n]/F");
192  _tree->Branch("BH_py",_BH_py,"BH_py[BH_n]/F");
193  _tree->Branch("BH_pz",_BH_pz,"BH_pz[BH_n]/F");
194  _tree->Branch("BH_track_id",_BH_track_id,"BH_track_id[BH_n]/I");
195  _tree->Branch("BH_e",_BH_e,"BH_e[BH_n]/F");
196 
197 
198  return 0;
199 }
200 
201 int TruthCaloTree::process_event(PHCompositeNode *topNode) {
202 
203  GetNodes(topNode);
204 
205  _b_event = _ievent;
206  if (_ievent %5000==0) std::cout<<"Event: "<<_ievent<<std::endl;
207 
208 
209  if (_debug) std::cout<<"Processing Event: "<< _b_event<<std::endl;
210  if (_debug) std::cout << "hello";
211 
213  // BLACKHOLE INFO
215 
216  _nBH = 0;
217  PHG4HitContainer::ConstRange bh_hit_range = blackhole->getHits();
218  for (PHG4HitContainer::ConstIterator hit_iter = bh_hit_range.first;
219  hit_iter != bh_hit_range.second; hit_iter++)
220  {
221  PHG4Hit *this_hit = hit_iter->second;
222  assert(this_hit);
223  _BH_e[_nBH] = this_hit->get_edep();
224  _BH_px[_nBH] = this_hit->get_px(0);
225  _BH_py[_nBH] = this_hit->get_py(0);
226  _BH_pz[_nBH] = this_hit->get_pz(0);
227  PHG4Particle *p = truthinfo->GetParticle(this_hit->get_trkid());
228  _BH_track_id[_nBH] = p->get_track_id();
229  _nBH++;
230  }
231 
233  // OUTER HCAL
235 
236  _b_tower_sim_n=0;
237  RawTowerContainer::ConstRange begin_end_sim = _towersSimOH->getTowers();
238  if (_debug) std::cout<<"Got the iterator"<<std::endl;
239 
240  for (RawTowerContainer::ConstIterator rtiter = begin_end_sim.first; rtiter != begin_end_sim.second; ++rtiter) {
241  if ((_b_tower_sim_n%10==0)&&(_debug)) std::cout<<"At sim tower: "<< _b_tower_sim_n<<std::endl;
242 
243  RawTower *tower = rtiter->second;
244  RawTowerGeom *tower_geom = _geomOH->get_tower_geometry(tower->get_key());
245  _b_tower_sim_E [ _b_tower_sim_n ] = tower->get_energy();
246  _b_tower_sim_eta [ _b_tower_sim_n ] = tower_geom->get_eta();
247  _b_tower_sim_phi [ _b_tower_sim_n ] = tower_geom->get_phi();
248  _b_tower_sim_ieta[ _b_tower_sim_n ] = tower_geom->get_bineta();
249  _b_tower_sim_iphi[ _b_tower_sim_n ] = tower_geom->get_binphi();
250  _b_tower_sim_n++;
251  if(_b_tower_sim_n >= nTowers){
252  std::cout << __FILE__ << " ERROR: _b_tower_sim_n has hit cap of " << nTowers << "!!!" << std::endl;
253  }
254  }
255 
257  // INNER HCAL
259 
260  _b_hcalIN_sim_n = 0;
261  RawTowerContainer::ConstRange begin_end_simIN = _towersSimIH->getTowers();
262  for (RawTowerContainer::ConstIterator rtiter = begin_end_simIN.first; rtiter != begin_end_simIN.second; ++rtiter) {
263  RawTower *tower = rtiter->second;
264  RawTowerGeom *tower_geom = _geomIH->get_tower_geometry(tower->get_key());
265  _b_hcalIN_sim_E [ _b_hcalIN_sim_n ] = tower->get_energy();
266  _b_hcalIN_sim_eta [ _b_hcalIN_sim_n ] = tower_geom->get_eta();
267  _b_hcalIN_sim_phi [ _b_hcalIN_sim_n ] = tower_geom->get_phi();
268  _b_hcalIN_sim_ieta [ _b_hcalIN_sim_n ] = tower_geom->get_bineta();
269  _b_hcalIN_sim_iphi [ _b_hcalIN_sim_n ] = tower_geom->get_binphi();
270  _b_hcalIN_sim_n++;
271 
272  }
273 
275  // EM CAL
277 
278  _b_EMcal_sim_n = 0;
279  RawTowerContainer::ConstRange begin_end_simEM = _towersSimEM->getTowers();
280  for (RawTowerContainer::ConstIterator rtiter = begin_end_simEM.first; rtiter != begin_end_simEM.second; ++rtiter) {
281  RawTower *tower = rtiter->second;
282  RawTowerGeom *tower_geom = _geomEM->get_tower_geometry(tower->get_key());
283  _b_EMcal_sim_E [ _b_EMcal_sim_n ] = tower->get_energy();
284  _b_EMcal_sim_eta [ _b_EMcal_sim_n ] = tower_geom->get_eta();
285  _b_EMcal_sim_phi [ _b_EMcal_sim_n ] = tower_geom->get_phi();
286  _b_EMcal_sim_ieta [ _b_EMcal_sim_n ] = tower_geom->get_bineta();
287  _b_EMcal_sim_iphi [ _b_EMcal_sim_n ] = tower_geom->get_binphi();
288  _b_EMcal_sim_n++;
289 
290  }
291 
293  // PRIMARY TRUTH INFO
295 
296 
297  n_child = 0;
298  //std::set <int> primary;
299  std::set<int> vertex;
300  //if (!primary.empty()) primary.clear();
301  if (!vertex.empty()) vertex.clear();
302  PHG4TruthInfoContainer::Range range = truthinfo->GetPrimaryParticleRange();
303  _b_n_truth = 0;
304  for (PHG4TruthInfoContainer::ConstIterator iter = range.first;
305  iter != range.second; ++iter) {
306  const PHG4Particle *truth = iter->second;
307  if (!truthinfo->is_primary(truth)) continue;
308  _b_truthp[_b_n_truth] = sqrt(truth->get_px() * truth->get_px()
309  + truth->get_py() * truth->get_py()
310  + truth->get_pz() * truth->get_pz());
311  _b_truthpt[_b_n_truth] = sqrt(truth->get_px() * truth->get_px()
312  + truth->get_py() * truth->get_py());
313  _b_truthenergy[_b_n_truth] = truth->get_e();
314  _b_truthphi[_b_n_truth] = atan2(truth->get_py(),truth->get_px());
315  _b_trutheta[_b_n_truth] = atanh(truth->get_pz() / _b_truthenergy[_b_n_truth]);
316 
318  if (_b_trutheta[_b_n_truth] != _b_trutheta[_b_n_truth])
319  _b_trutheta[_b_n_truth] = -99;
320  _b_truthpid[_b_n_truth] = truth->get_pid();
321  _b_truth_trackid[_b_n_truth] = truth->get_track_id();
322 
323  //if (truth->get_e() > 2) primary.insert(truth->get_track_id());
324  _b_n_truth++;
325  if (_b_n_truth >= 50000) break;
326  }
327 
328  // for finding all particles of a hadron shower recursively
329  /*
330  PHG4TruthInfoContainer::Range second_range = truthinfo->GetSecondaryParticleRange();
331  for (PHG4TruthInfoContainer::ConstIterator iter = second_range.second; iter != second_range.first; --iter) {
332  PHG4Particle *child = iter->second;
333  if (primary.find(child->get_parent_id()) != primary.end()) {
334  primary.insert(child->get_track_id());
335  vertex.insert(child->get_vtx_id());
336  child_pid[n_child] = child->get_pid();
337  child_parent_id[n_child] = child->get_parent_id();
338  child_vertex_id[n_child] = child->get_vtx_id();
339  child_px[n_child] = child->get_px();
340  child_py[n_child] = child->get_py();
341  child_pz[n_child] = child->get_pz();
342  child_energy[n_child] = child->get_e();
343  n_child++;
344  }
345  }
346  */
347 
349  // SECONDARY AND VERTEX TRUTH INFO
351 
352 
353  PHG4TruthInfoContainer::Range second_range = truthinfo->GetSecondaryParticleRange();
354 
355  for (PHG4TruthInfoContainer::ConstIterator iter = second_range.first; iter != second_range.second; ++iter) {
356  const PHG4Particle *child = iter->second;
357  if (child->get_parent_id() > 0) {
358  PHG4Particle *parent = truthinfo->GetParticle(child->get_parent_id());
359  if (parent->get_e() > 2.0) {
360  vertex.insert(child->get_vtx_id());
361  child_pid[n_child] = child->get_pid();
362  child_parent_id[n_child] = child->get_parent_id();
363  child_vertex_id[n_child] = child->get_vtx_id();
364  child_px[n_child] = child->get_px();
365  child_py[n_child] = child->get_py();
366  child_pz[n_child] = child->get_pz();
367  child_energy[n_child] = child->get_e();
368  n_child++;
369  }
370  }
371  }
372 
373 
374  PHG4TruthInfoContainer::VtxRange vtxrange = truthinfo->GetVtxRange();
375  n_vertex = 0;
376  for (PHG4TruthInfoContainer::ConstVtxIterator iter = vtxrange.first; iter != vtxrange.second; ++iter) {
377  PHG4VtxPoint *vtx = iter->second;
378  if (vertex.find(vtx->get_id()) != vertex.end()) {
379  vertex_x[n_vertex] = vtx->get_x();
380  vertex_y[n_vertex] = vtx->get_y();
381  vertex_z[n_vertex] = vtx->get_z();
382  vertex_id[n_vertex] = vtx->get_id();
383  n_vertex++;
384  }
385  }
386 
387  _file->cd();
388  _tree->Fill();
389  _ievent++;
390 
391  return 0;
392 }
393 
394 int TruthCaloTree::End(PHCompositeNode *topNode)
395 {
396  if (_debug) std::cout<<"Writing File"<<std::endl;
397  _file->cd();
398  _file->Write();
399  _file->Close();
400 
401  return 0;
402 }