STEER/STEERBase/TTreeStream.cxx

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