DataUtil/src/TreeWriter.cc

// $Id: TreeWriter.cxx 4071 2007-06-06 08:39:13Z loizides $

#include "MitAna/DataUtil/interface/TreeWriter.h"

#include <Riostream.h>
#include <TObject.h>
#include <TSystem.h>
#include <TProcessID.h>

using namespace mithep;

ClassImp(mithep::TreeWriter)

//__________________________________________________________________________________________________
TreeWriter::TreeWriter(const char *tname, Bool_t doreset)
  : TNamed(tname,Form("%s written by mithep::TreeWriter", tname)),
    fBaseURL("."), 
    fPrefix("mithep"), 
    fFileNumber(0),
    fCompressLevel(9),
    fDefBrSize(64*1024),
    fDefSL(99),
    fMaxSize((Long64_t)(0.99 * TTree::GetMaxTreeSize())),
    fkMinFreeSpace(1024*1024), 
    fkMinAvgSize(10*1024), 
    fEvtObjNum(-1), 
    fIsInit(kFALSE), 
    fDoObjNumReset(doreset), 
    fFile(0), 
    fTree(0)
{
  // Constructor.

  TDirectory::TContext context(0); 
  fTree = new TTree(GetName(), GetTitle());
  fTree->SetDirectory(0);
}

//__________________________________________________________________________________________________
TreeWriter::~TreeWriter()
{
  // Destructor.
   
  if(fIsInit) {
    CloseFile();
  }

  TDirectory::TContext context(0); 
  delete fTree;
}

//__________________________________________________________________________________________________
void TreeWriter::AddBranch(const char *name, const char *cname, void *obj, Int_t bsize, Int_t level)
{
  // Add branch with name "name" into tree and set its address to object pointer 
  // for class name "cname" using a given buffer size and splitlevel.

  fTree->Bronch(name, cname, obj, bsize, level);
}

//__________________________________________________________________________________________________
void TreeWriter::AddBranch(const char *name, const char *cname, void *obj, Int_t bsize)
{
  // Add branch with name "name" into tree and set its address to object pointer 
  // for class name "cname" using a given buffer size and default splitlevel.

  fTree->Bronch(name, cname, obj, bsize, fDefSL);
}

//__________________________________________________________________________________________________
void TreeWriter::AddBranch(const char *name, const char *cname, void *obj)
{
  // Add branch with name "name" into tree and set its address to object pointer 
  // for class name "cname" using a given buffer size and default splitlevel.

  fTree->Bronch(name, cname, obj, fDefBrSize, fDefSL);
}

//__________________________________________________________________________________________________
Bool_t TreeWriter::BeginEvent(Bool_t doreset)
{
  // Prepare for the next event. If doreset or fDoObjNumReset is kTRUE
  // store the current object number to be reset in FillEvent().

  if (!fIsInit) {
    OpenFile();
  }

  if(doreset || fDoObjNumReset) {
    fEvtObjNum = TProcessID::GetObjectCount();
  }

  return kTRUE;
}

//__________________________________________________________________________________________________
void TreeWriter::CloseFile()
{
  // Write tree and close file.

  if (!fIsInit) {
    Fatal("CloseFile", "File was not opened, call OpenFile() first!");
    return;
  }

  TDirectory::TContext context(fFile); // cd fFile && 
                                       // automatically restore gDirectory

  fTree->Write(fTree->GetName(),TObject::kOverwrite);
  fTree->Reset();
  fTree->SetDirectory(0);   

  fFile->Close();
  delete fFile; 
  fFile = 0;

  fIsInit = kFALSE; 
  fFileNumber++;  
}

//__________________________________________________________________________________________________
Bool_t TreeWriter::EndEvent(Bool_t doreset)
{
  // Store the event in the tree. If doreset or fDoObjNumReset is kTRUE
  // restore the stored object number at the time BeginEvent(kTRUE) 
  // was called.
   
  if (!fIsInit) {
    Fatal("EndEvent", "File is not open, did you call BeginEvent?");
    return kFALSE;
  }

  //TDirectory::TContext context(fFile);
  Int_t r = fTree->Fill();

  if(IsFull())
    CloseFile();

  if (doreset || fDoObjNumReset) {
    if (fEvtObjNum<0) {
      Error("EndEvent", "Object counter is zero. Did you call BeginEvent(kTRUE)?");
    } else {
      // Reset the TRef table. keep it from growing with each event (see doc)
      TProcessID::SetObjectCount(fEvtObjNum);
    }
  }

  return (r >= 0);
}

