ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/MitAna/DataUtil/src/TreeWriter.cc
Revision: 1.12
Committed: Mon Oct 6 16:53:26 2008 UTC (16 years, 7 months ago) by loizides
Content type: text/plain
Branch: MAIN
CVS Tags: Mit_006, Mit_005
Changes since 1.11: +6 -4 lines
Log Message:
Add do BranchRef and GetTree

File Contents

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