ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/MitAna/DataUtil/src/TreeWriter.cc
(Generate patch)

Comparing UserCode/MitAna/DataUtil/src/TreeWriter.cc (file contents):
Revision 1.1 by loizides, Tue May 27 19:36:05 2008 UTC vs.
Revision 1.13 by loizides, Mon Dec 1 17:39:20 2008 UTC

# Line 1 | Line 1
1   // $Id$
2  
3   #include "MitAna/DataUtil/interface/TreeWriter.h"
4
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 <    fFile(0),
31 <    fTree(0)
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 <  TDirectory::TContext context(0);
35 <  fTree = new TTree(GetName(), GetTitle());
36 <  fTree->SetDirectory(0);
36 >  fTrees.SetOwner();
37   }
38  
39   //__________________________________________________________________________________________________
# Line 41 | Line 41 | TreeWriter::~TreeWriter()
41   {
42    // Destructor.
43    
44 <  if(fIsInit) {
45 <    CloseFile();
46 <  }
47 <
48 <  TDirectory::TContext context(0);
49 <  delete fTree;
44 >  Terminate();
45   }
46  
47   //__________________________________________________________________________________________________
48 < void TreeWriter::AddBranch(const char *name, const char *cname, void *obj, Int_t bsize, Int_t level)
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 and set its address to object pointer
52 <  // for class name "cname" using a given buffer size and splitlevel.
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 <  fTree->Bronch(name, cname, obj, bsize, level);
54 >  MyTree  *t = AddOrGetMyTree(GetName());
55 >  TBranch *b = t->Bronch(name, cname, obj, bsize, level);
56 >  b->SetCompressionLevel(GetCompressLevel());
57   }
58  
59 < //__________________________________________________________________________________________________
60 < void TreeWriter::AddBranch(const char *name, const char *cname, void *obj, Int_t bsize)
59 >
60 > //--------------------------------------------------------------------------------------------------
61 > void TreeWriter::AddBranch(const char *name, void *obj, Int_t bsize, Int_t level)
62   {
63 <  // Add branch with name "name" into tree and set its address to object pointer
64 <  // for class name "cname" using a given buffer size and default splitlevel.
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 <  fTree->Bronch(name, cname, obj, bsize, fDefSL);
66 >  AddBranch(name, CName(obj), obj, bsize, level);
67   }
68  
69 < //__________________________________________________________________________________________________
70 < void TreeWriter::AddBranch(const char *name, const char *cname, void *obj)
69 > //--------------------------------------------------------------------------------------------------
70 > void TreeWriter::AddBranch(const char *name, const char *cname,
71 >                           void *obj, Int_t bsize)
72   {
73 <  // Add branch with name "name" into tree and set its address to object pointer
74 <  // for class name "cname" using a given buffer size and default splitlevel.
73 >  // 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  
76 <  fTree->Bronch(name, cname, obj, fDefBrSize, fDefSL);
76 >  MyTree  *t = AddOrGetMyTree(GetName());
77 >  TBranch *b = t->Bronch(name, cname, obj, bsize, fDefSL);
78 >  b->SetCompressionLevel(GetCompressLevel());
79   }
80  
81 < //__________________________________________________________________________________________________
81 > //--------------------------------------------------------------------------------------------------
82 > 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 > //--------------------------------------------------------------------------------------------------
91 > void TreeWriter::AddBranch(const char *name, const char *cname,
92 >                           void *obj)
93 > {
94 >  // Add branch with name "name" into tree with name "tname" and set its address
95 >  // to object pointer for class name "cname" using a default buffer size and splitlevel.
96 >
97 >  MyTree  *t = AddOrGetMyTree(GetName());
98 >  TBranch *b = t->Bronch(name, cname, obj, fDefBrSize, fDefSL);
99 >  b->SetCompressionLevel(GetCompressLevel());
100 > }
101 >
102 > //--------------------------------------------------------------------------------------------------
103 > 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 > //--------------------------------------------------------------------------------------------------
112 > void TreeWriter::AddBranchToTree(const char *tname, const char *name, const char *cname,
113 >                                 void *obj, Int_t bsize, Int_t level)
114 > {
115 >  // 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 >
118 >  MyTree  *t = AddOrGetMyTree(tname);
119 >  TBranch *b = t->Bronch(name, cname, obj, bsize, level);
120 >  b->SetCompressionLevel(GetCompressLevel());
121 > }
122 >
123 > //--------------------------------------------------------------------------------------------------
124 > 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 > //--------------------------------------------------------------------------------------------------
134 > void TreeWriter::AddBranchToTree(const char *tname, const char *name, const char *cname,
135 >                                 void *obj, Int_t bsize)
136 > {
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 >  MyTree  *t = AddOrGetMyTree(tname);
141 >  TBranch *b = t->Bronch(name, cname, obj, bsize, fDefSL);
142 >  b->SetCompressionLevel(GetCompressLevel());
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 >  TBranch *b = t->Bronch(name, cname, obj, fDefBrSize, fDefSL);
164 >  b->SetCompressionLevel(GetCompressLevel());
165 > }
166 >
167 > //--------------------------------------------------------------------------------------------------
168 > 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 > //--------------------------------------------------------------------------------------------------
177 > 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 > 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 >  MyTree *tree = GetMyTree(tn);
191 >  if (tree)
192 >    return tree;
193 >
194 >  TDirectory::TContext context(fFile);
195 >  tree = new MyTree(tn, tn);
196 >  tree->SetDirectory(fFile);
197 >  if (fDoBranchRef)
198 >    tree->BranchRef();
199 >  fTrees.AddLast(tree);  
200 >  return tree;
201 > }
202 >
203 > //--------------------------------------------------------------------------------------------------
204   Bool_t TreeWriter::BeginEvent(Bool_t doreset)
205   {
206    // Prepare for the next event. If doreset or fDoObjNumReset is kTRUE
# Line 86 | Line 210 | Bool_t TreeWriter::BeginEvent(Bool_t dor
210      OpenFile();
211    }
212  
213 <  if(doreset || fDoObjNumReset) {
213 >  if (doreset || fDoObjNumReset) {
214      fEvtObjNum = TProcessID::GetObjectCount();
215    }
216  
217    return kTRUE;
218   }
219  
220 < //__________________________________________________________________________________________________
220 > //--------------------------------------------------------------------------------------------------
221   void TreeWriter::CloseFile()
222   {
223 <  // Write tree and close file.
223 >  // Write tree(s) and close file.
224  
225    if (!fIsInit) {
226      Fatal("CloseFile", "File was not opened, call OpenFile() first!");
# Line 106 | Line 230 | void TreeWriter::CloseFile()
230    TDirectory::TContext context(fFile); // cd fFile &&
231                                         // automatically restore gDirectory
232  
233 <  fTree->Write(fTree->GetName(),TObject::kOverwrite);
234 <  fTree->Reset();
235 <  fTree->SetDirectory(0);  
233 >  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 >
237 >    // backup and restore list of branch pointers from TRefTable (needed for autoloading)
238 >    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 >    mt->SetDirectory(0);  
248 >  }
249  
250    fFile->Close();
251    delete fFile;
# Line 118 | Line 255 | void TreeWriter::CloseFile()
255    fFileNumber++;  
256   }
257  
258 < //__________________________________________________________________________________________________
258 > //--------------------------------------------------------------------------------------------------
259 > 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 >  if (tobj==0) {
265 >    Fatal("CName", "Given void* ptr can not be dereferenced into TObject*");
266 >  }
267 >  return tobj->ClassName();
268 > }
269 >
270 > //--------------------------------------------------------------------------------------------------
271 > 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   Bool_t TreeWriter::EndEvent(Bool_t doreset)
289   {
290    // Store the event in the tree. If doreset or fDoObjNumReset is kTRUE
# Line 130 | Line 296 | Bool_t TreeWriter::EndEvent(Bool_t dores
296      return kFALSE;
297    }
298  
299 <  //TDirectory::TContext context(fFile);
300 <  Int_t r = fTree->Fill();
299 >  Int_t r = 0;
300 >  for (Int_t i=0;i<fTrees.GetEntries();++i) {
301 >    MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
302 >    if (mt->GetAutoFill()==0)
303 >      continue;
304 >    r += mt->Fill();
305 >  }
306  
307 <  if(IsFull())
307 >  if (IsFull())
308      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 <      // Reset the TRef table. keep it from growing with each event (see doc)
314 >      // Reset the TRef table. Keep it from growing with each event (see doc)
315        TProcessID::SetObjectCount(fEvtObjNum);
316      }
317    }
# Line 148 | Line 319 | Bool_t TreeWriter::EndEvent(Bool_t dores
319    return (r >= 0);
320   }
321  
322 < //__________________________________________________________________________________________________
322 > //--------------------------------------------------------------------------------------------------
323 > Long64_t TreeWriter::GetEntries(const char *tn) const
324 > {
325 >  // Return entries of tree with given name. If no tree is given, return sum of entries
326 >  // of all trees.
327 >
328 >   if (fTrees.GetEntries()==0) return -1;
329 >
330 >   if (tn) {
331 >     const TTree *mt=GetTree(tn);
332 >     if (mt)
333 >       return mt->GetEntries();
334 >     else
335 >       return -1;
336 >   }
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 > //--------------------------------------------------------------------------------------------------
347 > 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 > //--------------------------------------------------------------------------------------------------
367 > 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 > //--------------------------------------------------------------------------------------------------
387 > 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 > //--------------------------------------------------------------------------------------------------
407   Bool_t TreeWriter::IsFull() const
408   {
409    // Check if the maximum file size has been reached.
410  
411    Long64_t entries = GetEntries();
157  
412    if (entries < 1) return kFALSE;
413    
414    Long64_t avgSize = GetFileSize() / entries;
# Line 165 | Line 419 | Bool_t TreeWriter::IsFull() const
419    return (GetFileSize() + avgSize + fkMinFreeSpace) > fMaxSize;
420   }
421  
422 < //__________________________________________________________________________________________________
422 > //--------------------------------------------------------------------------------------------------
423   void TreeWriter::OpenFile()
424   {
425    // Open the file and attach the tree.
# Line 187 | Line 441 | void TreeWriter::OpenFile()
441    }
442  
443    fFile->SetCompressionLevel(fCompressLevel);
444 <  fTree->SetDirectory(fFile);
444 >
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    fIsInit = kTRUE;
451   }
452  
453 < //__________________________________________________________________________________________________
453 > //--------------------------------------------------------------------------------------------------
454   void TreeWriter::Print(Option_t *option) const
455   {
456    // Print the contents of the tree writer.
457  
458 <  if(option) {
458 >  if (option) {
459      cout << ClassName() << " with members " << endl;
460      cout << "   fBaseURL:       " << fBaseURL << endl;
461      cout << "   fPreFix:        " << fPrefix << endl;
# Line 213 | Line 472 | void TreeWriter::Print(Option_t *option)
472         << (GetEntries() == 1 ? " event" : " events") << endl;
473   }
474  
475 < //__________________________________________________________________________________________________
475 > //--------------------------------------------------------------------------------------------------
476 > void TreeWriter::SetAutoFill(const char *tn, Bool_t b)
477 > {
478 >  // Set auto-fill mode of tree with given name.
479 >
480 >  if (fTrees.GetEntries()==0)
481 >    return;
482 >
483 >  MyTree *mt = GetMyTree(tn);
484 >  if (!mt)
485 >    return;
486 >
487 >  mt->SetAutoFill(b);
488 > }
489 >
490 > //--------------------------------------------------------------------------------------------------
491 > 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 > //--------------------------------------------------------------------------------------------------
507   void TreeWriter::StoreObject(const TObject *obj)
508   {
509    // Store object next to tree in file. Used to store the
# Line 231 | Line 521 | void TreeWriter::StoreObject(const TObje
521  
522    fFile->WriteTObject(obj,obj->GetName(),"WriteDelete");
523   }
524 +
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 + }

Diff Legend

Removed lines
+ Added lines
< Changed lines
> Changed lines