STrackCutStudy/macros/MergeFiles.C

0001 // ----------------------------------------------------------------------------
0002 // 'MergeFiles.C'
0003 // Derek Anderson
0004 // 09.25.2020
0005 //
0006 // Merges root files.  Adapted from code
0007 // by Nihar Sahoo.
0008 // ----------------------------------------------------------------------------
0009
0010 #include <fstream>
0011 #include <string.h>
0012 #include "TChain.h"
0013 #include "TFile.h"
0014 #include "TH1.h"
0015 #include "TTree.h"
0016 #include "TKey.h"
0017 #include "Riostream.h"
0018
0019
0020
0021 // forward declaration of 'MergeRootFile(TDirectory*, TList*)'
0022 void MergeRootfile(TDirectory *target, TList *sourcelist);
0023
0024
0025
0026 void MergeFiles(const Int_t nFiles=100, const TString filelist= "file.list", const TString outputrootfiles ="test.root") {
0027
0028   ifstream files;
0029   //files.open("file2.list");
0030   files.open(filelist);
0031
0032   char file_stm[800];
0033   int line=0;
0034
0035   TList *FileList;
0036   TFile *Target;
0037
0038
0039   TString filename_output;
0040   TString chain_;
0041   Int_t start_num = 0;
0042   Int_t end_num = 4;
0043
0044   Target = TFile::Open(outputrootfiles, "RECREATE");
0045   FileList = new TList();
0046
0047   while(files.good()) {
0048         files >> file_stm;
0049
0050             // quick fix [Derek, 06.02.2017]
0051             if ((line + 1) > nFiles) {
0052               cout << "WARNING: Whoah there, tiger! Only need " << nFiles << " files!\n"
0053                    << "         Moving on to processing files."
0054                    << endl;
0055               break;
0056             }
0057
0058         //___________________
0059
0060         TFile f(file_stm);
0061
0062         //      Target = TFile::Open( chain_, "RECREATE" );
0063
0064
0065         Bool_t fOk = f.IsZombie();
0066         cout<<"file# "<<line<<"  "<< file_stm<<"  zombi= "<<fOk<<endl;
0067
0068         if(!fOk)
0069           {
0070         FileList->Add(TFile::Open(file_stm));
0071           }
0072
0073
0074
0075         if (!files.good()) break;
0076         line++;
0077       }
0078     MergeRootfile( Target, FileList );
0079
0080
0081 }  // end 'TriggerResolutionStudy_MergeFiles(Int_t, TString, TString)'
0082
0083
0084
0085 void MergeRootfile( TDirectory *target, TList *sourcelist ) {
0086
0087   //  cout << "Target path: " << target->GetPath() << endl;
0088   TString path( (char*)strstr( target->GetPath(), ":" ) );
0089   path.Remove( 0, 2 );
0090
0091   TFile *first_source = (TFile*)sourcelist->First();
0092   first_source->cd( path );
0093   TDirectory *current_sourcedir = gDirectory;
0094   //gain time, do not add the objects in the list in memory
0095   Bool_t status = TH1::AddDirectoryStatus();
0096   TH1::AddDirectory(kFALSE);
0097
0098   // loop over all keys in this directory
0099   TChain *globChain = 0;
0100   TIter nextkey( current_sourcedir->GetListOfKeys() );
0101   TKey *key, *oldkey=0;
0102   while ( (key = (TKey*)nextkey())) {
0103
0104     //keep only the highest cycle number for each key
0105     if (oldkey && !strcmp(oldkey->GetName(),key->GetName())) continue;
0106
0107     // read object from first source file
0108     first_source->cd( path );
0109     TObject *obj = key->ReadObj();
0110
0111     if ( obj->IsA()->InheritsFrom( TH1::Class() ) ) {
0112       // descendant of TH1 -> merge it
0113
0114       //      cout << "Merging histogram " << obj->GetName() << endl;
0115       TH1 *h1 = (TH1*)obj;
0116
0117       // loop over all source files and add the content of the
0118       // correspondant histogram to the one pointed to by "h1"
0119       TFile *nextsource = (TFile*)sourcelist->After( first_source );
0120       while ( nextsource ) {
0121
0122     // make sure we are at the correct directory level by cd'ing to path
0123     nextsource->cd( path );
0124     TKey *key2 = (TKey*)gDirectory->GetListOfKeys()->FindObject(h1->GetName());
0125     if (key2) {
0126       TH1 *h2 = (TH1*)key2->ReadObj();
0127       h1->Add( h2 );
0128       delete h2;
0129     }
0130
0131     nextsource = (TFile*)sourcelist->After( nextsource );
0132       }
0133     }
0134     else if ( obj->IsA()->InheritsFrom( TTree::Class() ) ) {
0135
0136       // loop over all source files create a chain of Trees "globChain"
0137       const char* obj_name= obj->GetName();
0138
0139       globChain = new TChain(obj_name);
0140       globChain->Add(first_source->GetName());
0141       TFile *nextsource = (TFile*)sourcelist->After( first_source );
0142       //      const char* file_name = nextsource->GetName();
0143       // cout << "file name  " << file_name << endl;
0144       while ( nextsource ) {
0145
0146     globChain->Add(nextsource->GetName());
0147     nextsource = (TFile*)sourcelist->After( nextsource );
0148       }
0149
0150     } else if ( obj->IsA()->InheritsFrom( TDirectory::Class() ) ) {
0151       // it's a subdirectory
0152
0153       cout << "Found subdirectory " << obj->GetName() << endl;
0154
0155       // create a new subdir of same name and title in the target file
0156       target->cd();
0157       TDirectory *newdir = target->mkdir( obj->GetName(), obj->GetTitle() );
0158
0159       // newdir is now the starting point of another round of merging
0160       // newdir still knows its depth within the target file via
0161       // GetPath(), so we can still figure out where we are in the recursion
0162       MergeRootfile( newdir, sourcelist );
0163
0164     } else {
0165
0166       // object is of no type that we know or can handle
0167       cout << "Unknown object type, name: "
0168        << obj->GetName() << " title: " << obj->GetTitle() << endl;
0169     }
0170
0171     // now write the merged histogram (which is "in" obj) to the target file
0172     // note that this will just store obj in the current directory level,
0173     // which is not persistent until the complete directory itself is stored
0174     // by "target->Write()" below
0175     if ( obj ) {
0176       target->cd();
0177
0178       //!!if the object is a tree, it is stored in globChain...
0179       if(obj->IsA()->InheritsFrom( TTree::Class() ))
0180     globChain->Merge(target->GetFile(),0,"keep");
0181       else
0182     obj->Write( key->GetName() );
0183     }
0184
0185   } // while ( ( TKey *key = (TKey*)nextkey() ) )
0186
0187   // save modifications to target file
0188   target->SaveSelf(kTRUE);
0189   TH1::AddDirectory(status);
0190 }  // end 'MergeRootFiles(TDirectory*, TList*)'
0191
0192 // end ------------------------------------------------------------------------