Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
TreeReader.h
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file TreeReader.h
1 // This file is part of the Acts project.
2 //
3 // Copyright (C) 2021 CERN for the benefit of the Acts project
4 //
5 // This Source Code Form is subject to the terms of the Mozilla Public
6 // License, v. 2.0. If a copy of the MPL was not distributed with this
7 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
8 
9 #pragma once
10 
11 #include <array>
12 #include <string>
13 #include <vector>
14 
15 #include <TCanvas.h>
16 #include <TEfficiency.h>
17 #include <TFile.h>
18 #include <TTree.h>
19 
20 struct RecoTrackInfo {
21  double eta;
22  double pt;
23  unsigned int nMajorityHits;
24  unsigned int nMeasurements;
25 };
26 
27 struct ParticleInfo {
28  ULong64_t particleId = 0;
29  double eta = 0;
30  double p = 0;
31  double pt = 0;
32  UShort_t nHits = 0;
33 };
34 
37 struct TreeReader {
38  // The constructor
39  TreeReader(TTree* tree_) : tree(tree_) {}
40 
41  // Get entry
42  void getEntry(unsigned int i) const {
43  if (entryNumbers.size() > i) {
44  tree->GetEntry(entryNumbers[i]);
45  } else {
46  tree->GetEntry(i);
47  }
48  };
49 
50  // The tree being read
51  TTree* tree = nullptr;
52 
53  protected:
56  std::vector<long long> entryNumbers = {};
57 };
58 
62 struct TrackStatesReader : public TreeReader {
63  // Delete the default constructor
64  TrackStatesReader() = delete;
65 
66  // The constructor
67  TrackStatesReader(TTree* tree_, bool sortEvents) : TreeReader(tree_) {
68  tree->SetBranchAddress("event_nr", &eventId);
69  tree->SetBranchAddress("eLOC0_prt", &LOC0_prt);
70  tree->SetBranchAddress("eLOC1_prt", &LOC1_prt);
71  tree->SetBranchAddress("ePHI_prt", &PHI_prt);
72  tree->SetBranchAddress("eTHETA_prt", &THETA_prt);
73  tree->SetBranchAddress("eQOP_prt", &QOP_prt);
74  tree->SetBranchAddress("eT_prt", &T_prt);
75  tree->SetBranchAddress("eLOC0_flt", &LOC0_flt);
76  tree->SetBranchAddress("eLOC1_flt", &LOC1_flt);
77  tree->SetBranchAddress("ePHI_flt", &PHI_flt);
78  tree->SetBranchAddress("eTHETA_flt", &THETA_flt);
79  tree->SetBranchAddress("eQOP_flt", &QOP_flt);
80  tree->SetBranchAddress("eT_flt", &T_flt);
81  tree->SetBranchAddress("eLOC0_smt", &LOC0_smt);
82  tree->SetBranchAddress("eLOC1_smt", &LOC1_smt);
83  tree->SetBranchAddress("ePHI_smt", &PHI_smt);
84  tree->SetBranchAddress("eTHETA_smt", &THETA_smt);
85  tree->SetBranchAddress("eQOP_smt", &QOP_smt);
86  tree->SetBranchAddress("eT_smt", &T_smt);
87 
88  tree->SetBranchAddress("res_eLOC0_prt", &res_LOC0_prt);
89  tree->SetBranchAddress("res_eLOC1_prt", &res_LOC1_prt);
90  tree->SetBranchAddress("res_ePHI_prt", &res_PHI_prt);
91  tree->SetBranchAddress("res_eTHETA_prt", &res_THETA_prt);
92  tree->SetBranchAddress("res_eQOP_prt", &res_QOP_prt);
93  tree->SetBranchAddress("res_eT_prt", &res_T_prt);
94  tree->SetBranchAddress("res_eLOC0_flt", &res_LOC0_flt);
95  tree->SetBranchAddress("res_eLOC1_flt", &res_LOC1_flt);
96  tree->SetBranchAddress("res_ePHI_flt", &res_PHI_flt);
97  tree->SetBranchAddress("res_eTHETA_flt", &res_THETA_flt);
98  tree->SetBranchAddress("res_eQOP_flt", &res_QOP_flt);
99  tree->SetBranchAddress("res_eT_flt", &res_T_flt);
100  tree->SetBranchAddress("res_eLOC0_smt", &res_LOC0_smt);
101  tree->SetBranchAddress("res_eLOC1_smt", &res_LOC1_smt);
102  tree->SetBranchAddress("res_ePHI_smt", &res_PHI_smt);
103  tree->SetBranchAddress("res_eTHETA_smt", &res_THETA_smt);
104  tree->SetBranchAddress("res_eQOP_smt", &res_QOP_smt);
105  tree->SetBranchAddress("res_eT_smt", &res_T_smt);
106 
107  tree->SetBranchAddress("pull_eLOC0_prt", &pull_LOC0_prt);
108  tree->SetBranchAddress("pull_eLOC1_prt", &pull_LOC1_prt);
109  tree->SetBranchAddress("pull_ePHI_prt", &pull_PHI_prt);
110  tree->SetBranchAddress("pull_eTHETA_prt", &pull_THETA_prt);
111  tree->SetBranchAddress("pull_eQOP_prt", &pull_QOP_prt);
112  tree->SetBranchAddress("pull_eT_prt", &pull_T_prt);
113  tree->SetBranchAddress("pull_eLOC0_flt", &pull_LOC0_flt);
114  tree->SetBranchAddress("pull_eLOC1_flt", &pull_LOC1_flt);
115  tree->SetBranchAddress("pull_ePHI_flt", &pull_PHI_flt);
116  tree->SetBranchAddress("pull_eTHETA_flt", &pull_THETA_flt);
117  tree->SetBranchAddress("pull_eQOP_flt", &pull_QOP_flt);
118  tree->SetBranchAddress("pull_eT_flt", &pull_T_flt);
119  tree->SetBranchAddress("pull_eLOC0_smt", &pull_LOC0_smt);
120  tree->SetBranchAddress("pull_eLOC1_smt", &pull_LOC1_smt);
121  tree->SetBranchAddress("pull_ePHI_smt", &pull_PHI_smt);
122  tree->SetBranchAddress("pull_eTHETA_smt", &pull_THETA_smt);
123  tree->SetBranchAddress("pull_eQOP_smt", &pull_QOP_smt);
124  tree->SetBranchAddress("pull_eT_smt", &pull_T_smt);
125 
126  tree->SetBranchAddress("g_x_prt", &g_x_prt);
127  tree->SetBranchAddress("g_y_prt", &g_y_prt);
128  tree->SetBranchAddress("g_z_prt", &g_z_prt);
129  tree->SetBranchAddress("g_x_flt", &g_x_flt);
130  tree->SetBranchAddress("g_y_flt", &g_y_flt);
131  tree->SetBranchAddress("g_z_flt", &g_z_flt);
132  tree->SetBranchAddress("g_x_smt", &g_x_smt);
133  tree->SetBranchAddress("g_y_smt", &g_y_smt);
134  tree->SetBranchAddress("g_z_smt", &g_z_smt);
135 
136  tree->SetBranchAddress("nStates", &nStates);
137  tree->SetBranchAddress("nMeasurements", &nMeasurements);
138  tree->SetBranchAddress("volume_id", &volume_id);
139  tree->SetBranchAddress("layer_id", &layer_id);
140  tree->SetBranchAddress("module_id", &module_id);
141  tree->SetBranchAddress("predicted", &predicted);
142  tree->SetBranchAddress("filtered", &filtered);
143  tree->SetBranchAddress("smoothed", &smoothed);
144 
145  // It's not necessary if you just need to read one file, but please do it to
146  // synchronize events if multiple root files are read
147  if (sortEvents) {
148  entryNumbers.resize(tree->GetEntries());
149  tree->Draw("event_nr", "", "goff");
150  // Sort to get the entry numbers of the ordered events
151  TMath::Sort(tree->GetEntries(), tree->GetV1(), entryNumbers.data(),
152  false);
153  }
154  }
155 
156  // The variables
157  uint32_t eventId = 0;
158  std::vector<float>* LOC0_prt =
159  new std::vector<float>;
160  std::vector<float>* LOC1_prt =
161  new std::vector<float>;
162  std::vector<float>* PHI_prt =
163  new std::vector<float>;
164  std::vector<float>* THETA_prt =
165  new std::vector<float>;
166  std::vector<float>* QOP_prt =
167  new std::vector<float>;
168  std::vector<float>* T_prt =
169  new std::vector<float>;
170  std::vector<float>* LOC0_flt =
171  new std::vector<float>;
172  std::vector<float>* LOC1_flt =
173  new std::vector<float>;
174  std::vector<float>* PHI_flt =
175  new std::vector<float>;
176  std::vector<float>* THETA_flt =
177  new std::vector<float>;
178  std::vector<float>* QOP_flt =
179  new std::vector<float>;
180  std::vector<float>* T_flt = new std::vector<float>;
181  std::vector<float>* LOC0_smt =
182  new std::vector<float>;
183  std::vector<float>* LOC1_smt =
184  new std::vector<float>;
185  std::vector<float>* PHI_smt =
186  new std::vector<float>;
187  std::vector<float>* THETA_smt =
188  new std::vector<float>;
189  std::vector<float>* QOP_smt =
190  new std::vector<float>;
191  std::vector<float>* T_smt = new std::vector<float>;
192 
193  std::vector<float>* res_LOC0_prt =
194  new std::vector<float>;
195  std::vector<float>* res_LOC1_prt =
196  new std::vector<float>;
197  std::vector<float>* res_PHI_prt =
198  new std::vector<float>;
199  std::vector<float>* res_THETA_prt =
200  new std::vector<float>;
201  std::vector<float>* res_QOP_prt =
202  new std::vector<float>;
203  std::vector<float>* res_T_prt =
204  new std::vector<float>;
205  std::vector<float>* res_LOC0_flt =
206  new std::vector<float>;
207  std::vector<float>* res_LOC1_flt =
208  new std::vector<float>;
209  std::vector<float>* res_PHI_flt =
210  new std::vector<float>;
211  std::vector<float>* res_THETA_flt =
212  new std::vector<float>;
213  std::vector<float>* res_QOP_flt =
214  new std::vector<float>;
215  std::vector<float>* res_T_flt =
216  new std::vector<float>;
217  std::vector<float>* res_LOC0_smt =
218  new std::vector<float>;
219  std::vector<float>* res_LOC1_smt =
220  new std::vector<float>;
221  std::vector<float>* res_PHI_smt =
222  new std::vector<float>;
223  std::vector<float>* res_THETA_smt =
224  new std::vector<float>;
225  std::vector<float>* res_QOP_smt =
226  new std::vector<float>;
227  std::vector<float>* res_T_smt =
228  new std::vector<float>;
229 
230  std::vector<float>* pull_LOC0_prt =
231  new std::vector<float>;
232  std::vector<float>* pull_LOC1_prt =
233  new std::vector<float>;
234  std::vector<float>* pull_PHI_prt =
235  new std::vector<float>;
236  std::vector<float>* pull_THETA_prt =
237  new std::vector<float>;
238  std::vector<float>* pull_QOP_prt =
239  new std::vector<float>;
240  std::vector<float>* pull_T_prt =
241  new std::vector<float>;
242  std::vector<float>* pull_LOC0_flt =
243  new std::vector<float>;
244  std::vector<float>* pull_LOC1_flt =
245  new std::vector<float>;
246  std::vector<float>* pull_PHI_flt =
247  new std::vector<float>;
248  std::vector<float>* pull_THETA_flt =
249  new std::vector<float>;
250  std::vector<float>* pull_QOP_flt =
251  new std::vector<float>;
252  std::vector<float>* pull_T_flt =
253  new std::vector<float>;
254  std::vector<float>* pull_LOC0_smt =
255  new std::vector<float>;
256  std::vector<float>* pull_LOC1_smt =
257  new std::vector<float>;
258  std::vector<float>* pull_PHI_smt =
259  new std::vector<float>;
260  std::vector<float>* pull_THETA_smt =
261  new std::vector<float>;
262  std::vector<float>* pull_QOP_smt =
263  new std::vector<float>;
264  std::vector<float>* pull_T_smt =
265  new std::vector<float>;
266 
267  std::vector<float>* g_x_prt = new std::vector<float>;
268  std::vector<float>* g_y_prt = new std::vector<float>;
269  std::vector<float>* g_z_prt = new std::vector<float>;
270  std::vector<float>* g_x_flt = new std::vector<float>;
271  std::vector<float>* g_y_flt = new std::vector<float>;
272  std::vector<float>* g_z_flt = new std::vector<float>;
273  std::vector<float>* g_x_smt = new std::vector<float>;
274  std::vector<float>* g_y_smt = new std::vector<float>;
275  std::vector<float>* g_z_smt = new std::vector<float>;
276 
277  std::vector<int>* volume_id = new std::vector<int>;
278  std::vector<int>* layer_id = new std::vector<int>;
279  std::vector<int>* module_id = new std::vector<int>;
280 
281  std::vector<bool>* predicted = new std::vector<bool>;
282  std::vector<bool>* filtered = new std::vector<bool>;
283  std::vector<bool>* smoothed = new std::vector<bool>;
284 
285  unsigned int nStates = 0, nMeasurements = 0;
286 };
287 
291 struct TrackSummaryReader : public TreeReader {
292  // Delete the default constructor
293  TrackSummaryReader() = delete;
294 
295  // The constructor
296  TrackSummaryReader(TTree* tree_, bool sortEvents) : TreeReader(tree_) {
297  tree->SetBranchAddress("event_nr", &eventId);
298  tree->SetBranchAddress("nStates", &nStates);
299  tree->SetBranchAddress("nMeasurements", &nMeasurements);
300  tree->SetBranchAddress("nOutliers", &nOutliers);
301  tree->SetBranchAddress("nHoles", &nHoles);
302  tree->SetBranchAddress("chi2Sum", &chi2Sum);
303  tree->SetBranchAddress("measurementChi2", &measurementChi2);
304  tree->SetBranchAddress("NDF", &NDF);
305  tree->SetBranchAddress("measurementVolume", &measurementVolume);
306  tree->SetBranchAddress("measurementLayer", &measurementLayer);
307  tree->SetBranchAddress("outlierVolume", &outlierVolume);
308  tree->SetBranchAddress("outlierLayer", &outlierLayer);
309  tree->SetBranchAddress("nMajorityHits", &nMajorityHits);
310  tree->SetBranchAddress("nSharedHits", &nSharedHits);
311  tree->SetBranchAddress("majorityParticleId", &majorityParticleId);
312 
313  tree->SetBranchAddress("hasFittedParams", &hasFittedParams);
314 
315  tree->SetBranchAddress("t_theta", &t_theta);
316  tree->SetBranchAddress("t_phi", &t_phi);
317  tree->SetBranchAddress("t_eta", &t_eta);
318  tree->SetBranchAddress("t_p", &t_p);
319  tree->SetBranchAddress("t_pT", &t_pT);
320  tree->SetBranchAddress("t_d0", &t_d0);
321  tree->SetBranchAddress("t_z0", &t_z0);
322  tree->SetBranchAddress("t_charge", &t_charge);
323  tree->SetBranchAddress("t_time", &t_time);
324 
325  tree->SetBranchAddress("eLOC0_fit", &eLOC0_fit);
326  tree->SetBranchAddress("eLOC1_fit", &eLOC1_fit);
327  tree->SetBranchAddress("ePHI_fit", &ePHI_fit);
328  tree->SetBranchAddress("eTHETA_fit", &eTHETA_fit);
329  tree->SetBranchAddress("eQOP_fit", &eQOP_fit);
330  tree->SetBranchAddress("eT_fit", &eT_fit);
331  tree->SetBranchAddress("err_eLOC0_fit", &err_eLOC0_fit);
332  tree->SetBranchAddress("err_eLOC1_fit", &err_eLOC1_fit);
333  tree->SetBranchAddress("err_ePHI_fit", &err_ePHI_fit);
334  tree->SetBranchAddress("err_eTHETA_fit", &err_eTHETA_fit);
335  tree->SetBranchAddress("err_eQOP_fit", &err_eQOP_fit);
336  tree->SetBranchAddress("err_eT_fit", &err_eT_fit);
337 
338  // It's not necessary if you just need to read one file, but please do it to
339  // synchronize events if multiple root files are read
340  if (sortEvents) {
341  entryNumbers.resize(tree->GetEntries());
342  tree->Draw("event_nr", "", "goff");
343  // Sort to get the entry numbers of the ordered events
344  TMath::Sort(tree->GetEntries(), tree->GetV1(), entryNumbers.data(),
345  false);
346  }
347  }
348 
349  // The variables
350  uint32_t eventId = 0;
351  std::vector<unsigned int>* nStates = new std::vector<unsigned int>;
352  std::vector<unsigned int>* nMeasurements = new std::vector<unsigned int>;
353  std::vector<unsigned int>* nOutliers = new std::vector<unsigned int>;
354  std::vector<unsigned int>* nHoles = new std::vector<unsigned int>;
355  std::vector<unsigned int>* nSharedHits = new std::vector<unsigned int>;
356  std::vector<float>* chi2Sum = new std::vector<float>;
357  std::vector<unsigned int>* NDF = new std::vector<unsigned int>;
358  std::vector<std::vector<double>>* measurementChi2 =
359  new std::vector<std::vector<double>>;
360  std::vector<std::vector<double>>* outlierChi2 =
361  new std::vector<std::vector<double>>;
362  std::vector<std::vector<double>>* measurementVolume =
363  new std::vector<std::vector<double>>;
364  std::vector<std::vector<double>>* measurementLayer =
365  new std::vector<std::vector<double>>;
366  std::vector<std::vector<double>>* outlierVolume =
367  new std::vector<std::vector<double>>;
368  std::vector<std::vector<double>>* outlierLayer =
369  new std::vector<std::vector<double>>;
370  std::vector<unsigned int>* nMajorityHits = new std::vector<unsigned int>;
371  std::vector<uint64_t>* majorityParticleId = new std::vector<uint64_t>;
372 
373  std::vector<bool>* hasFittedParams = new std::vector<bool>;
374 
375  // True parameters
376  std::vector<float>* t_d0 = new std::vector<float>;
377  std::vector<float>* t_z0 = new std::vector<float>;
378  std::vector<float>* t_phi = new std::vector<float>;
379  std::vector<float>* t_theta = new std::vector<float>;
380  std::vector<float>* t_eta = new std::vector<float>;
381  std::vector<float>* t_p = new std::vector<float>;
382  std::vector<float>* t_pT = new std::vector<float>;
383  std::vector<float>* t_time = new std::vector<float>;
384  std::vector<int>* t_charge = new std::vector<int>;
385 
386  // Estimated parameters
387  std::vector<float>* eLOC0_fit = new std::vector<float>;
388  std::vector<float>* eLOC1_fit = new std::vector<float>;
389  std::vector<float>* ePHI_fit = new std::vector<float>;
390  std::vector<float>* eTHETA_fit = new std::vector<float>;
391  std::vector<float>* eQOP_fit = new std::vector<float>;
392  std::vector<float>* eT_fit = new std::vector<float>;
393 
394  std::vector<float>* err_eLOC0_fit = new std::vector<float>;
395  std::vector<float>* err_eLOC1_fit = new std::vector<float>;
396  std::vector<float>* err_ePHI_fit = new std::vector<float>;
397  std::vector<float>* err_eTHETA_fit = new std::vector<float>;
398  std::vector<float>* err_eQOP_fit = new std::vector<float>;
399  std::vector<float>* err_eT_fit = new std::vector<float>;
400 };
401 
405 struct ParticleReader : public TreeReader {
406  // Delete the default constructor
407  ParticleReader() = delete;
408 
409  // The constructor
410  ParticleReader(TTree* tree_, bool sortEvents) : TreeReader(tree_) {
411  tree->SetBranchAddress("event_id", &eventId);
412  tree->SetBranchAddress("particle_id", &particleId);
413  tree->SetBranchAddress("particle_type", &particleType);
414  tree->SetBranchAddress("vx", &vx);
415  tree->SetBranchAddress("vy", &vy);
416  tree->SetBranchAddress("vz", &vz);
417  tree->SetBranchAddress("vt", &vt);
418  tree->SetBranchAddress("px", &px);
419  tree->SetBranchAddress("py", &py);
420  tree->SetBranchAddress("pz", &pz);
421  tree->SetBranchAddress("m", &m);
422  tree->SetBranchAddress("q", &q);
423  tree->SetBranchAddress("nhits", &nHits);
424  tree->SetBranchAddress("ntracks", &nTracks);
425  tree->SetBranchAddress("ntracks_majority", &nTracksMajority);
426 
427  // It's not necessary if you just need to read one file, but please do it to
428  // synchronize events if multiple root files are read
429  if (sortEvents) {
430  entryNumbers.resize(tree->GetEntries());
431  tree->Draw("event_id", "", "goff");
432  // Sort to get the entry numbers of the ordered events
433  TMath::Sort(tree->GetEntries(), tree->GetV1(), entryNumbers.data(),
434  false);
435  }
436  }
437 
438  // Get all the particles with requested event id
439  std::vector<ParticleInfo> getParticles(const uint32_t& eventNumber) const {
440  // Find the start entry and the batch size for this event
441  std::string eventNumberStr = std::to_string(eventNumber);
442  std::string findStartEntry = "event_id<" + eventNumberStr;
443  std::string findParticlesSize = "event_id==" + eventNumberStr;
444  size_t startEntry = tree->GetEntries(findStartEntry.c_str());
445  size_t nParticles = tree->GetEntries(findParticlesSize.c_str());
446  if (nParticles == 0) {
447  throw std::invalid_argument(
448  "No particles found. Please check the input file.");
449  }
450  std::vector<ParticleInfo> particles;
451  particles.reserve(nParticles);
452  for (unsigned int i = 0; i < nParticles; ++i) {
453  getEntry(startEntry + i);
454  auto pt = std::hypot(px, py);
455  auto p = std::hypot(pt, pz);
456  auto eta = std::atanh(pz / p * 1.);
457  particles.push_back({particleId, eta, p, pt, nHits});
458  }
459  return particles;
460  }
461 
462  // The variables
463  ULong64_t eventId = 0;
464  ULong64_t particleId = 0;
465  Int_t particleType = 0;
466  float vx = 0, vy = 0, vz = 0;
467  float vt = 0;
468  float px = 0, py = 0, pz = 0;
469  float m = 0;
470  float q = 0;
471  UShort_t nHits = 0;
472  UShort_t nTracks = 0;
473  UShort_t nTracksMajority = 0;
474 };