ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/MitAna/DataUtil/src/TreeWriter.cc
Revision: 1.2
Committed: Mon Jun 2 08:58:52 2008 UTC (16 years, 11 months ago) by loizides
Content type: text/plain
Branch: MAIN
Changes since 1.1: +188 -30 lines
Log Message:
Added AddBranchToTree functionality.

File Contents

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