ROOTPWA
fillUdstDataIntoMassBins_example.C
Go to the documentation of this file.
1 
2 //
3 // Copyright 2010
4 //
5 // This file is part of rootpwa
6 //
7 // rootpwa is free software: you can redistribute it and/or modify
8 // it under the terms of the GNU General Public License as published by
9 // the Free Software Foundation, either version 3 of the License, or
10 // (at your option) any later version.
11 //
12 // rootpwa is distributed in the hope that it will be useful,
13 // but WITHOUT ANY WARRANTY; without even the implied warranty of
14 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15 // GNU General Public License for more details.
16 //
17 // You should have received a copy of the GNU General Public License
18 // along with rootpwa. If not, see <http://www.gnu.org/licenses/>.
19 //
21 //-------------------------------------------------------------------------
22 // File and Version Information:
23 // $Rev:: $: revision of last commit
24 // $Author:: $: author of last commit
25 // $Date:: $: date of last commit
26 //
27 // Description:
28 // example macro that reads uDST tree and generates mass bin
29 // directory structure with .root files in ROOTPWA format
30 // parts that depend on the analyzed final state and the uDST
31 // tree structure are marked accordingly
32 //
33 // to run this macro type
34 // root rootlogon.C fillUdstDataIntoMassBins_example.C+
35 //
36 //
37 // Author List:
38 // Stefan Pflueger TUM (original author)
39 //
40 //
41 //-------------------------------------------------------------------------
42 
43 
44 #include <iostream>
45 #include <iomanip>
46 #include <sstream>
47 #include <string>
48 #include <vector>
49 #include <sys/stat.h>
50 #include <sys/types.h>
51 #include <algorithm>
52 
53 #include <boost/progress.hpp>
54 
55 #include "TStopwatch.h"
56 #include "TFile.h"
57 #include "TTree.h"
58 #include "TChain.h"
59 #include "TClonesArray.h"
60 #include "TLorentzVector.h"
61 
62 #include "reportingUtils.hpp"
63 
64 
65 using namespace std;
66 using namespace boost;
67 using namespace rpwa;
68 
69 
70 // generic function that creates mass bin directory structure and
71 // returns array of event files and trees, one for each mass bin
72 bool
73 createMassBinFiles(vector<TFile*>& pwaFiles,
74  vector<TTree*>& pwaTrees,
75  const string& dirBaseName = "/tmp/",
76  const unsigned int nmbMassBins = 50,
77  const double massBinWidth = 50, // [MeV/c^2]
78  const double massRangeMin = 500, // [MeV/c^2]
79  const string& treeName = "rootPwaEvtTree")
80 {
81  printInfo << "creating mass bin directories/files in '" << dirBaseName << "':" << endl;
82  // cleanup
83  for (unsigned int i = 0; i < pwaTrees.size(); ++i)
84  if (pwaTrees[i])
85  delete pwaTrees[i];
86  pwaTrees.clear();
87  for (unsigned int i = 0; i < pwaFiles.size(); ++i)
88  if (pwaFiles[i]) {
89  pwaFiles[i]->Close();
90  delete pwaFiles[i];
91  }
92  pwaFiles.clear();
93  pwaFiles.resize(nmbMassBins, 0);
94  pwaTrees.resize(nmbMassBins, 0);
95  bool success = true;
96  for (unsigned int i = 0; i < nmbMassBins; ++i) {
97  const double binMin = massRangeMin + i * massBinWidth;
98  const double binMax = binMin + massBinWidth;
99  // create mass bin directory
100  stringstream n;
101  n << binMin << "." << binMax;
102  const string dirName = dirBaseName + "/" + n.str();
103  mkdir(dirName.c_str(), S_IRWXU | S_IRWXG); // create directory read/writable by owner and group
104  // create output file
105  const string pwaFileName = dirName + "/" + n.str() + ".root";
106  TFile* pwaFile = TFile::Open(pwaFileName.c_str(), "RECREATE");
107  if (not pwaFile or pwaFile->IsZombie()) {
108  printWarn << "problems creating file '" << pwaFileName << "'" << endl;
109  success = false;
110  } else {
111  pwaFiles[i] = pwaFile;
112  pwaTrees[i] = new TTree(treeName.c_str(), treeName.c_str());
113  if (not pwaTrees[i]) {
114  printWarn << "problems creating tree '" << treeName << "' " << "in file "
115  << "'" << pwaFileName << "'" << endl;
116  success = false;
117  } else {
118  pwaTrees[i]->SetDirectory(pwaFile);
119  printSucc << "created tree '" << treeName << "' in file " << "'" << pwaFileName
120  << "'" << endl;
121  }
122  }
123  }
124  return success;
125 }
126 
127 
128 // fills event data into correct mass bin
129 bool
130 writeEvent(vector<TTree*>& pwaTrees,
131  const TLorentzVector& beamLv,
132  // !!! <channel-dependent part> !!!
133  const TLorentzVector& piZero0,
134  const TLorentzVector& piZero1,
135  const TLorentzVector& piMinus,
136  // !!! </channel-dependent part> !!!
137  const double XMass, // [GeV/c^2]
138  const unsigned int nmbMassBins = 50,
139  const double massBinWidth = 50, // [MeV/c^2]
140  const double massRangeMin = 500, // [MeV/c^2]
141  const string& prodKinMomentaLeafName = "prodKinMomenta",
142  const string& decayKinMomentaLeafName = "decayKinMomenta",
143  const bool debug = false)
144 {
145 
146  const double mass = 1000 * XMass; // convert from GeV/c^2 to MeV/c^2
147  // make sure that mass is in range
148  if ((mass < massRangeMin) or (mass > (massRangeMin + nmbMassBins * massBinWidth)))
149  return false;
150  const unsigned int bin = (unsigned int) ((mass - massRangeMin) / massBinWidth);
151  if (not pwaTrees[bin]) {
152  printWarn << "null pointer for tree for mass bin [" << massRangeMin + bin * massBinWidth << ", "
153  << massRangeMin + (bin + 1) * massBinWidth << "]" << endl;
154  return false;
155  }
156  // fill tree
157  if (debug)
158  printDebug << "filling tree for bin " << bin << " = ["
159  << massRangeMin + bin * massBinWidth << ", "
160  << massRangeMin + (bin + 1) * massBinWidth << "] MeV/c^2" << endl;
161 
162  // create tree leafs
163  static TClonesArray* prodKinMomenta = new TClonesArray("TVector3");
164  static TClonesArray* decayKinMomenta = new TClonesArray("TVector3");
165 
166  // connect leaf variables to tree branches or create branches, if they don't exist yet
167  TTree* outTree = pwaTrees[bin];
168  if (outTree->SetBranchAddress(prodKinMomentaLeafName.c_str(), &prodKinMomenta) < 0) {
169  printWarn << "could not connect variable to branch '" << prodKinMomentaLeafName << "'. "
170  << "skipping." << endl;
171  return false;
172  }
173  if (outTree->SetBranchAddress(decayKinMomentaLeafName.c_str(), &decayKinMomenta) < 0) {
174  printWarn << "could not connect variable to branch '" << prodKinMomentaLeafName << "'. "
175  << "skipping." << endl;
176  return false;
177  }
178 
179  // clear arrays
180  prodKinMomenta->Clear ();
181  decayKinMomenta->Clear();
182 
183  // set leaf variables
184  // beam particle
185  new ((*prodKinMomenta)[0]) TVector3(beamLv.Vect());
186 
187  // !!! <channel-dependent part> !!!
188  // for target particle elastic scattering is assumed
189  // outgoing hadrons
190  new ((*decayKinMomenta)[0]) TVector3(piZero0.Vect());
191  new ((*decayKinMomenta)[1]) TVector3(piZero1.Vect());
192  new ((*decayKinMomenta)[2]) TVector3(piMinus.Vect());
193  // !!! </channel-dependent part> !!!
194 
195  // fill tree
196  outTree->Fill();
197  return true;
198 }
199 
200 
201 void
202 fillUdstDataIntoMassBins_example(const string& inFileNamePattern = "fillUdstDataIntoMassBins_example.root",
203  const string& dirName = "./test",
204  const long int maxNmbEvents = -1,
205  const unsigned int nmbMassBins = 50,
206  const double massBinWidth = 40, // [MeV/c^2]
207  const double massRangeMin = 500, // [MeV/c^2]
208  const string& uDstTreeName = "pwaDataTree",
209  const string& pwaTreeName = "rootPwaEvtTree",
210  const long int treeCacheSize = 25000000, // 25 MByte ROOT tree read cache
211  const bool debug = false)
212 {
213  const string prodKinPartNamesObjName = "prodKinParticles";
214  const string prodKinMomentaLeafName = "prodKinMomenta";
215  const string decayKinPartNamesObjName = "decayKinParticles";
216  const string decayKinMomentaLeafName = "decayKinMomenta";
217 
218  TStopwatch timer;
219  timer.Start();
220 
221  printInfo << "reading uDST file(s) '" << inFileNamePattern << "'" << endl
222  << " writing " << nmbMassBins << " mass bins in mass interval "
223  << "[" << massRangeMin << ", " << massRangeMin + nmbMassBins * massBinWidth << "] "
224  << "MeV/c^2 to '" << dirName << "'" << endl
225  << " reading uDST data from tree '" << uDstTreeName << "'" << endl
226  << " writing PWA data to tree '" << pwaTreeName << "'" << endl;
227 
228  // create chain and connect tree leaf variables to branches
229  TChain uDstChain(uDstTreeName.c_str());
230  if (uDstChain.Add(inFileNamePattern.c_str()) < 1)
231  printWarn << "no events in uDST input file(s) '" << inFileNamePattern << "'" << endl;
232  const long int nmbEventsUdstChain = uDstChain.GetEntries();
233  uDstChain.GetListOfFiles()->ls();
234 
235  // !!! <channel-dependent part> !!!
236  // connect tree leafs
237  TLorentzVector* photons[4] = {0, 0, 0, 0};
238  TLorentzVector* piMinus = 0;
239  TLorentzVector* beam = 0;
240  TLorentzVector* recoil = 0;
241  uDstChain.SetBranchAddress("gamma1", &(photons[0]));
242  uDstChain.SetBranchAddress("gamma2", &(photons[1]));
243  uDstChain.SetBranchAddress("gamma3", &(photons[2]));
244  uDstChain.SetBranchAddress("gamma4", &(photons[3]));
245  uDstChain.SetBranchAddress("pi_out", &piMinus);
246  uDstChain.SetBranchAddress("pi_in", &beam);
247  uDstChain.SetBranchAddress("proton", &recoil);
248  uDstChain.SetCacheSize(treeCacheSize);
249  uDstChain.AddBranchToCache("gamma1", true);
250  uDstChain.AddBranchToCache("gamma2", true);
251  uDstChain.AddBranchToCache("gamma3", true);
252  uDstChain.AddBranchToCache("gamma4", true);
253  uDstChain.AddBranchToCache("pi_out", true);
254  uDstChain.AddBranchToCache("pi_in", true);
255  uDstChain.AddBranchToCache("proton", true);
256  uDstChain.StopCacheLearningPhase();
257  // !!! </channel-dependent part> !!!
258 
259  // create directories and .root files
260  vector<TFile*> pwaDataFiles;
261  vector<TTree*> pwaDataTrees;
262  if (not createMassBinFiles(pwaDataFiles, pwaDataTrees, dirName, nmbMassBins, massBinWidth,
263  massRangeMin, pwaTreeName)) {
264  printErr << "there were problems creating the mass bin directories/files. aborting." << endl;
265  return;
266  }
267  printSucc << "created " << pwaDataFiles.size() << " directories/files" << endl;
268 
269  // write arrays with production and decay particle names to root files
270  {
271  TClonesArray prodKinPartNames ("TObjString", 1);
272  TClonesArray decayKinPartNames("TObjString", 3);
273 
274  // !!! <channel-dependent part> !!!
275  new (prodKinPartNames [0]) TObjString("pi-"); // beam particle
276  new (decayKinPartNames[0]) TObjString("pi0");
277  new (decayKinPartNames[1]) TObjString("pi0");
278  new (decayKinPartNames[2]) TObjString("pi-");
279  // !!! </channel-dependent part> !!!
280 
281  for (unsigned int i = 0; i < pwaDataFiles.size(); ++i) {
282  pwaDataFiles[i]->cd();
283  prodKinPartNames.Write (prodKinPartNamesObjName.c_str (), TObject::kSingleKey);
284  decayKinPartNames.Write(decayKinPartNamesObjName.c_str(), TObject::kSingleKey);
285  }
286  printSucc << "wrote particle name arrays to all files. "
287  << "beam = 'pi-', decay = {'pi0', 'pi0', 'pi-'}." << endl;
288  }
289 
290  // create tree leafs
291  {
292  TClonesArray* prodKinMomenta = new TClonesArray("TVector3");
293  TClonesArray* decayKinMomenta = new TClonesArray("TVector3");
294  const int splitLevel = 99;
295  const int bufSize = 256000;
296  for (unsigned int i = 0; i < pwaDataTrees.size(); ++i) {
297  pwaDataTrees[i]->Branch(prodKinMomentaLeafName.c_str(), "TClonesArray", &prodKinMomenta,
298  bufSize, splitLevel);
299  pwaDataTrees[i]->Branch(decayKinMomentaLeafName.c_str(), "TClonesArray", &decayKinMomenta,
300  bufSize, splitLevel);
301  }
302  printSucc << "created branches for all trees" << endl;
303  }
304 
305  // loop over events
306  const long int nmbEvents = ((maxNmbEvents > 0) ? maxNmbEvents : nmbEventsUdstChain);
307  printInfo << "writing events into mass bin files" << endl
308  << " looping over " << nmbEvents << " tree entries" << endl;
309  unsigned long int countEvWritten = 0;
310  progress_display progressIndicator(nmbEvents, cout, "");
311  for (long int eventIndex = 0; eventIndex < nmbEvents; ++eventIndex) {
312  ++progressIndicator;
313  if ((uDstChain.LoadTree(eventIndex) < 0) or (uDstChain.GetEntry(eventIndex) == 0)) {
314  printWarn << "error reading event " << eventIndex << " from tree. skipping." << endl;
315  continue;
316  }
317 
318  // !!! <channel-dependent part> !!!
319  // now just make some minor calculations before
320  // construct two pi0s
321  const TLorentzVector piZeros[2] = {*(photons[0]) + *(photons[1]), *(photons[2]) + *(photons[3])};
322  // construct intermediate state X
323  const TLorentzVector X = piZeros[0] + piZeros[1] + *piMinus;
324  // calculate t'
325  const double t = (*beam - X) * (*beam - X);
326  const double tPrime = fabs(t) - fabs((X.M() * X.M() - beam->M() * beam->M())*(X.M() * X.M() - beam->M() * beam->M())
327  / (4 * (beam->Vect() * beam->Vect())));
328  // cut on t'
329  if ((tPrime < 0.1) or (tPrime > 1.0))
330  continue;
331  // write out PWA data
332  const double fsEnergy = X.E() + recoil->E() - recoil->M(); // measured total energy of final state
333  const double scaleFactor = fsEnergy / beam->E();
334  const TLorentzVector beamScaled(beam->Vect() * scaleFactor, fsEnergy);
335  if (writeEvent(pwaDataTrees, beamScaled, piZeros[0], piZeros[1], *piMinus,
336  X.M(), nmbMassBins, massBinWidth, massRangeMin,
337  prodKinMomentaLeafName, decayKinMomentaLeafName, debug))
338  ++countEvWritten;
339  // !!! </channel-dependent part> !!!
340  }
341 
342  // write trees
343  long unsigned int countTreeEvents = 0;
344  for (unsigned int i = 0; i < nmbMassBins; ++i) {
345  pwaDataTrees[i]->GetCurrentFile()->Write();
346  long unsigned int nmbEvents = pwaDataTrees[i]->GetEntries();
347  printSucc << "written " << setw(10) << nmbEvents << " events to file " << "'"
348  << pwaDataTrees[i]->GetCurrentFile()->GetName() << "'" << endl;
349  countTreeEvents += nmbEvents;
350  pwaDataTrees[i]->GetCurrentFile()->Close();
351  }
352  pwaDataFiles.clear();
353  pwaDataTrees.clear();
354 
355  printInfo << "wrote " << min(countEvWritten, countTreeEvents)
356  << " out of " << nmbEvents << " events" <<endl;
357  timer.Stop();
358  printInfo << "this job consumed: ";
359  timer.Print();
360 }
Generated by   doxygen 1.8.1.2