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   fId(0),
 275   fNextName(),
 276   fNextNameCounter(),
 277   fStatus(0)
 278 {
 279   //
 280   // Standard ctor
 281   //
 282 }
 283
 284 TTreeStream::~TTreeStream()
 285 {
 286   //
 287   // Class dtor
 288   //
 289   fElements->Delete();
 290   fBranches->Clear();
 291   delete fElements;
 292   delete fBranches;
 293 }
 294
 295 void TTreeStream::Close()
 296 {
 297   //
 298   // Flush data to disk and close
 299   //
 300   fTree->Write();
 301 }
 302
 303 Int_t TTreeStream::CheckIn(Char_t type, void *pointer)
 304 {
 305   //
 306   // Insert object of given type
 307   //
 308   if (!fElements) fElements = new TObjArray(1000);
 309   TTreeDataElement* element = (TTreeDataElement*)fElements->At(fCurrentIndex);
 310   if (!element) {
 311     element = new TTreeDataElement(type);
 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(pointer);
 328     fElements->AddAt(element,fCurrentIndex);
 329     fCurrentIndex++;
 330     return 0; //new element added
 331   }
 332   if (element->GetType()!=type){
 333     fStatus++;
 334     return 1; //mismatched data element
 335   }
 336   element->SetPointer(pointer);
 337   fCurrentIndex++;
 338   return 0;
 339 }
 340
 341 Int_t TTreeStream::CheckIn(TObject *o){
 342   //
 343   // Insert TObject
 344   //
 345   if (!o) return 0;
 346   if (!fElements) fElements = new TObjArray(1000);
 347   TTreeDataElement* element = (TTreeDataElement*)fElements->At(fCurrentIndex);
 348   if (!element) {
 349     element = new TTreeDataElement(o->IsA());
 350     //
 351     char name[1000];
 352     if (fNextName.Length()>0){
 353       if (fNextNameCounter==0){
 354         sprintf(name,"%s",(const char*)fNextName);
 355       }
 356       if (fNextNameCounter>0){
 357         sprintf(name,"%s%d",(const char*)fNextName,fNextNameCounter);
 358       }
 359     }
 360     else{
 361       sprintf(name,"B%d",fCurrentIndex);
 362     }
 363     element->SetName(name);
 364
 365     element->SetPointer(o);
 366     fElements->AddAt(element,fCurrentIndex);
 367     fCurrentIndex++;
 368     return 0; //new element added
 369   }
 370   if (element->fClass!=o->IsA()){
 371     fStatus++;
 372     return 1; //mismatched data element
 373   }
 374   element->SetPointer(o);
 375   fCurrentIndex++;
 376   return 0;
 377 }
 378
 379 void TTreeStream::BuildTree(){
 380   //
 381   // Build the Tree
 382   //
 383   if (fTree->GetEntries()>0) return;
 384   fTree = new TTree(GetName(),GetName());
 385   Int_t entries = fElements->GetEntriesFast();
 386   fBranches = new TObjArray(entries);
 387
 388   for (Int_t i=0;i<entries;i++){
 389     //
 390     TTreeDataElement* element = (TTreeDataElement*)fElements->At(i);
 391     char bname1[1000];
 392     if (element->GetName()[0]==0){
 393       sprintf(bname1,"B%d",i);
 394     }
 395     else{
 396       sprintf(bname1,element->GetName());
 397     }
 398     if (element->fClass){
 399       if (element->fClass->GetBaseClass("TClonesArray")){
 400         TBranch * br = fTree->Branch(bname1,element->fClass->GetName(),&(element->fPointer));
 401         fBranches->AddAt(br,i);
 402       }else
 403         {
 404           TBranch * br = fTree->Branch(bname1,element->fClass->GetName(),&(element->fPointer));
 405           fBranches->AddAt(br,i);
 406         }
 407     }
 408     if (element->GetType()>0){
 409       char bname2[1000];
 410       sprintf(bname2,"B%d/%c",i,element->GetType());
 411       TBranch * br = fTree->Branch(bname1,element->fPointer,bname2);
 412       fBranches->AddAt(br,i);
 413     }
 414   }
 415 }
 416
 417 void TTreeStream::Fill(){
 418   //
 419   // Fill the tree
 420   //
 421   if (fTree) {
 422     Int_t entries=fElements->GetEntriesFast();
 423     if (entries>fTree->GetNbranches()) BuildTree();
 424     for (Int_t i=0;i<entries;i++){
 425       TTreeDataElement* el  = (TTreeDataElement*)fElements->At(i);
 426       if (!el) continue;
 427       if (!el->GetType()) continue;
 428       TBranch      * br  = (TBranch*)fBranches->At(i);
 429       if (br &&el){
 430         if (el->GetType())  br->SetAddress(el->fPointer);
 431       }
 432     }
 433     if (fStatus==0) fTree->Fill(); //fill only in case of non conflicts
 434     fStatus=0;
 435   }
 436 }
 437
 438 TTreeStream & TTreeStream::Endl()
 439 {
 440   //
 441   // Perform pseudo endl operation
 442   //
 443   if (fTree->GetNbranches()==0) BuildTree();
 444   Fill();
 445   fStatus =0;
 446   fCurrentIndex=0;
 447   return *this;
 448 }
 449
 450
 451 TTreeStream  &TTreeStream::operator<<(Char_t *name)
 452 {
 453   //
 454   // Endl
 455   //
 456   if (name[0]=='\n'){
 457     return Endl();
 458   }
 459   //
 460   //if tree was already defined ignore
 461   if (fTree->GetEntries()>0) return *this;
 462   //check branch name if tree was not
 463   //
 464   Int_t last=0;
 465   for (last=0;;last++){
 466     if (name[last]==0) break;
 467   }
 468
 469   if (last>0&&name[last-1]=='='){
 470     fNextName = name;
 471     fNextName[last-1]=0;
 472     fNextNameCounter=0;
 473   }
 474   return *this;
 475 }
 476