DataUtil/src/TreeWriter.cc

// $Id: TreeWriter.cc,v 1.14 2009/03/03 08:37:12 loizides Exp $

#include "MitAna/DataUtil/interface/TreeWriter.h"
#include <Riostream.h>
#include <TObject.h>
#include <TSystem.h>
#include <TProcessID.h>
#include <TBranchRef.h>
#include "MitAna/DataUtil/interface/Debug.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), 
  fDoBranchRef(0), 
  fFile(0), 
  fTrees(0)
{
  // Constructor.

  fTrees.SetOwner();
}

//__________________________________________________________________________________________________
TreeWriter::~TreeWriter()
{
  // Destructor.
   
  Terminate();
}

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

  MyTree  *t = AddOrGetMyTree(GetName());
  TBranch *b = t->Bronch(name, cname, obj, bsize, level);
  b->SetCompressionLevel(GetCompressLevel());
}

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

  AddBranch(name, CName(obj), obj, bsize, level);
}

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

  MyTree  *t = AddOrGetMyTree(GetName());
  TBranch *b = t->Bronch(name, cname, obj, bsize, fDefSL);
  b->SetCompressionLevel(GetCompressLevel());
}

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

  AddBranch(name, CName(obj), obj, bsize);
}

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

  MyTree  *t = AddOrGetMyTree(GetName());
  TBranch *b = t->Bronch(name, cname, obj, fDefBrSize, fDefSL);
  b->SetCompressionLevel(GetCompressLevel());
}

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

  AddBranch(name, CName(obj), obj);
}

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

  MyTree  *t = AddOrGetMyTree(tname);
  TBranch *b = t->Bronch(name, cname, obj, bsize, level);
  b->SetCompressionLevel(GetCompressLevel());
}

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

  AddBranchToTree(tname, name, CName(obj), obj, bsize, level);
}

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

  MyTree  *t = AddOrGetMyTree(tname);
  TBranch *b = t->Bronch(name, cname, obj, bsize, fDefSL);
  b->SetCompressionLevel(GetCompressLevel());
}

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

  AddBranchToTree(tname, name, CName(obj), obj, bsize);
}

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

  MyTree  *t = AddOrGetMyTree(tname);
  TBranch *b = t->Bronch(name, cname, obj, fDefBrSize, fDefSL);
  b->SetCompressionLevel(GetCompressLevel());
}

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

  AddBranchToTree(tname, name, CName(obj), obj);
}

//--------------------------------------------------------------------------------------------------
void TreeWriter::AddTree(const char *tname)
{
  // Add tree with name "name" into tree with name "tname".

  AddOrGetMyTree(tname);
}

//--------------------------------------------------------------------------------------------------
MyTree *TreeWriter::AddOrGetMyTree(const char *tn)
{
  // Add new tree if not present in array of trees or return
  // present tree.

  MyTree *tree = GetMyTree(tn);
  if (tree)
    return tree;

  TDirectory::TContext context(fFile); 
  tree = new MyTree(tn, tn);
  tree->SetDirectory(fFile);
  if (fDoBranchRef)
    tree->BranchRef();
  fTrees.AddLast(tree);  
  return tree;
}

//--------------------------------------------------------------------------------------------------
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(s) and close file.

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

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

  for (Int_t i=0;i<fTrees.GetEntries();++i) {
    MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
    mt->Write(mt->GetName(),TObject::kOverwrite);

    // backup and restore list of branch pointers from TRefTable (needed for autoloading)
    if (mt->GetBranchRef()) {
      TObjArray *parents = mt->GetBranchRef()->GetRefTable()->GetParents();
      TObjArray parentsBak(*parents);
      mt->Reset();
      for (Int_t j=0; j<parentsBak.GetEntries(); ++j)
        parents->Add(parentsBak.At(j));
    }
    else
      mt->Reset();
    mt->SetDirectory(0);   
  }

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

  fIsInit = kFALSE; 
  fFileNumber++;  
}

