STEER/TTreeStream.cxx

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