STEER/TTreeStream.cxx

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