//--------------------------------------------------------------------------------------------------
const char *TreeWriter::CName(void *obj) const 
{
  // Dereference void* pointer into TObject* pointer

  TObject **sobj = static_cast<TObject**>(obj);
  TObject *tobj = dynamic_cast<TObject*>(*sobj);
  if (tobj==0) {
    Fatal("CName", "Given void* ptr can not be dereferenced into TObject*");
  }
  return tobj->ClassName();
}

//--------------------------------------------------------------------------------------------------
void TreeWriter::DoBranchRef(const char *tn)
{
  // Fill BranchRef for given tree.

  if (fTrees.GetEntries()==0) {
    Error("DoBranchRef", "Tree with name %s not found!", tn); 
    return;
  }

  MyTree *mt = GetMyTree(tn);
  if (!mt)
    return;

  mt->BranchRef();
}

//--------------------------------------------------------------------------------------------------
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;
  }

  Int_t r = 0;
  for (Int_t i=0;i<fTrees.GetEntries();++i) {
    MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
    if (mt->GetAutoFill()==0)
      continue;
    r += mt->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);
}

//--------------------------------------------------------------------------------------------------
Long64_t TreeWriter::GetEntries(const char *tn) const
{ 
  // Return entries of tree with given name. If no tree is given, return sum of entries
  // of all trees.

   if (fTrees.GetEntries()==0) return -1;

   if (tn) {
     const TTree *mt=GetTree(tn);
     if (mt)
       return mt->GetEntries();
     else
       return -1;
   }

   Long64_t ret = 0;
   for (Int_t i=0;i<fTrees.GetEntries();++i) {
      const MyTree *mt = static_cast<const MyTree*>(fTrees.At(i));
      ret += mt->GetEntries();
   }
   return ret;
}

//--------------------------------------------------------------------------------------------------
MyTree *mithep::TreeWriter::GetMyTree(const char *tn)
{
  // Return MyTree with given name from array.

  if (fTrees.GetEntries()==0) 
    return 0;

  TObject *obj = 0;
  if (tn==0) {
    obj = fTrees.At(0);
  } else {
    obj = fTrees.FindObject(tn);
  }
  
  if (obj)  
    return static_cast<MyTree*>(obj);
  return 0;
}   

//--------------------------------------------------------------------------------------------------
const TTree *mithep::TreeWriter::GetTree(const char *tn) const
{
  // Return TTree with given name from array.

  if (fTrees.GetEntries()==0) 
    return 0;

  TObject *obj = 0;
  if (tn==0) {
    obj = fTrees.At(0);
  } else {
    obj = fTrees.FindObject(tn);
  }
  
  if (obj)  
    return dynamic_cast<const TTree*>(obj);
  return 0;
}   

//--------------------------------------------------------------------------------------------------
TTree *mithep::TreeWriter::GetTree(const char *tn)
{
  // Return TTree with given name from array.

  if (fTrees.GetEntries()==0) 
    return 0;

  TObject *obj = 0;
  if (tn==0) {
    obj = fTrees.At(0);
  } else {
    obj = fTrees.FindObject(tn);
  }
  
  if (obj)  
    return dynamic_cast<TTree*>(obj);
  return 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);

  for (Int_t i=0;i<fTrees.GetEntries();++i) {
    MyTree *mt = static_cast<MyTree*>(fTrees.At(i));
    mt->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::SetAutoFill(const char *tn, Bool_t b)
{
  // Set auto-fill mode of tree with given name.

  if (fTrees.GetEntries()==0)
    return;

  MyTree *mt = GetMyTree(tn);
  if (!mt)
    return;

  mt->SetAutoFill(b);
}

//--------------------------------------------------------------------------------------------------
void TreeWriter::SetMaxSize(Long64_t s) 
{ 
  // Set maximum file size. Check if this exceeds the ROOT file size and if, 
  // print a warning and adjust it.

  if (s>=(Long64_t)(0.99 * TTree::GetMaxTreeSize())) {
    Long64_t news = (Long64_t)(s/0.99);
    Warning("SetMaxSize", "Maximum tree size increased from %lld to %lld",
            TTree::GetMaxTreeSize(), news);
    TTree::SetMaxTreeSize(news);
  }

  fMaxSize=s; 
}

//--------------------------------------------------------------------------------------------------
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");
}

