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