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.15 by loizides, Thu Mar 12 18:26:08 2009 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 > void TreeWriter::AddBranch(const char *name, void *obj, Int_t bsize, Int_t level)
61   {
62 <  // Add branch with name "name" into tree and set its address to object pointer
63 <  // for class name "cname" using a given buffer size and default splitlevel.
62 >  // Add branch with name "name" into tree with name "tname" and set its address
63 >  // to object pointer using a given buffer size and splitlevel.
64  
65 <  fTree->Bronch(name, cname, obj, bsize, fDefSL);
65 >  AddBranch(name, CName(obj), obj, bsize, level);
66   }
67  
68 < //__________________________________________________________________________________________________
69 < void TreeWriter::AddBranch(const char *name, const char *cname, void *obj)
68 > //--------------------------------------------------------------------------------------------------
69 > void TreeWriter::AddBranch(const char *name, const char *cname,
70 >                           void *obj, Int_t bsize)
71   {
72 <  // Add branch with name "name" into tree and set its address to object pointer
73 <  // for class name "cname" using a given buffer size and default splitlevel.
72 >  // Add branch with name "name" into tree with name "tname" and set its address
73 >  // to object pointer for class name "cname" using a given buffer size and default splitlevel.
74  
75 <  fTree->Bronch(name, cname, obj, fDefBrSize, fDefSL);
75 >  MyTree  *t = AddOrGetMyTree(GetName());
76 >  TBranch *b = t->Bronch(name, cname, obj, bsize, fDefSL);
77 >  b->SetCompressionLevel(GetCompressLevel());
78   }
79  
80 < //__________________________________________________________________________________________________
80 > //--------------------------------------------------------------------------------------------------
81 > void TreeWriter::AddBranch(const char *name, void *obj, Int_t bsize)
82 > {
83 >  // Add branch with name "name" into tree with name "tname" and set its address
84 >  // to object pointer using a given buffer size and default splitlevel.
85 >
86 >  AddBranch(name, CName(obj), obj, bsize);
87 > }
88 >
89 > //--------------------------------------------------------------------------------------------------
90 > void TreeWriter::AddBranch(const char *name, const char *cname,
91 >                           void *obj)
92 > {
93 >  // Add branch with name "name" into tree with name "tname" and set its address
94 >  // to object pointer for class name "cname" using a default buffer size and splitlevel.
95 >
96 >  MyTree  *t = AddOrGetMyTree(GetName());
97 >  TBranch *b = t->Bronch(name, cname, obj, fDefBrSize, fDefSL);
98 >  b->SetCompressionLevel(GetCompressLevel());
99 > }
100 >
101 > //--------------------------------------------------------------------------------------------------
102 > void TreeWriter::AddBranch(const char *name, void *obj)
103 > {
104 >  // Add branch with name "name" into tree with name "tname" and set its address
105 >  // to object pointer using a default buffer size and splitlevel.
106 >
107 >  AddBranch(name, CName(obj), obj);
108 > }
109 >
110 > //--------------------------------------------------------------------------------------------------
111 > void TreeWriter::AddBranchToTree(const char *tname, const char *name, const char *cname,
112 >                                 void *obj, Int_t bsize, Int_t level)
113 > {
114 >  // Add branch with name "name" into tree with name "tname" and set its address
115 >  // to object pointer for class name "cname" using a given buffer size and splitlevel.
116 >
117 >  MyTree  *t = AddOrGetMyTree(tname);
118 >  TBranch *b = t->Bronch(name, cname, obj, bsize, level);
119 >  b->SetCompressionLevel(GetCompressLevel());
120 > }
121 >
122 > //--------------------------------------------------------------------------------------------------
123 > void TreeWriter::AddBranchToTree(const char *tname, const char *name, void *obj,
124 >                                 Int_t bsize, Int_t level)
125 > {
126 >  // Add branch with name "name" into tree with name "tname" and set its address
127 >  // to object pointer using a given buffer size and splitlevel.
128 >
129 >  AddBranchToTree(tname, name, CName(obj), obj, bsize, level);
130 > }
131 >
132 > //--------------------------------------------------------------------------------------------------
133 > void TreeWriter::AddBranchToTree(const char *tname, const char *name, const char *cname,
134 >                                 void *obj, Int_t bsize)
135 > {
136 >  // Add branch with name "name" into tree with name "tname" and set its address
137 >  // to object pointer for class name "cname" using a given buffer size and default splitlevel.
138 >
139 >  MyTree  *t = AddOrGetMyTree(tname);
140 >  TBranch *b = t->Bronch(name, cname, obj, bsize, fDefSL);
141 >  b->SetCompressionLevel(GetCompressLevel());
142 > }
143 >
144 > //--------------------------------------------------------------------------------------------------
145 > void TreeWriter::AddBranchToTree(const char *tname, const char *name, void *obj,
146 >                                 Int_t bsize)
147 > {
148 >  // Add branch with name "name" into tree with name "tname" and set its address
149 >  // to object pointer using a given buffer size and default splitlevel.
150 >
151 >  AddBranchToTree(tname, name, CName(obj), obj, bsize);
152 > }
153 >
154 > //--------------------------------------------------------------------------------------------------
155 > void TreeWriter::AddBranchToTree(const char *tname, const char *name, const char *cname,
156 >                                 void *obj)
157 > {
158 >  // Add branch with name "name" into tree with name "tname" and set its address
159 >  // to object pointer for class name "cname" using a default buffer size and splitlevel.
160 >
161 >  MyTree  *t = AddOrGetMyTree(tname);
162 >  TBranch *b = t->Bronch(name, cname, obj, fDefBrSize, fDefSL);
163 >  b->SetCompressionLevel(GetCompressLevel());
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 > void TreeWriter::AddTree(const char *tname)
177 > {
178 >  // Add tree with name "name" into tree with name "tname".
179 >
180 >  AddOrGetMyTree(tname);
181 > }
182 >
183 > //--------------------------------------------------------------------------------------------------
184 > MyTree *TreeWriter::AddOrGetMyTree(const char *tn)
185 > {
186 >  // Add new tree if not present in array of trees or return
187 >  // present tree.
188 >
189 >  MyTree *tree = GetMyTree(tn);
190 >  if (tree)
191 >    return tree;
192 >
193 >  TDirectory::TContext context(fFile);
194 >  tree = new MyTree(tn, tn);
195 >  tree->SetDirectory(fFile);
196 >  if (fDoBranchRef)
197 >    tree->BranchRef();
198 >  fTrees.AddLast(tree);  
199 >  return tree;
200 > }
201 >
202 > //--------------------------------------------------------------------------------------------------
203   Bool_t TreeWriter::BeginEvent(Bool_t doreset)
204   {
205    // Prepare for the next event. If doreset or fDoObjNumReset is kTRUE
# Line 86 | Line 209 | Bool_t TreeWriter::BeginEvent(Bool_t dor
209      OpenFile();
210    }
211  
212 <  if(doreset || fDoObjNumReset) {
212 >  if (doreset || fDoObjNumReset) {
213      fEvtObjNum = TProcessID::GetObjectCount();
214    }
215  
216    return kTRUE;
217   }
218  
219 < //__________________________________________________________________________________________________
219 > //--------------------------------------------------------------------------------------------------
220   void TreeWriter::CloseFile()
221   {
222 <  // Write tree and close file.
222 >  // Write tree(s) and close file.
223  
224    if (!fIsInit) {
225      Fatal("CloseFile", "File was not opened, call OpenFile() first!");
# Line 106 | Line 229 | void TreeWriter::CloseFile()
229    TDirectory::TContext context(fFile); // cd fFile &&
230                                         // automatically restore gDirectory
231  
232 <  fTree->Write(fTree->GetName(),TObject::kOverwrite);
233 <  fTree->Reset();
234 <  fTree->SetDirectory(0);  
232 >  for (Int_t i=0;i<fTrees.GetEntries();++i) {
233 >    MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
234 >    mt->Write(mt->GetName(),TObject::kOverwrite);
235 >
236 >    // backup and restore list of branch pointers from TRefTable (needed for autoloading)
237 >    if (mt->GetBranchRef()) {
238 >      TObjArray *parents = mt->GetBranchRef()->GetRefTable()->GetParents();
239 >      TObjArray parentsBak(*parents);
240 >      mt->Reset();
241 >      for (Int_t j=0; j<parentsBak.GetEntries(); ++j)
242 >        parents->Add(parentsBak.At(j));
243 >    }
244 >    else
245 >      mt->Reset();
246 >    mt->SetDirectory(0);  
247 >  }
248  
249    fFile->Close();
250    delete fFile;
# Line 118 | Line 254 | void TreeWriter::CloseFile()
254    fFileNumber++;  
255   }
256  
257 < //__________________________________________________________________________________________________
257 > //--------------------------------------------------------------------------------------------------
258 > const char *TreeWriter::CName(void *obj) const
259 > {
260 >  // Dereference void* pointer into TObject* pointer
261 >
262 >  TObject **sobj = static_cast<TObject**>(obj);
263 >  TObject *tobj = dynamic_cast<TObject*>(*sobj);
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