//__________________________________________________________________________________________________
void TreeWriter::Terminate()
{
  // Terminate tree file writing.
   
  if (fIsInit) {
    CloseFile();
  }

  TDirectory::TContext context(0); 
  fTrees.Clear();
}
Revision:	1.15
Committed:	Thu Mar 12 18:26:08 2009 UTC (16 years, 1 month ago) by loizides
Content type:	text/plain
Branch:	MAIN
CVS Tags:	Mit_008pre2
Changes since 1.14:	+1 -2 lines
Log Message:	Cosmetics
#	Content
1	// $Id: TreeWriter.cc,v 1.14 2009/03/03 08:37:12 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->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;
251	fFile = 0;
252
253	fIsInit = kFALSE;
254	fFileNumber++;
255	}
256
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
291	// restore the stored object number at the time BeginEvent(kTRUE)
292	// was called.
293
294	if (!fIsInit) {
295	Fatal("EndEvent", "File is not open, did you call BeginEvent?");
296	return kFALSE;
297	}
298
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())
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)
315	TProcessID::SetObjectCount(fEvtObjNum);
316	}
317	}
318
319	return (r >= 0);
320	}
321
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();
412	if (entries < 1) return kFALSE;
413
414	Long64_t avgSize = GetFileSize() / entries;
415
416	if (avgSize < fkMinAvgSize)
417	avgSize = fkMinAvgSize;
418
419	return (GetFileSize() + avgSize + fkMinFreeSpace) > fMaxSize;
420	}
421
422	//--------------------------------------------------------------------------------------------------
423	void TreeWriter::OpenFile()
424	{
425	// Open the file and attach the tree.
426
427	if (fIsInit) {
428	Fatal("OpenFile", "File is already open, call CloseFile first!");
429	return;
430	}
431
432	TDirectory::TContext context(0);
433
434	TString pathname=GetFullName();
435	gSystem->ExpandPathName(pathname);
436
437	fFile = TFile::Open(pathname, "RECREATE");
438	if (fFile == 0) {
439	Fatal("OpenFile", "Could not open file %s", pathname.Data());
440	return;
441	}
442
443	fFile->SetCompressionLevel(fCompressLevel);
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	//--------------------------------------------------------------------------------------------------
454	void TreeWriter::Print(Option_t *option) const
455	{
456	// Print the contents of the tree writer.
457
458	if (option) {
459	cout << ClassName() << " with members " << endl;
460	cout << " fBaseURL: " << fBaseURL << endl;
461	cout << " fPreFix: " << fPrefix << endl;
462	cout << " fFileNumber: " << fFileNumber << endl;
463	cout << " fCompressLevel: " << fCompressLevel << endl;
464	cout << " fDefBrSize: " << fDefBrSize << endl;
465	cout << " fDefSL: " << fDefSL << endl;
466	cout << " fMaxSize: " << fMaxSize << endl;
467	cout << " fDoObjNumReset: " << fDoObjNumReset << endl;
468	return;
469	}
470
471	cout << ClassName() << ": " << GetEntries()
472	<< (GetEntries() == 1 ? " event" : " events") << endl;
473	}
474
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
510	// settings of how the tree was created.
511
512	if (!fIsInit) {
513	Fatal("StoreObject", "Tree is not created, call create first!");
514	return;
515	}
516
517	if (!obj) {
518	Fatal("StoreObject", "Ptr to TObject is null!");
519	return;
520	}
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	}