New classes to facilitate features extraction (M.Ivanov)
[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
140 TTreeSRedirector::~TTreeSRedirector(){
141   //
142   //
143   Close();       //write the tree to the selected file
144   fFile->Close();
145   delete fFile;
146 }
147
148 TTreeStream  & TTreeSRedirector::operator<<(Int_t id)
149 {
150   //
151   // return reference to the data layout with given identifier
152   // if not existing - creates new
153   if (!fDataLayouts) fDataLayouts = new TObjArray(10000);
154   TTreeStream *clayout=0;
155   Int_t entries = fDataLayouts->GetEntriesFast();
156   for (Int_t i=0;i<entries;i++){
157     TTreeStream * layout = (TTreeStream*)fDataLayouts->At(i);
158     if (!layout) continue;
159     if (layout->fId==id) {
160       clayout = layout;
161       break;
162     }
163   }
164   if (!clayout){
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=0;
182   Int_t hash = TMath::Hash(name);
183   Int_t entries = fDataLayouts->GetEntriesFast();
184   for (Int_t i=0;i<entries;i++){
185     TTreeStream * layout = (TTreeStream*)fDataLayouts->At(i);
186     if (!layout) continue;
187     if (layout->fHash==hash) {
188       clayout = layout;
189       break;
190     }
191   }
192   if (!clayout){
193     clayout = new TTreeStream(name);
194     clayout->fId=-1;
195     clayout->SetName(name);
196     clayout->fHash = hash;
197     fDataLayouts->AddAt(clayout,entries);    
198   }
199   return *clayout;
200 }
201
202
203
204
205 void TTreeSRedirector::Close(){
206   //
207   //
208   Int_t entries = fDataLayouts->GetEntriesFast();
209   for (Int_t i=0;i<entries;i++){
210      TTreeStream * layout = (TTreeStream*)fDataLayouts->At(i);
211      if (layout){
212        if (layout->fTree) layout->fTree->Write();
213      }
214   }
215 }
216
217
218
219 //-------------------------------------------------------------
220 TTreeDataElement:: TTreeDataElement(Char_t type){
221   //
222   //
223   fType   = type;
224   fDType  = 0;
225   fClass  = 0;
226   fPointer= 0;
227 }
228 TTreeDataElement:: TTreeDataElement(TDataType* type){
229   //
230   //
231   fType   = 0;
232   fDType  = type;
233   fClass  = 0;
234   fPointer= 0;
235 }
236 TTreeDataElement:: TTreeDataElement(TClass* cl){
237   //
238   //
239   fType   = 0;
240   fDType  = 0;
241   fClass  = cl;
242   fPointer= 0;
243 }
244
245 //-------------------------------------------------------------------
246 TTreeStream::TTreeStream(const char *treename):TNamed(treename,treename){
247   fElements =0;
248   fTree =0;
249   fCurrentIndex =0;
250   fNextName="";
251   fNextNameCounter=0;
252   fTree = new TTree(treename, treename);
253 }
254
255 TTreeStream::~TTreeStream()
256 {
257   fElements->Delete();
258   fBranches->Clear();
259   delete fElements;
260   delete fBranches;
261 }
262
263 void TTreeStream::Close()
264 {
265   //
266   fTree->Write();
267 }
268
269 Int_t TTreeStream::CheckIn(Char_t type, void *pointer)
270 {
271   if (!fElements) fElements = new TObjArray(1000);
272   TTreeDataElement* element = (TTreeDataElement*)fElements->At(fCurrentIndex);
273   if (!element) {
274     element = new TTreeDataElement(type);
275     //
276     char name[1000];
277     if (fNextName.Length()>0){
278       if (fNextNameCounter==0){
279         sprintf(name,"%s",(const char*)fNextName);
280       }
281       if (fNextNameCounter>0){
282         sprintf(name,"%s%d",(const char*)fNextName,fNextNameCounter);
283       }      
284     }
285     else{
286       sprintf(name,"B%d.",fCurrentIndex);
287     }
288     element->SetName(name);
289     //
290     element->SetPointer(pointer);
291     fElements->AddAt(element,fCurrentIndex);
292     fCurrentIndex++;
293     return 0; //new element added
294   }
295   if (element->GetType()!=type){
296     fStatus++;
297     return 1; //mismatched data element
298   }
299   element->SetPointer(pointer);
300   fCurrentIndex++;
301   return 0;
302 }
303
304 Int_t TTreeStream::CheckIn(TObject *o){
305   //
306   //
307   if (!o) return 0;
308   if (!fElements) fElements = new TObjArray(1000);
309   TTreeDataElement* element = (TTreeDataElement*)fElements->At(fCurrentIndex);
310   if (!element) {
311     element = new TTreeDataElement(o->IsA());
312     //
313     char name[1000];
314     if (fNextName.Length()>0){
315       if (fNextNameCounter==0){
316         sprintf(name,"%s",(const char*)fNextName);
317       }
318       if (fNextNameCounter>0){
319         sprintf(name,"%s%d",(const char*)fNextName,fNextNameCounter);
320       }      
321     }
322     else{
323       sprintf(name,"B%d",fCurrentIndex);
324     }
325     element->SetName(name);
326
327     element->SetPointer(o);
328     fElements->AddAt(element,fCurrentIndex);
329     fCurrentIndex++;
330     return 0; //new element added
331   }
332   if (element->fClass!=o->IsA()){
333     fStatus++;
334     return 1; //mismatched data element
335   }
336   element->SetPointer(o);
337   fCurrentIndex++;
338   return 0;  
339 }
340
341 void TTreeStream::BuildTree(){
342   //
343   //
344   if (fTree->GetEntries()>0) return;
345   fTree = new TTree(GetName(),GetName());
346   Int_t entries = fElements->GetEntriesFast();  
347   fBranches = new TObjArray(entries);
348   
349   for (Int_t i=0;i<entries;i++){
350     //
351     TTreeDataElement* element = (TTreeDataElement*)fElements->At(i);
352     char bname1[1000];
353     if (element->GetName()[0]==0){
354       sprintf(bname1,"B%d",i);
355     }
356     else{
357       sprintf(bname1,element->GetName());
358     }
359     if (element->fClass){
360       TBranch * br = fTree->Branch(bname1,element->fClass->GetName(),&(element->fPointer));
361       fBranches->AddAt(br,i);
362     }
363     if (element->GetType()>0){
364       char bname2[1000];
365       sprintf(bname2,"B%d/%c",i,element->GetType());
366       TBranch * br = fTree->Branch(bname1,element->fPointer,bname2);
367       fBranches->AddAt(br,i);
368     }
369   }
370 }
371
372 void TTreeStream::Fill(){
373   //
374   //
375   if (fTree) { 
376     Int_t entries=fElements->GetEntriesFast();
377     if (entries>fTree->GetNbranches()) BuildTree();
378     for (Int_t i=0;i<entries;i++){    
379       TTreeDataElement* el  = (TTreeDataElement*)fElements->At(i);
380       if (!el) continue;
381       if (!el->GetType()) continue;
382       TBranch      * br  = (TBranch*)fBranches->At(i);
383       if (br &&el){
384         if (el->GetType())  br->SetAddress(el->fPointer);
385       }
386     }
387     if (fStatus==0) fTree->Fill(); //fill only in case of non conflicts
388     fStatus=0;
389   }
390 }
391
392 TTreeStream & TTreeStream::Endl(){
393   if (fTree->GetNbranches()==0) BuildTree();
394   Fill();
395   fStatus =0;
396   fCurrentIndex=0;
397   return *this;
398 }
399
400
401 TTreeStream  &TTreeStream::operator<<(Char_t *name)
402 {
403   //
404   //
405   //
406   //Endl 
407   if (name[0]=='\n'){
408     return Endl();
409   }
410   //
411   //if tree was already defined ignore
412   if (fTree->GetEntries()>0) return *this;
413   //check branch name if tree was not 
414   //
415   Int_t last=0;
416   for (last=0;;last++){
417     if (name[last]==0) break;    
418   }
419   
420   if (last>0&&name[last-1]=='='){
421     fNextName = name;
422     fNextName[last-1]=0;
423     fNextNameCounter=0;
424   }
425   return *this;
426 }
427