690164465ed256acaf6c787708ccf78fedff1ac3
[u/mrichter/AliRoot.git] / STEER / TTreeStream.cxx
1 #include "TObjArray.h"
2 #include "TFile.h"
3 #include "TTree.h"
4 #include "TBrowser.h"
5 #include "TTreeStream.h"
6
7
8 ClassImp(TTreeDataElement)
9 ClassImp(TTreeStream)
10 ClassImp(TTreeSRedirector)
11
12 /*
13   marian.ivanov@cern.ch
14   //
15   ------------------------------------------------------------------------------------------------
16   TTreeStream
17   Standard stream (cout) like input for the tree
18   Run and see TTreeStreamer::Test() - to see TTreeStreamer functionality
19   ------------------------------------------------------------------------------------------------  
20   //
21   -------------------------------------------------------------------------------------------------
22   TTreeSRedirector
23   Redirect file to  different TTreeStreams  
24   Run and see   TTreeSRedirector::Test() as an example of TTreeSRedirectorer functionality 
25   // 
26 */
27
28
29 void TTreeStream::Test()
30 {
31   //
32   // 
33   TFile *ftest = new TFile("teststreamer.root","recreate");
34   if (!ftest) ftest = new TFile("teststreamer.root","new");
35   //
36   //create to streems Tree1 and Tree2
37   TTreeStream stream1("Tree1");
38   TTreeStream stream2("Tree2");
39   //
40   Char_t ch='s';
41   Float_t f=3.;
42   Float_t f2=1;
43   TObject *po  = new TObject;
44   TObject *po2 = new TObject;
45   for (Int_t i=0;i<100000;i++) {
46     f=i*100;
47     po->SetUniqueID(i);
48     po2->SetUniqueID(i*100);
49     ch=i%120;
50     //
51     //    Stream the data
52     //    The data layout of stream is defined during first invocation of streamer.
53     //    Endl is the trigger which define the end of structure.
54     // 
55     //    The name of branch can be specified using strings with = at the the end
56     //    if string is not specified automatic convention is u (sed B0, B1, ...Bn)
57     stream1<<"i="<<i<<"ch="<<ch<<"f="<<f<<"po="<<po<<"\n";
58     f  = 1./(100.1+i);
59     f2 = -f;     
60     //3.) just another example - we can fill the same tree with different objects
61     //
62     stream2<<f<<po<<"\n";
63     stream2<<f2<<po2<<"\n";
64   }
65   //
66   //4.) Close the streeamers (Write the streamed tree's to the file) and close the corresponding file.
67   //
68   stream1.Close();
69   stream2.Close();
70   ftest->Close();
71   delete ftest;
72   //
73   //5.) and now see results  in file tteststreamer.root
74 }
75
76
77
78 void TTreeSRedirector::Test()
79 {
80   //
81   //Example test function to show functionality of TTreeSRedirector
82   //
83   //
84   //1.)create the  redirector associated with file (testredirector.root)
85   //
86   //
87   TTreeSRedirector *pmistream= new TTreeSRedirector("testredirector.root");
88   TTreeSRedirector &mistream = *pmistream;
89   Char_t ch='s';
90   Float_t f=3.;
91   Float_t f2=1;
92   TObject *po  = new TObject;
93   TObject *po2 = new TObject;
94   for (Int_t i=0;i<100000;i++) {
95     f=i*100;
96     po->SetUniqueID(i);
97     po2->SetUniqueID(i*100);
98     ch=i%120;
99     //
100     //2.) create the tree with identifier specified by first argument
101     //                                layout specified by sequence of arguments
102     //                                Tree identifier has to be specified as first argument !!! 
103     //    if the tree and layout was already defined the consistency if layout is checked
104     //                                if the data are consisten fill given tree 
105     //    the name of branch can be specified using strings with = at the the end
106     //    if string is not specified use automatic convention  B0, B1, ...Bn
107     mistream<<"TreeIdentifier"<<"i="<<i<<"ch="<<ch<<"f="<<f<<"po="<<po<<"\n";
108     f  = 1./(100.1+i);
109     f2 = -f; 
110     
111     //3.) just another example - we can fill the same tree with different objects
112     //
113     mistream<<"TreeK"<<f<<po<<"\n";
114     mistream<<"TreeK"<<f2<<po2<<"\n";
115   }
116   //
117   //4.) write the streamed tree's to the file and close the corresponding file in destructor
118   //
119   delete pmistream;
120   //
121   //5.) and now see results in file testredirector.root 
122 }
123
124
125
126
127
128
129 TTreeSRedirector::TTreeSRedirector(const char *fname){
130   //
131   //
132   fFile = new TFile(fname,"recreate");
133   if (!fFile){
134     fFile = new TFile(fname,"new");
135   }
136   fDataLayouts =0;
137 }
138
139 TTreeSRedirector::~TTreeSRedirector(){
140   //
141   //
142   Close();       //write the tree to the selected file
143   fFile->Close();
144   delete fFile;
145 }
146
147 TTreeStream  & TTreeSRedirector::operator<<(Int_t id)
148 {
149   //
150   // return reference to the data layout with given identifier
151   // if not existing - creates new
152   if (!fDataLayouts) fDataLayouts = new TObjArray(10000);
153   TTreeStream *clayout=0;
154   Int_t entries = fDataLayouts->GetEntriesFast();
155   for (Int_t i=0;i<entries;i++){
156     TTreeStream * layout = (TTreeStream*)fDataLayouts->At(i);
157     if (!layout) continue;
158     if (layout->fId==id) {
159       clayout = layout;
160       break;
161     }
162   }
163   if (!clayout){
164     fFile->cd();
165     char chname[100];
166     sprintf(chname,"Tree%d",id);
167     clayout = new TTreeStream(chname);
168     clayout->fId=id;
169     fDataLayouts->AddAt(clayout,entries);
170   }
171   return *clayout;
172 }
173
174
175 TTreeStream  & TTreeSRedirector::operator<<(const char* name)
176 {
177   //
178   // return reference to the data layout with given identifier
179   // if not existing - creates new
180   if (!fDataLayouts) fDataLayouts = new TObjArray(10000);
181   TTreeStream *clayout=(TTreeStream*)fDataLayouts->FindObject(name);
182   Int_t entries = fDataLayouts->GetEntriesFast();
183
184   if (!clayout){
185     fFile->cd();
186     clayout = new TTreeStream(name);
187     clayout->fId=-1;
188     clayout->SetName(name);
189     fDataLayouts->AddAt(clayout,entries);    
190   }
191   return *clayout;
192 }
193
194
195
196
197 void TTreeSRedirector::Close(){
198   //
199   //
200   TFile * backup = gFile;
201   fFile->cd();
202   if (fDataLayouts){
203     Int_t entries = fDataLayouts->GetEntriesFast();
204     for (Int_t i=0;i<entries;i++){
205       TTreeStream * layout = (TTreeStream*)fDataLayouts->At(i);
206       if (layout){
207         if (layout->fTree) layout->fTree->Write(layout->GetName());
208       }
209     }
210     delete fDataLayouts;
211     fDataLayouts=0;
212   }
213   backup->cd();
214 }
215
216
217
218 //-------------------------------------------------------------
219 TTreeDataElement:: TTreeDataElement(Char_t type){
220   //
221   //
222   fType   = type;
223   fDType  = 0;
224   fClass  = 0;
225   fPointer= 0;
226 }
227 TTreeDataElement:: TTreeDataElement(TDataType* type){
228   //
229   //
230   fType   = 0;
231   fDType  = type;
232   fClass  = 0;
233   fPointer= 0;
234 }
235 TTreeDataElement:: TTreeDataElement(TClass* cl){
236   //
237   //
238   fType   = 0;
239   fDType  = 0;
240   fClass  = cl;
241   fPointer= 0;
242 }
243
244 //-------------------------------------------------------------------
245 TTreeStream::TTreeStream(const char *treename):TNamed(treename,treename){
246   fElements =0;
247   fTree =0;
248   fCurrentIndex =0;
249   fNextName="";
250   fNextNameCounter=0;
251   fTree = new TTree(treename, treename);
252 }
253
254 TTreeStream::~TTreeStream()
255 {
256   fElements->Delete();
257   fBranches->Clear();
258   delete fElements;
259   delete fBranches;
260 }
261
262 void TTreeStream::Close()
263 {
264   //
265   fTree->Write();
266 }
267
268 Int_t TTreeStream::CheckIn(Char_t type, void *pointer)
269 {
270   if (!fElements) fElements = new TObjArray(1000);
271   TTreeDataElement* element = (TTreeDataElement*)fElements->At(fCurrentIndex);
272   if (!element) {
273     element = new TTreeDataElement(type);
274     //
275     char name[1000];
276     if (fNextName.Length()>0){
277       if (fNextNameCounter==0){
278         sprintf(name,"%s",(const char*)fNextName);
279       }
280       if (fNextNameCounter>0){
281         sprintf(name,"%s%d",(const char*)fNextName,fNextNameCounter);
282       }      
283     }
284     else{
285       sprintf(name,"B%d.",fCurrentIndex);
286     }
287     element->SetName(name);
288     //
289     element->SetPointer(pointer);
290     fElements->AddAt(element,fCurrentIndex);
291     fCurrentIndex++;
292     return 0; //new element added
293   }
294   if (element->GetType()!=type){
295     fStatus++;
296     return 1; //mismatched data element
297   }
298   element->SetPointer(pointer);
299   fCurrentIndex++;
300   return 0;
301 }
302
303 Int_t TTreeStream::CheckIn(TObject *o){
304   //
305   //
306   if (!o) return 0;
307   if (!fElements) fElements = new TObjArray(1000);
308   TTreeDataElement* element = (TTreeDataElement*)fElements->At(fCurrentIndex);
309   if (!element) {
310     element = new TTreeDataElement(o->IsA());
311     //
312     char name[1000];
313     if (fNextName.Length()>0){
314       if (fNextNameCounter==0){
315         sprintf(name,"%s",(const char*)fNextName);
316       }
317       if (fNextNameCounter>0){
318         sprintf(name,"%s%d",(const char*)fNextName,fNextNameCounter);
319       }      
320     }
321     else{
322       sprintf(name,"B%d",fCurrentIndex);
323     }
324     element->SetName(name);
325
326     element->SetPointer(o);
327     fElements->AddAt(element,fCurrentIndex);
328     fCurrentIndex++;
329     return 0; //new element added
330   }
331   if (element->fClass!=o->IsA()){
332     fStatus++;
333     return 1; //mismatched data element
334   }
335   element->SetPointer(o);
336   fCurrentIndex++;
337   return 0;  
338 }
339
340 void TTreeStream::BuildTree(){
341   //
342   //
343   if (fTree->GetEntries()>0) return;
344   fTree = new TTree(GetName(),GetName());
345   Int_t entries = fElements->GetEntriesFast();  
346   fBranches = new TObjArray(entries);
347   
348   for (Int_t i=0;i<entries;i++){
349     //
350     TTreeDataElement* element = (TTreeDataElement*)fElements->At(i);
351     char bname1[1000];
352     if (element->GetName()[0]==0){
353       sprintf(bname1,"B%d",i);
354     }
355     else{
356       sprintf(bname1,element->GetName());
357     }
358     if (element->fClass){
359       if (element->fClass->GetBaseClass("TClonesArray")){
360         TBranch * br = fTree->Branch(bname1,&(element->fPointer));
361         fBranches->AddAt(br,i);
362       }else
363         {
364           TBranch * br = fTree->Branch(bname1,element->fClass->GetName(),&(element->fPointer));
365           fBranches->AddAt(br,i);
366         }
367     }
368     if (element->GetType()>0){
369       char bname2[1000];
370       sprintf(bname2,"B%d/%c",i,element->GetType());
371       TBranch * br = fTree->Branch(bname1,element->fPointer,bname2);
372       fBranches->AddAt(br,i);
373     }
374   }
375 }
376
377 void TTreeStream::Fill(){
378   //
379   //
380   if (fTree) { 
381     Int_t entries=fElements->GetEntriesFast();
382     if (entries>fTree->GetNbranches()) BuildTree();
383     for (Int_t i=0;i<entries;i++){    
384       TTreeDataElement* el  = (TTreeDataElement*)fElements->At(i);
385       if (!el) continue;
386       if (!el->GetType()) continue;
387       TBranch      * br  = (TBranch*)fBranches->At(i);
388       if (br &&el){
389         if (el->GetType())  br->SetAddress(el->fPointer);
390       }
391     }
392     if (fStatus==0) fTree->Fill(); //fill only in case of non conflicts
393     fStatus=0;
394   }
395 }
396
397 TTreeStream & TTreeStream::Endl(){
398   if (fTree->GetNbranches()==0) BuildTree();
399   Fill();
400   fStatus =0;
401   fCurrentIndex=0;
402   return *this;
403 }
404
405
406 TTreeStream  &TTreeStream::operator<<(Char_t *name)
407 {
408   //
409   //
410   //
411   //Endl 
412   if (name[0]=='\n'){
413     return Endl();
414   }
415   //
416   //if tree was already defined ignore
417   if (fTree->GetEntries()>0) return *this;
418   //check branch name if tree was not 
419   //
420   Int_t last=0;
421   for (last=0;;last++){
422     if (name[last]==0) break;    
423   }
424   
425   if (last>0&&name[last-1]=='='){
426     fNextName = name;
427     fNextName[last-1]=0;
428     fNextNameCounter=0;
429   }
430   return *this;
431 }
432