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