ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/MitAna/DataUtil/src/TreeWriter.cc
Revision: 1.13
Committed: Mon Dec 1 17:39:20 2008 UTC (16 years, 5 months ago) by loizides
Content type: text/plain
Branch: MAIN
CVS Tags: Mit_006b, Mit_006a
Changes since 1.12: +42 -9 lines
Log Message:
Added AddTree and DoBranchRef

File Contents

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