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
 148 TTreeStream  & TTreeSRedirector::operator<<(Int_t id)
 149 {
 150   //
 151   // return reference to the data layout with given identifier
 152   // if not existing - creates new
 153   if (!fDataLayouts) fDataLayouts = new TObjArray(10000);
 154   TTreeStream *clayout=0;
 155   Int_t entries = fDataLayouts->GetEntriesFast();
 156   for (Int_t i=0;i<entries;i++){
 157     TTreeStream * layout = (TTreeStream*)fDataLayouts->At(i);
 158     if (!layout) continue;
 159     if (layout->fId==id) {
 160       clayout = layout;
 161       break;
 162     }
 163   }
 164   if (!clayout){
 165     fFile->cd();
 166     char chname[100];
 167     sprintf(chname,"Tree%d",id);
 168     clayout = new TTreeStream(chname);
 169     clayout->fId=id;
 170     fDataLayouts->AddAt(clayout,entries);
 171   }
 172   return *clayout;
 173 }
 174
 175
 176 TTreeStream  & TTreeSRedirector::operator<<(const char* name)
 177 {
 178   //
 179   // return reference to the data layout with given identifier
 180   // if not existing - creates new
 181   if (!fDataLayouts) fDataLayouts = new TObjArray(10000);
 182   TTreeStream *clayout=(TTreeStream*)fDataLayouts->FindObject(name);
 183   Int_t entries = fDataLayouts->GetEntriesFast();
 184
 185   if (!clayout){
 186     fFile->cd();
 187     clayout = new TTreeStream(name);
 188     clayout->fId=-1;
 189     clayout->SetName(name);
 190     fDataLayouts->AddAt(clayout,entries);
 191   }
 192   return *clayout;
 193 }
 194
 195
 196
 197
 198 void TTreeSRedirector::Close(){
 199   //
 200   //
 201   TFile * backup = gFile;
 202   fFile->cd();
 203   if (fDataLayouts){
 204     Int_t entries = fDataLayouts->GetEntriesFast();
 205     for (Int_t i=0;i<entries;i++){
 206       TTreeStream * layout = (TTreeStream*)fDataLayouts->At(i);
 207       if (layout){
 208         if (layout->fTree) layout->fTree->Write(layout->GetName());
 209       }
 210     }
 211     delete fDataLayouts;
 212     fDataLayouts=0;
 213   }
 214   if (backup) backup->cd();
 215 }
 216
 217
 218
 219 //-------------------------------------------------------------
 220 TTreeDataElement:: TTreeDataElement(Char_t type) :
 221   TNamed(),
 222   fType(type),
 223   fDType(0),
 224   fClass(0),
 225   fPointer(0)
 226 {
 227   //
 228   //
 229   //
 230 }
 231
 232 TTreeDataElement:: TTreeDataElement(TDataType* type) :
 233   TNamed(),
 234   fType(0),
 235   fDType(type),
 236   fClass(0),
 237   fPointer(0)
 238 {
 239   //
 240   //
 241   //
 242 }
 243
 244 TTreeDataElement:: TTreeDataElement(TClass* cl) :
 245   TNamed(),
 246   fType(0),
 247   fDType(0),
 248   fClass(cl),
 249   fPointer(0)
 250 {
 251   //
 252   //
 253   //
 254 }
 255
 256 //-------------------------------------------------------------------
 257 TTreeStream::TTreeStream(const char *treename):
 258   TNamed(treename,treename),
 259   fElements(0),
 260   fBranches(0),
 261   fTree(new TTree(treename, treename)),
 262   fCurrentIndex(0),
 263   fNextName(),
 264   fNextNameCounter(),
 265   fStatus(0)
 266 {
 267   //
 268   // Standard ctor
 269   //
 270 }
 271
 272 TTreeStream::~TTreeStream()
 273 {
 274   //
 275   // Class dtor
 276   //
 277   fElements->Delete();
 278   fBranches->Clear();
 279   delete fElements;
 280   delete fBranches;
 281 }
 282
 283 void TTreeStream::Close()
 284 {
 285   //
 286   // Flush data to disk and close
 287   //
 288   fTree->Write();
 289 }
 290
 291 Int_t TTreeStream::CheckIn(Char_t type, void *pointer)
 292 {
 293   //
 294   // Insert object of given type
 295   //
 296   if (!fElements) fElements = new TObjArray(1000);
 297   TTreeDataElement* element = (TTreeDataElement*)fElements->At(fCurrentIndex);
 298   if (!element) {
 299     element = new TTreeDataElement(type);
 300     //
 301     char name[1000];
 302     if (fNextName.Length()>0){
 303       if (fNextNameCounter==0){
 304         sprintf(name,"%s",(const char*)fNextName);
 305       }
 306       if (fNextNameCounter>0){
 307         sprintf(name,"%s%d",(const char*)fNextName,fNextNameCounter);
 308       }
 309     }
 310     else{
 311       sprintf(name,"B%d.",fCurrentIndex);
 312     }
 313     element->SetName(name);
 314     //
 315     element->SetPointer(pointer);
 316     fElements->AddAt(element,fCurrentIndex);
 317     fCurrentIndex++;
 318     return 0; //new element added
 319   }
 320   if (element->GetType()!=type){
 321     fStatus++;
 322     return 1; //mismatched data element
 323   }
 324   element->SetPointer(pointer);
 325   fCurrentIndex++;
 326   return 0;
 327 }
 328
 329 Int_t TTreeStream::CheckIn(TObject *o){
 330   //
 331   // Insert TObject
 332   //
 333   if (!o) return 0;
 334   if (!fElements) fElements = new TObjArray(1000);
 335   TTreeDataElement* element = (TTreeDataElement*)fElements->At(fCurrentIndex);
 336   if (!element) {
 337     element = new TTreeDataElement(o->IsA());
 338     //
 339     char name[1000];
 340     if (fNextName.Length()>0){
 341       if (fNextNameCounter==0){
 342         sprintf(name,"%s",(const char*)fNextName);
 343       }
 344       if (fNextNameCounter>0){
 345         sprintf(name,"%s%d",(const char*)fNextName,fNextNameCounter);
 346       }
 347     }
 348     else{
 349       sprintf(name,"B%d",fCurrentIndex);
 350     }
 351     element->SetName(name);
 352
 353     element->SetPointer(o);
 354     fElements->AddAt(element,fCurrentIndex);
 355     fCurrentIndex++;
 356     return 0; //new element added
 357   }
 358   if (element->fClass!=o->IsA()){
 359     fStatus++;
 360     return 1; //mismatched data element
 361   }
 362   element->SetPointer(o);
 363   fCurrentIndex++;
 364   return 0;
 365 }
 366
 367 void TTreeStream::BuildTree(){
 368   //
 369   // Build the Tree
 370   //
 371   if (fTree->GetEntries()>0) return;
 372   fTree = new TTree(GetName(),GetName());
 373   Int_t entries = fElements->GetEntriesFast();
 374   fBranches = new TObjArray(entries);
 375
 376   for (Int_t i=0;i<entries;i++){
 377     //
 378     TTreeDataElement* element = (TTreeDataElement*)fElements->At(i);
 379     char bname1[1000];
 380     if (element->GetName()[0]==0){
 381       sprintf(bname1,"B%d",i);
 382     }
 383     else{
 384       sprintf(bname1,element->GetName());
 385     }
 386     if (element->fClass){
 387       if (element->fClass->GetBaseClass("TClonesArray")){
 388         TBranch * br = fTree->Branch(bname1,element->fClass->GetName(),&(element->fPointer));
 389         fBranches->AddAt(br,i);
 390       }else
 391         {
 392           TBranch * br = fTree->Branch(bname1,element->fClass->GetName(),&(element->fPointer));
 393           fBranches->AddAt(br,i);
 394         }
 395     }
 396     if (element->GetType()>0){
 397       char bname2[1000];
 398       sprintf(bname2,"B%d/%c",i,element->GetType());
 399       TBranch * br = fTree->Branch(bname1,element->fPointer,bname2);
 400       fBranches->AddAt(br,i);
 401     }
 402   }
 403 }
 404
 405 void TTreeStream::Fill(){
 406   //
 407   // Fill the tree
 408   //
 409   if (fTree) {
 410     Int_t entries=fElements->GetEntriesFast();
 411     if (entries>fTree->GetNbranches()) BuildTree();
 412     for (Int_t i=0;i<entries;i++){
 413       TTreeDataElement* el  = (TTreeDataElement*)fElements->At(i);
 414       if (!el) continue;
 415       if (!el->GetType()) continue;
 416       TBranch      * br  = (TBranch*)fBranches->At(i);
 417       if (br &&el){
 418         if (el->GetType())  br->SetAddress(el->fPointer);
 419       }
 420     }
 421     if (fStatus==0) fTree->Fill(); //fill only in case of non conflicts
 422     fStatus=0;
 423   }
 424 }
 425
 426 TTreeStream & TTreeStream::Endl()
 427 {
 428   //
 429   // Perform pseudo endl operation
 430   //
 431   if (fTree->GetNbranches()==0) BuildTree();
 432   Fill();
 433   fStatus =0;
 434   fCurrentIndex=0;
 435   return *this;
 436 }
 437
 438
 439 TTreeStream  &TTreeStream::operator<<(Char_t *name)
 440 {
 441   //
 442   // Endl
 443   //
 444   if (name[0]=='\n'){
 445     return Endl();
 446   }
 447   //
 448   //if tree was already defined ignore
 449   if (fTree->GetEntries()>0) return *this;
 450   //check branch name if tree was not
 451   //
 452   Int_t last=0;
 453   for (last=0;;last++){
 454     if (name[last]==0) break;
 455   }
 456
 457   if (last>0&&name[last-1]=='='){
 458     fNextName = name;
 459     fNextName[last-1]=0;
 460     fNextNameCounter=0;
 461   }
 462   return *this;
 463 }
 464