ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/MitAna/DataUtil/src/TreeWriter.cc
Revision: 1.6
Committed: Tue Jun 24 14:05:03 2008 UTC (16 years, 10 months ago) by loizides
Content type: text/plain
Branch: MAIN
Changes since 1.5: +3 -2 lines
Log Message:
Cosmetics.

File Contents

# Content
1 // $Id: TreeWriter.cc,v 1.5 2008/06/17 08:22:41 loizides Exp $
2
3 #include "MitAna/DataUtil/interface/TreeWriter.h"
4
5 #include <Riostream.h>
6 #include <TObject.h>
7 #include <TSystem.h>
8 #include <TProcessID.h>
9
10 #include "MitAna/DataUtil/interface/Debug.h"
11
12 using namespace mithep;
13
14 ClassImp(mithep::TreeWriter)
15
16 //__________________________________________________________________________________________________
17 TreeWriter::TreeWriter(const char *tname, Bool_t doreset) :
18 TNamed(tname,Form("%s written by mithep::TreeWriter", tname)),
19 fBaseURL("."),
20 fPrefix("mithep"),
21 fFileNumber(0),
22 fCompressLevel(9),
23 fDefBrSize(64*1024),
24 fDefSL(99),
25 fMaxSize((Long64_t)(0.99 * TTree::GetMaxTreeSize())),
26 fkMinFreeSpace(1024*1024),
27 fkMinAvgSize(10*1024),
28 fEvtObjNum(-1),
29 fIsInit(kFALSE),
30 fDoObjNumReset(doreset),
31 fFile(0),
32 fTrees(0)
33 {
34 // Constructor.
35
36 fTrees.SetOwner();
37 }
38
39 //__________________________________________________________________________________________________
40 TreeWriter::~TreeWriter()
41 {
42 // Destructor.
43
44 if (fIsInit) {
45 CloseFile();
46 }
47
48 TDirectory::TContext context(0);
49 fTrees.Clear();
50 }
51
52 //__________________________________________________________________________________________________
53 void TreeWriter::AddBranch(const char *name, const char *cname,
54 void *obj, Int_t bsize, Int_t level)
55 {
56 // Add branch with name "name" into tree with name "tname" and set its address
57 // to object pointer for class name "cname" using a given buffer size and splitlevel.
58
59 MyTree *t = AddOrGetMyTree(GetName());
60 t->Bronch(name, cname, obj, bsize, level);
61 }
62
63
64 //__________________________________________________________________________________________________
65 void TreeWriter::AddBranch(const char *name, void *obj, Int_t bsize, Int_t level)
66 {
67 // Add branch with name "name" into tree with name "tname" and set its address
68 // to object pointer using a given buffer size and splitlevel.
69
70 AddBranch(name, CName(obj), obj, bsize, level);
71 }
72
73 //__________________________________________________________________________________________________
74 void TreeWriter::AddBranch(const char *name, const char *cname,
75 void *obj, Int_t bsize)
76 {
77 // Add branch with name "name" into tree with name "tname" and set its address
78 // to object pointer for class name "cname" using a given buffer size and default splitlevel.
79
80 MyTree *t = AddOrGetMyTree(GetName());
81 t->Bronch(name, cname, obj, bsize, fDefSL);
82 }
83
84 //__________________________________________________________________________________________________
85 void TreeWriter::AddBranch(const char *name, void *obj, Int_t bsize)
86 {
87 // Add branch with name "name" into tree with name "tname" and set its address
88 // to object pointer using a given buffer size and default splitlevel.
89
90 AddBranch(name, CName(obj), obj, bsize);
91 }
92
93 //__________________________________________________________________________________________________
94 void TreeWriter::AddBranch(const char *name, const char *cname,
95 void *obj)
96 {
97 // Add branch with name "name" into tree with name "tname" and set its address
98 // to object pointer for class name "cname" using a default buffer size and splitlevel.
99
100 MyTree *t = AddOrGetMyTree(GetName());
101 t->Bronch(name, cname, obj, fDefBrSize, fDefSL);
102 }
103
104 //__________________________________________________________________________________________________
105 void TreeWriter::AddBranch(const char *name, void *obj)
106 {
107 // Add branch with name "name" into tree with name "tname" and set its address
108 // to object pointer using a default buffer size and splitlevel.
109
110 AddBranch(name, CName(obj), obj);
111 }
112
113 //__________________________________________________________________________________________________
114 void TreeWriter::AddBranchToTree(const char *tname, const char *name, const char *cname,
115 void *obj, Int_t bsize, Int_t level)
116 {
117 // Add branch with name "name" into tree with name "tname" and set its address
118 // to object pointer for class name "cname" using a given buffer size and splitlevel.
119
120 MyTree *t = AddOrGetMyTree(tname);
121 t->Bronch(name, cname, obj, bsize, level);
122 }
123
124 //__________________________________________________________________________________________________
125 void TreeWriter::AddBranchToTree(const char *tname, const char *name, void *obj,
126 Int_t bsize, Int_t level)
127 {
128 // Add branch with name "name" into tree with name "tname" and set its address
129 // to object pointer using a given buffer size and splitlevel.
130
131 AddBranchToTree(tname, name, CName(obj), obj, bsize, level);
132 }
133
134 //__________________________________________________________________________________________________
135 void TreeWriter::AddBranchToTree(const char *tname, const char *name, const char *cname,
136 void *obj, Int_t bsize)
137 {
138 // Add branch with name "name" into tree with name "tname" and set its address
139 // to object pointer for class name "cname" using a given buffer size and default splitlevel.
140
141 MyTree *t = AddOrGetMyTree(tname);
142 t->Bronch(name, cname, obj, bsize, fDefSL);
143 }
144
145 //__________________________________________________________________________________________________
146 void TreeWriter::AddBranchToTree(const char *tname, const char *name, void *obj,
147 Int_t bsize)
148 {
149 // Add branch with name "name" into tree with name "tname" and set its address
150 // to object pointer using a given buffer size and default splitlevel.
151
152 AddBranchToTree(tname, name, CName(obj), obj, bsize);
153 }
154
155 //__________________________________________________________________________________________________
156 void TreeWriter::AddBranchToTree(const char *tname, const char *name, const char *cname,
157 void *obj)
158 {
159 // Add branch with name "name" into tree with name "tname" and set its address
160 // to object pointer for class name "cname" using a default buffer size and splitlevel.
161
162 MyTree *t = AddOrGetMyTree(tname);
163 t->Bronch(name, cname, obj, fDefBrSize, fDefSL);
164 }
165
166 //__________________________________________________________________________________________________
167 void TreeWriter::AddBranchToTree(const char *tname, const char *name, void *obj)
168 {
169 // Add branch with name "name" into tree with name "tname" and set its address
170 // to object pointer for class name "cname" using a default buffer size and splitlevel.
171
172 AddBranchToTree(tname, name, CName(obj), obj);
173 }
174
175 //__________________________________________________________________________________________________
176 MyTree *TreeWriter::AddOrGetMyTree(const char *tn)
177 {
178 // Add new tree if not present in array of trees or return
179 // present tree.
180
181 MyTree *tree = dynamic_cast<MyTree*>(fTrees.FindObject(tn));
182 if (tree)
183 return tree;
184
185 TDirectory::TContext context(fFile);
186 tree = new MyTree(tn, tn);
187 tree->SetDirectory(fFile);
188 fTrees.AddLast(tree);
189 return tree;
190 }
191
192 //__________________________________________________________________________________________________
193 Bool_t TreeWriter::BeginEvent(Bool_t doreset)
194 {
195 // Prepare for the next event. If doreset or fDoObjNumReset is kTRUE
196 // store the current object number to be reset in FillEvent().
197
198 if (!fIsInit) {
199 OpenFile();
200 }
201
202 if (doreset || fDoObjNumReset) {
203 fEvtObjNum = TProcessID::GetObjectCount();
204 }
205
206 return kTRUE;
207 }
208
209 //__________________________________________________________________________________________________
210 void TreeWriter::CloseFile()
211 {
212 // Write tree(s) and close file.
213
214 if (!fIsInit) {
215 Fatal("CloseFile", "File was not opened, call OpenFile() first!");
216 return;
217 }
218
219 TDirectory::TContext context(fFile); // cd fFile &&
220 // automatically restore gDirectory
221
222 for (Int_t i=0;i<fTrees.GetEntries();++i) {
223 MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
224 mt->Write(mt->GetName(),TObject::kOverwrite);
225 mt->Reset();
226 mt->SetDirectory(0);
227 }
228
229 fFile->Close();
230 delete fFile;
231 fFile = 0;
232
233 fIsInit = kFALSE;
234 fFileNumber++;
235 }
236
237 //__________________________________________________________________________________________________
238 const char *TreeWriter::CName(void *obj) const
239 {
240 // Dereference void* pointer into TObject* pointer
241
242 TObject *tobj = dynamic_cast<TObject*>(*(TObject**)obj);
243 if (tobj==0) {
244 Fatal("ClassName", "Given void* ptr can not be dereferenced into TObject*");
245 }
246 return tobj->ClassName();
247 }
248
249 //__________________________________________________________________________________________________
250 Bool_t TreeWriter::EndEvent(Bool_t doreset)
251 {
252 // Store the event in the tree. If doreset or fDoObjNumReset is kTRUE
253 // restore the stored object number at the time BeginEvent(kTRUE)
254 // was called.
255
256 if (!fIsInit) {
257 Fatal("EndEvent", "File is not open, did you call BeginEvent?");
258 return kFALSE;
259 }
260
261 Int_t r = 0;
262 for (Int_t i=0;i<fTrees.GetEntries();++i) {
263 MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
264 if (mt->GetAutoFill()==0)
265 continue;
266 r += mt->Fill();
267 }
268
269 if (IsFull())
270 CloseFile();
271
272 if (doreset || fDoObjNumReset) {
273 if (fEvtObjNum<0) {
274 Error("EndEvent", "Object counter is zero. Did you call BeginEvent(kTRUE)?");
275 } else {
276 // Reset the TRef table. keep it from growing with each event (see doc)
277 TProcessID::SetObjectCount(fEvtObjNum);
278 }
279 }
280
281 return (r >= 0);
282 }
283
284 //-------------------------------------------------------------------------------------------------
285 Long64_t TreeWriter::GetEntries(const char *tn) const
286 {
287 // Return entries of tree with given name. If no tree is given, return sum of entries
288 // of all trees.
289
290 if (fTrees.GetEntries()==0) return -1;
291
292 if (tn) {
293 const TTree *mt=GetTree(tn);
294 if (mt)
295 return mt->GetEntries();
296 else
297 return -1;
298 }
299
300 Long64_t ret = 0;
301 for (Int_t i=0;i<fTrees.GetEntries();++i) {
302 const MyTree *mt = static_cast<const MyTree*>(fTrees.At(i));
303 ret += mt->GetEntries();
304 }
305 return ret;
306 }
307
308 //-------------------------------------------------------------------------------------------------
309 MyTree *mithep::TreeWriter::GetMyTree(const char *tn)
310 {
311 // Return MyTree with given name from array.
312
313 if (fTrees.GetEntries()==0)
314 return 0;
315
316 TObject *obj = 0;
317 if (tn==0) {
318 obj = fTrees.At(0);
319 } else {
320 obj = fTrees.FindObject(tn);
321 }
322
323 if (obj)
324 return static_cast<MyTree*>(obj);
325 return 0;
326 }
327
328 //-------------------------------------------------------------------------------------------------
329 const TTree *mithep::TreeWriter::GetTree(const char *tn) const
330 {
331 // Return TTree with given name from array.
332
333 if (fTrees.GetEntries()==0)
334 return 0;
335
336 TObject *obj = 0;
337 if (tn==0) {
338 obj = fTrees.At(0);
339 } else {
340 obj = fTrees.FindObject(tn);
341 }
342
343 if (obj)
344 return dynamic_cast<const TTree*>(obj);
345 return 0;
346 }
347
348 //-------------------------------------------------------------------------------------------------
349 TTree *mithep::TreeWriter::GetTree(const char *tn)
350 {
351 // Return TTree with given name from array.
352
353 if (fTrees.GetEntries()==0)
354 return 0;
355
356 TObject *obj = 0;
357 if (tn==0) {
358 obj = fTrees.At(0);
359 } else {
360 obj = fTrees.FindObject(tn);
361 }
362
363 if (obj)
364 return dynamic_cast<TTree*>(obj);
365 return 0;
366 }
367
368 //__________________________________________________________________________________________________
369 Bool_t TreeWriter::IsFull() const
370 {
371 // Check if the maximum file size has been reached.
372
373 Long64_t entries = GetEntries();
374 if (entries < 1) return kFALSE;
375
376 Long64_t avgSize = GetFileSize() / entries;
377
378 if (avgSize < fkMinAvgSize)
379 avgSize = fkMinAvgSize;
380
381 return (GetFileSize() + avgSize + fkMinFreeSpace) > fMaxSize;
382 }
383
384 //__________________________________________________________________________________________________
385 void TreeWriter::OpenFile()
386 {
387 // Open the file and attach the tree.
388
389 if (fIsInit) {
390 Fatal("OpenFile", "File is already open, call CloseFile first!");
391 return;
392 }
393
394 TDirectory::TContext context(0);
395
396 TString pathname=GetFullName();
397 gSystem->ExpandPathName(pathname);
398
399 fFile = TFile::Open(pathname, "RECREATE");
400 if (fFile == 0) {
401 Fatal("OpenFile", "Could not open file %s", pathname.Data());
402 return;
403 }
404
405 fFile->SetCompressionLevel(fCompressLevel);
406
407 for (Int_t i=0;i<fTrees.GetEntries();++i) {
408 MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
409 mt->SetDirectory(fFile);
410 }
411
412 fIsInit = kTRUE;
413 }
414
415 //__________________________________________________________________________________________________
416 void TreeWriter::Print(Option_t *option) const
417 {
418 // Print the contents of the tree writer.
419
420 if (option) {
421 cout << ClassName() << " with members " << endl;
422 cout << " fBaseURL: " << fBaseURL << endl;
423 cout << " fPreFix: " << fPrefix << endl;
424 cout << " fFileNumber: " << fFileNumber << endl;
425 cout << " fCompressLevel: " << fCompressLevel << endl;
426 cout << " fDefBrSize: " << fDefBrSize << endl;
427 cout << " fDefSL: " << fDefSL << endl;
428 cout << " fMaxSize: " << fMaxSize << endl;
429 cout << " fDoObjNumReset: " << fDoObjNumReset << endl;
430 return;
431 }
432
433 cout << ClassName() << ": " << GetEntries()
434 << (GetEntries() == 1 ? " event" : " events") << endl;
435 }
436
437 //-------------------------------------------------------------------------------------------------
438 void TreeWriter::SetAutoFill(const char *tn, Bool_t b)
439 {
440 // Set auto-fill mode of tree with given name.
441
442 if (fTrees.GetEntries()==0)
443 return;
444
445 MyTree *mt = GetMyTree(tn);
446 if (!mt)
447 return;
448
449 mt->SetAutoFill(b);
450 }
451
452
453 //-------------------------------------------------------------------------------------------------
454 void TreeWriter::SetMaxSize(Long64_t s)
455 {
456 // Set maximum file size. Check if this exceeds the ROOT file size and if,
457 // print a warning and adjust it.
458
459 if (s>=(Long64_t)(0.99 * TTree::GetMaxTreeSize())) {
460 Long64_t news = (Long64_t)(s/0.99);
461 Warning("SetMaxSize", "Maximum tree size increased from %lld to %lld",
462 TTree::GetMaxTreeSize(), news);
463 TTree::SetMaxTreeSize(news);
464 }
465
466 fMaxSize=s;
467 }
468
469 //__________________________________________________________________________________________________
470 void TreeWriter::StoreObject(const TObject *obj)
471 {
472 // Store object next to tree in file. Used to store the
473 // settings of how the tree was created.
474
475 if (!fIsInit) {
476 Fatal("StoreObject", "Tree is not created, call create first!");
477 return;
478 }
479
480 if (!obj) {
481 Fatal("StoreObject", "Ptr to TObject is null!");
482 return;
483 }
484
485 fFile->WriteTObject(obj,obj->GetName(),"WriteDelete");
486 }