//__________________________________________________________________________________________________
Bool_t TreeWriter::IsFull() const
{
  // Check if the maximum file size has been reached.

  Long64_t entries = GetEntries();
   
  if (entries < 1) return kFALSE;
   
  Long64_t avgSize = GetFileSize() / entries;
   
  if (avgSize < fkMinAvgSize)
    avgSize = fkMinAvgSize;

  return (GetFileSize() + avgSize + fkMinFreeSpace) > fMaxSize;
}

//__________________________________________________________________________________________________
void TreeWriter::OpenFile()
{
  // Open the file and attach the tree.

  if (fIsInit) {
    Fatal("OpenFile", "File is already open, call CloseFile first!");
    return;
  }

  TDirectory::TContext context(0);
   
  TString pathname=GetFullName();
  gSystem->ExpandPathName(pathname);

  fFile = TFile::Open(pathname, "RECREATE");
  if (fFile == 0) {
    Fatal("OpenFile", "Could not open file %s", pathname.Data());
    return;
  }

  fFile->SetCompressionLevel(fCompressLevel);
  fTree->SetDirectory(fFile);
  fIsInit = kTRUE;
}

//__________________________________________________________________________________________________
void TreeWriter::Print(Option_t *option) const
{
  // Print the contents of the tree writer.

  if(option) {
    cout << ClassName() << " with members " << endl;
    cout << "   fBaseURL:       " << fBaseURL << endl; 
    cout << "   fPreFix:        " << fPrefix << endl; 
    cout << "   fFileNumber:    " << fFileNumber << endl; 
    cout << "   fCompressLevel: " << fCompressLevel << endl; 
    cout << "   fDefBrSize:     " << fDefBrSize << endl; 
    cout << "   fDefSL:         " << fDefSL << endl; 
    cout << "   fMaxSize:       " << fMaxSize << endl; 
    cout << "   fDoObjNumReset: " << fDoObjNumReset << endl;
    return;
  }

  cout << ClassName() << ": " << GetEntries() 
       << (GetEntries() == 1 ? " event" : " events") << endl;
}

