1 |
// $Id: TreeWriter.cc,v 1.16 2009/03/23 22:15:11 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 |
Terminate();
|
45 |
}
|
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 |
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, void *obj, Int_t bsize, Int_t level)
|
61 |
{
|
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 |
AddBranch(name, CName(obj), obj, bsize, level);
|
66 |
}
|
67 |
|
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 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 |
MyTree *t = AddOrGetMyTree(GetName());
|
76 |
TBranch *b = t->Bronch(name, cname, obj, bsize, fDefSL);
|
77 |
b->SetCompressionLevel(GetCompressLevel());
|
78 |
}
|
79 |
|
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
|
206 |
// store the current object number to be reset in FillEvent().
|
207 |
|
208 |
if (!fIsInit) {
|
209 |
OpenFile();
|
210 |
}
|
211 |
|
212 |
if (doreset || fDoObjNumReset) {
|
213 |
fEvtObjNum = TProcessID::GetObjectCount();
|
214 |
}
|
215 |
|
216 |
return kTRUE;
|
217 |
}
|
218 |
|
219 |
//--------------------------------------------------------------------------------------------------
|
220 |
void TreeWriter::CloseFile()
|
221 |
{
|
222 |
// Write tree(s) and close file.
|
223 |
|
224 |
if (!fIsInit) {
|
225 |
Fatal("CloseFile", "File was not opened, call OpenFile() first!");
|
226 |
return;
|
227 |
}
|
228 |
|
229 |
TDirectory::TContext context(fFile); // cd fFile &&
|
230 |
// automatically restore gDirectory
|
231 |
|
232 |
for (Int_t i=0;i<fTrees.GetEntries();++i) {
|
233 |
MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
|
234 |
mt->FlushBaskets();
|
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;
|
252 |
fFile = 0;
|
253 |
|
254 |
fIsInit = kFALSE;
|
255 |
fFileNumber++;
|
256 |
}
|
257 |
|
258 |
//--------------------------------------------------------------------------------------------------
|
259 |
const char *TreeWriter::CName(void *obj) const
|
260 |
{
|
261 |
// Dereference void* pointer into TObject* pointer
|
262 |
|
263 |
TObject **sobj = static_cast<TObject**>(obj);
|
264 |
TObject *tobj = dynamic_cast<TObject*>(*sobj);
|
265 |
if (tobj==0) {
|
266 |
Fatal("CName", "Given void* ptr cannot be dereferenced into TObject*");
|
267 |
}
|
268 |
return tobj->ClassName();
|
269 |
}
|
270 |
|
271 |
//--------------------------------------------------------------------------------------------------
|
272 |
void TreeWriter::DoBranchRef(const char *tn)
|
273 |
{
|
274 |
// Fill BranchRef for given tree.
|
275 |
|
276 |
if (fTrees.GetEntries()==0) {
|
277 |
Error("DoBranchRef", "Tree with name %s not found!", tn);
|
278 |
return;
|
279 |
}
|
280 |
|
281 |
MyTree *mt = GetMyTree(tn);
|
282 |
if (!mt)
|
283 |
return;
|
284 |
|
285 |
mt->BranchRef();
|
286 |
}
|
287 |
|
288 |
//--------------------------------------------------------------------------------------------------
|
289 |
Bool_t TreeWriter::EndEvent(Bool_t doreset)
|
290 |
{
|
291 |
// Store the event in the tree. If doreset or fDoObjNumReset is kTRUE
|
292 |
// restore the stored object number at the time BeginEvent(kTRUE)
|
293 |
// was called.
|
294 |
|
295 |
if (!fIsInit) {
|
296 |
Fatal("EndEvent", "File is not open, did you call BeginEvent?");
|
297 |
return kFALSE;
|
298 |
}
|
299 |
|
300 |
Int_t r = 0;
|
301 |
for (Int_t i=0;i<fTrees.GetEntries();++i) {
|
302 |
MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
|
303 |
if (mt->GetAutoFill()==0)
|
304 |
continue;
|
305 |
r += mt->Fill();
|
306 |
}
|
307 |
|
308 |
if (IsFull())
|
309 |
CloseFile();
|
310 |
|
311 |
if (doreset || fDoObjNumReset) {
|
312 |
if (fEvtObjNum<0) {
|
313 |
Error("EndEvent", "Object counter is zero. Did you call BeginEvent(kTRUE)?");
|
314 |
} else {
|
315 |
// Reset the TRef table. Keep it from growing with each event (see doc)
|
316 |
TProcessID::SetObjectCount(fEvtObjNum);
|
317 |
}
|
318 |
}
|
319 |
|
320 |
return (r >= 0);
|
321 |
}
|
322 |
|
323 |
//--------------------------------------------------------------------------------------------------
|
324 |
Long64_t TreeWriter::GetEntries(const char *tn) const
|
325 |
{
|
326 |
// Return entries of tree with given name. If no tree is given, return sum of entries
|
327 |
// of all trees.
|
328 |
|
329 |
if (fTrees.GetEntries()==0) return -1;
|
330 |
|
331 |
if (tn) {
|
332 |
const TTree *mt=GetTree(tn);
|
333 |
if (mt)
|
334 |
return mt->GetEntries();
|
335 |
else
|
336 |
return -1;
|
337 |
}
|
338 |
|
339 |
Long64_t ret = 0;
|
340 |
for (Int_t i=0;i<fTrees.GetEntries();++i) {
|
341 |
const MyTree *mt = static_cast<const MyTree*>(fTrees.At(i));
|
342 |
ret += mt->GetEntries();
|
343 |
}
|
344 |
return ret;
|
345 |
}
|
346 |
|
347 |
//--------------------------------------------------------------------------------------------------
|
348 |
MyTree *mithep::TreeWriter::GetMyTree(const char *tn)
|
349 |
{
|
350 |
// Return MyTree with given name from array.
|
351 |
|
352 |
if (fTrees.GetEntries()==0)
|
353 |
return 0;
|
354 |
|
355 |
TObject *obj = 0;
|
356 |
if (tn==0) {
|
357 |
obj = fTrees.At(0);
|
358 |
} else {
|
359 |
obj = fTrees.FindObject(tn);
|
360 |
}
|
361 |
|
362 |
if (obj)
|
363 |
return static_cast<MyTree*>(obj);
|
364 |
return 0;
|
365 |
}
|
366 |
|
367 |
//--------------------------------------------------------------------------------------------------
|
368 |
const TTree *mithep::TreeWriter::GetTree(const char *tn) const
|
369 |
{
|
370 |
// Return TTree with given name from array.
|
371 |
|
372 |
if (fTrees.GetEntries()==0)
|
373 |
return 0;
|
374 |
|
375 |
TObject *obj = 0;
|
376 |
if (tn==0) {
|
377 |
obj = fTrees.At(0);
|
378 |
} else {
|
379 |
obj = fTrees.FindObject(tn);
|
380 |
}
|
381 |
|
382 |
if (obj)
|
383 |
return dynamic_cast<const TTree*>(obj);
|
384 |
return 0;
|
385 |
}
|
386 |
|
387 |
//--------------------------------------------------------------------------------------------------
|
388 |
TTree *mithep::TreeWriter::GetTree(const char *tn)
|
389 |
{
|
390 |
// Return TTree with given name from array.
|
391 |
|
392 |
if (fTrees.GetEntries()==0)
|
393 |
return 0;
|
394 |
|
395 |
TObject *obj = 0;
|
396 |
if (tn==0) {
|
397 |
obj = fTrees.At(0);
|
398 |
} else {
|
399 |
obj = fTrees.FindObject(tn);
|
400 |
}
|
401 |
|
402 |
if (obj)
|
403 |
return dynamic_cast<TTree*>(obj);
|
404 |
return 0;
|
405 |
}
|
406 |
|
407 |
//--------------------------------------------------------------------------------------------------
|
408 |
Bool_t TreeWriter::IsFull() const
|
409 |
{
|
410 |
// Check if the maximum file size has been reached.
|
411 |
|
412 |
Long64_t entries = GetEntries();
|
413 |
if (entries < 1) return kFALSE;
|
414 |
|
415 |
Long64_t avgSize = GetFileSize() / entries;
|
416 |
|
417 |
if (avgSize < fkMinAvgSize)
|
418 |
avgSize = fkMinAvgSize;
|
419 |
|
420 |
return (GetFileSize() + avgSize + fkMinFreeSpace) > fMaxSize;
|
421 |
}
|
422 |
|
423 |
//--------------------------------------------------------------------------------------------------
|
424 |
void TreeWriter::OpenFile()
|
425 |
{
|
426 |
// Open the file and attach the tree.
|
427 |
|
428 |
if (fIsInit) {
|
429 |
Fatal("OpenFile", "File is already open, call CloseFile first!");
|
430 |
return;
|
431 |
}
|
432 |
|
433 |
TDirectory::TContext context(0);
|
434 |
|
435 |
TString pathname=GetFullName();
|
436 |
gSystem->ExpandPathName(pathname);
|
437 |
|
438 |
fFile = TFile::Open(pathname, "RECREATE");
|
439 |
if (fFile == 0) {
|
440 |
Fatal("OpenFile", "Could not open file %s", pathname.Data());
|
441 |
return;
|
442 |
}
|
443 |
|
444 |
fFile->SetCompressionLevel(fCompressLevel);
|
445 |
|
446 |
for (Int_t i=0;i<fTrees.GetEntries();++i) {
|
447 |
MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
|
448 |
mt->SetDirectory(fFile);
|
449 |
}
|
450 |
|
451 |
fIsInit = kTRUE;
|
452 |
}
|
453 |
|
454 |
//--------------------------------------------------------------------------------------------------
|
455 |
void TreeWriter::Print(Option_t *option) const
|
456 |
{
|
457 |
// Print the contents of the tree writer.
|
458 |
|
459 |
if (option) {
|
460 |
cout << ClassName() << " with members " << endl;
|
461 |
cout << " fBaseURL: " << fBaseURL << endl;
|
462 |
cout << " fPreFix: " << fPrefix << endl;
|
463 |
cout << " fFileNumber: " << fFileNumber << endl;
|
464 |
cout << " fCompressLevel: " << fCompressLevel << endl;
|
465 |
cout << " fDefBrSize: " << fDefBrSize << endl;
|
466 |
cout << " fDefSL: " << fDefSL << endl;
|
467 |
cout << " fMaxSize: " << fMaxSize << endl;
|
468 |
cout << " fDoObjNumReset: " << fDoObjNumReset << endl;
|
469 |
return;
|
470 |
}
|
471 |
|
472 |
cout << ClassName() << ": " << GetEntries()
|
473 |
<< (GetEntries() == 1 ? " event" : " events") << endl;
|
474 |
}
|
475 |
|
476 |
//--------------------------------------------------------------------------------------------------
|
477 |
void TreeWriter::SetAutoFill(const char *tn, Bool_t b)
|
478 |
{
|
479 |
// Set auto-fill mode of tree with given name.
|
480 |
|
481 |
if (fTrees.GetEntries()==0)
|
482 |
return;
|
483 |
|
484 |
MyTree *mt = GetMyTree(tn);
|
485 |
if (!mt)
|
486 |
return;
|
487 |
|
488 |
mt->SetAutoFill(b);
|
489 |
}
|
490 |
|
491 |
//--------------------------------------------------------------------------------------------------
|
492 |
void TreeWriter::SetMaxSize(Long64_t s)
|
493 |
{
|
494 |
// Set maximum file size. Check if this exceeds the ROOT file size and if,
|
495 |
// print a warning and adjust it.
|
496 |
|
497 |
if (s>=(Long64_t)(0.99 * TTree::GetMaxTreeSize())) {
|
498 |
Long64_t news = (Long64_t)(s/0.99);
|
499 |
Warning("SetMaxSize", "Maximum tree size increased from %lld to %lld",
|
500 |
TTree::GetMaxTreeSize(), news);
|
501 |
TTree::SetMaxTreeSize(news);
|
502 |
}
|
503 |
|
504 |
fMaxSize=s;
|
505 |
}
|
506 |
|
507 |
//--------------------------------------------------------------------------------------------------
|
508 |
void TreeWriter::StoreObject(const TObject *obj)
|
509 |
{
|
510 |
// Store object next to tree in file. Used to store the
|
511 |
// settings of how the tree was created.
|
512 |
|
513 |
if (!fIsInit) {
|
514 |
Fatal("StoreObject", "Tree is not created, call create first!");
|
515 |
return;
|
516 |
}
|
517 |
|
518 |
if (!obj) {
|
519 |
Fatal("StoreObject", "Ptr to TObject is null!");
|
520 |
return;
|
521 |
}
|
522 |
|
523 |
fFile->WriteTObject(obj,obj->GetName(),"WriteDelete");
|
524 |
}
|
525 |
|
526 |
//__________________________________________________________________________________________________
|
527 |
void TreeWriter::Terminate()
|
528 |
{
|
529 |
// Terminate tree file writing.
|
530 |
|
531 |
if (fIsInit) {
|
532 |
CloseFile();
|
533 |
}
|
534 |
|
535 |
TDirectory::TContext context(0);
|
536 |
fTrees.Clear();
|
537 |
}
|