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