//__________________________________________________________________________________________________
void TreeWriter::StoreObject(const TObject *obj)
{
  // Store object next to tree in file. Used to store the
  // settings of how the tree was created.

  if (!fIsInit) {
    Fatal("StoreObject", "Tree is not created, call create first!");
    return;
  }

  if (!obj) {
    Fatal("StoreObject", "Ptr to TObject is null!");
    return;
  }

  fFile->WriteTObject(obj,obj->GetName(),"WriteDelete");
}
Revision:	1.1
Committed:	Tue May 27 19:36:05 2008 UTC (16 years, 11 months ago) by loizides
Content type:	text/plain
Branch:	MAIN
Log Message:	First commit of Util classes.
#	User	Rev	Content
1	loizides	1.1	// $Id: TreeWriter.cxx 4071 2007-06-06 08:39:13Z loizides $
2
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			fTree(0)
31			{
32			// Constructor.
33
34			TDirectory::TContext context(0);
35			fTree = new TTree(GetName(), GetTitle());
36			fTree->SetDirectory(0);
37			}
38
39			//__________________________________________________________________________________________________
40			TreeWriter::~TreeWriter()
41			{
42			// Destructor.
43
44			if(fIsInit) {
45			CloseFile();
46			}
47
48			TDirectory::TContext context(0);
49			delete fTree;
50			}
51
52			//__________________________________________________________________________________________________
53			void TreeWriter::AddBranch(const char name, const char cname, void *obj, Int_t bsize, Int_t level)
54			{
55			// Add branch with name "name" into tree and set its address to object pointer
56			// for class name "cname" using a given buffer size and splitlevel.
57
58			fTree->Bronch(name, cname, obj, bsize, level);
59			}
60
61			//__________________________________________________________________________________________________
62			void TreeWriter::AddBranch(const char name, const char cname, void *obj, Int_t bsize)
63			{
64			// Add branch with name "name" into tree and set its address to object pointer
65			// for class name "cname" using a given buffer size and default splitlevel.
66
67			fTree->Bronch(name, cname, obj, bsize, fDefSL);
68			}
69
70			//__________________________________________________________________________________________________
71			void TreeWriter::AddBranch(const char name, const char cname, void *obj)
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.
75
76			fTree->Bronch(name, cname, obj, fDefBrSize, fDefSL);
77			}
78
79			//__________________________________________________________________________________________________
80			Bool_t TreeWriter::BeginEvent(Bool_t doreset)
81			{
82			// Prepare for the next event. If doreset or fDoObjNumReset is kTRUE
83			// store the current object number to be reset in FillEvent().
84
85			if (!fIsInit) {
86			OpenFile();
87			}
88
89			if(doreset \|\| fDoObjNumReset) {
90			fEvtObjNum = TProcessID::GetObjectCount();
91			}
92
93			return kTRUE;
94			}
95
96			//__________________________________________________________________________________________________
97			void TreeWriter::CloseFile()
98			{
99			// Write tree and close file.
100
101			if (!fIsInit) {
102			Fatal("CloseFile", "File was not opened, call OpenFile() first!");
103			return;
104			}
105
106			TDirectory::TContext context(fFile); // cd fFile &&
107			// automatically restore gDirectory
108
109			fTree->Write(fTree->GetName(),TObject::kOverwrite);
110			fTree->Reset();
111			fTree->SetDirectory(0);
112
113			fFile->Close();
114			delete fFile;
115			fFile = 0;
116
117			fIsInit = kFALSE;
118			fFileNumber++;
119			}
120
121			//__________________________________________________________________________________________________
122			Bool_t TreeWriter::EndEvent(Bool_t doreset)
123			{
124			// Store the event in the tree. If doreset or fDoObjNumReset is kTRUE
125			// restore the stored object number at the time BeginEvent(kTRUE)
126			// was called.
127
128			if (!fIsInit) {
129			Fatal("EndEvent", "File is not open, did you call BeginEvent?");
130			return kFALSE;
131			}
132
133			//TDirectory::TContext context(fFile);
134			Int_t r = fTree->Fill();
135
136			if(IsFull())
137			CloseFile();
138
139			if (doreset \|\| fDoObjNumReset) {
140			if (fEvtObjNum<0) {
141			Error("EndEvent", "Object counter is zero. Did you call BeginEvent(kTRUE)?");
142			} else {
143			// Reset the TRef table. keep it from growing with each event (see doc)
144			TProcessID::SetObjectCount(fEvtObjNum);
145			}
146			}
147
148			return (r >= 0);
149			}
150
151			//__________________________________________________________________________________________________
152			Bool_t TreeWriter::IsFull() const
153			{
154			// Check if the maximum file size has been reached.
155
156			Long64_t entries = GetEntries();
157
158			if (entries < 1) return kFALSE;
159
160			Long64_t avgSize = GetFileSize() / entries;
161
162			if (avgSize < fkMinAvgSize)
163			avgSize = fkMinAvgSize;
164
165			return (GetFileSize() + avgSize + fkMinFreeSpace) > fMaxSize;
166			}
167
168			//__________________________________________________________________________________________________
169			void TreeWriter::OpenFile()
170			{
171			// Open the file and attach the tree.
172
173			if (fIsInit) {
174			Fatal("OpenFile", "File is already open, call CloseFile first!");
175			return;
176			}
177
178			TDirectory::TContext context(0);
179
180			TString pathname=GetFullName();
181			gSystem->ExpandPathName(pathname);
182
183			fFile = TFile::Open(pathname, "RECREATE");
184			if (fFile == 0) {
185			Fatal("OpenFile", "Could not open file %s", pathname.Data());
186			return;
187			}
188
189			fFile->SetCompressionLevel(fCompressLevel);
190			fTree->SetDirectory(fFile);
191			fIsInit = kTRUE;
192			}
193
194			//__________________________________________________________________________________________________
195			void TreeWriter::Print(Option_t *option) const
196			{
197			// Print the contents of the tree writer.
198
199			if(option) {
200			cout << ClassName() << " with members " << endl;
201			cout << " fBaseURL: " << fBaseURL << endl;
202			cout << " fPreFix: " << fPrefix << endl;
203			cout << " fFileNumber: " << fFileNumber << endl;
204			cout << " fCompressLevel: " << fCompressLevel << endl;
205			cout << " fDefBrSize: " << fDefBrSize << endl;
206			cout << " fDefSL: " << fDefSL << endl;
207			cout << " fMaxSize: " << fMaxSize << endl;
208			cout << " fDoObjNumReset: " << fDoObjNumReset << endl;
209			return;
210			}
211
212			cout << ClassName() << ": " << GetEntries()
213			<< (GetEntries() == 1 ? " event" : " events") << endl;
214			}
215
216			//__________________________________________________________________________________________________
217			void TreeWriter::StoreObject(const TObject *obj)
218			{
219			// Store object next to tree in file. Used to store the
220			// settings of how the tree was created.
221
222			if (!fIsInit) {
223			Fatal("StoreObject", "Tree is not created, call create first!");
224			return;
225			}
226
227			if (!obj) {
228			Fatal("StoreObject", "Ptr to TObject is null!");
229			return;
230			}
231
232			fFile->WriteTObject(obj,obj->GetName(),"WriteDelete");
233			}