tests_dav/clitkMergeRootFiles.cxx

   1 /**
   2    =================================================
   3    * @file   clitkMergeRootFiles.cxx
   4    * @author David Sarrut <david.sarrut@creatis.insa-lyon.fr>
   5    * @date   01 Apr 2009
   6    *
   7    * @brief
   8    *
   9    =================================================*/
  10
  11 #include "clitkMergeRootFiles_ggo.h"
  12 #include "clitkCommon.h"
  13 #include <string>
  14 #include <TROOT.h>
  15 #include <TPluginManager.h>
  16 #include <TFile.h>
  17 #include <TKey.h>
  18 #include <TFileMerger.h>
  19 #include <TTree.h>
  20 #include <TChain.h>
  21 #include <TH1.h>
  22 #include <TH2.h>
  23 #include <iostream>
  24 #include <set>
  25
  26 using std::endl;
  27 using std::cout;
  28
  29 typedef std::list<std::string> Strings;
  30 typedef std::list<TFile*> Handles;
  31 typedef std::set<std::string> StringsSet;
  32
  33 void mergeSingleTree(TChain* chain, TFile* output_handle, const std::string& tree_name)
  34 {
  35         assert(chain->FindBranch("runID"));
  36         assert(chain->FindBranch("eventID"));
  37         assert(chain->FindBranch("time"));
  38
  39         int eventID = -1;
  40         int runID = -1;
  41         double time = -1;
  42
  43         chain->SetBranchAddress("eventID",&eventID);
  44         chain->SetBranchAddress("runID",&runID);
  45         chain->SetBranchAddress("time",&time);
  46
  47         output_handle->cd();
  48         TTree* output_tree = chain->CloneTree(0);
  49
  50         const long int nentries = chain->GetEntries();
  51         for (long int kk=0; kk<nentries; kk++)
  52         {
  53                 chain->GetEntry(kk);
  54                 //cout << kk << "/" << chain->GetTreeNumber() << endl;
  55                 output_tree->Fill();
  56         }
  57
  58         output_tree->Write();
  59         delete output_tree;
  60 }
  61
  62 void mergeGateTree(TChain* chain, TFile* output_handle, const std::string& tree_name)
  63 {
  64 }
  65
  66 void mergeCoinTree(TChain* chain, TFile* output_handle, const std::string& tree_name)
  67 {
  68 }
  69
  70 void mergePetRoot(const Strings& input_filenames, const std::string& output_filename)
  71 {
  72         TFile* output_handle = TFile::Open(output_filename.c_str(),"RECREATE");
  73
  74         Handles input_handles;
  75         StringsSet tree_names;
  76         StringsSet h1_names;
  77         StringsSet h2_names;
  78         for (Strings::const_iterator iter=input_filenames.begin(); iter!=input_filenames.end(); iter++)
  79         {
  80                 TFile* handle = TFile::Open(iter->c_str(),"READ");
  81                 assert(handle);
  82
  83                 TIter key_next(handle->GetListOfKeys());
  84                 while (TKey* key = static_cast<TKey*>(key_next()))
  85                 {
  86                         const std::string name = key->GetName();
  87                         const bool is_tree = key->ReadObj()->InheritsFrom("TTree");
  88                         const bool is_h1 = key->ReadObj()->InheritsFrom("TH1");
  89                         const bool is_h2 = key->ReadObj()->InheritsFrom("TH2");
  90                         if (is_tree) tree_names.insert(name);
  91                         if (is_h1) { h1_names.insert(name); cout << name << "## " << static_cast<TH1D*>(key->ReadObj())->GetMean() << endl; }
  92                         if (is_h2) h2_names.insert(name);
  93                 }
  94
  95                 input_handles.push_back(handle);
  96         }
  97
  98         for (Handles::const_iterator iter=input_handles.begin(); iter!=input_handles.end(); iter++)
  99         {
 100                 (*iter)->Close();
 101                 delete *iter;
 102         }
 103
 104         cout << "found " << tree_names.size() << " tree "<< h1_names.size() << " histo1d " << h2_names.size() << " histo2d" << endl;
 105
 106         {
 107                 cout << "merging trees" << endl;
 108                 for (StringsSet::const_iterator iter_tree=tree_names.begin(); iter_tree!=tree_names.end(); iter_tree++)
 109                 {
 110                         cout << "  tree " << *iter_tree;
 111                         TChain* chain = new TChain(iter_tree->c_str());
 112                         for (Strings::const_iterator iter_filename=input_filenames.begin(); iter_filename!=input_filenames.end(); iter_filename++)
 113                                 chain->Add(iter_filename->c_str());
 114                         cout << " nentries " << chain->GetEntries() << endl;
 115
 116                         if ((*iter_tree) == "Singles" || (*iter_tree) == "Hits") mergeSingleTree(chain,output_handle,*iter_tree);
 117                         if ((*iter_tree) == "Gate") mergeGateTree(chain,output_handle,*iter_tree);
 118                         if ((*iter_tree) == "Coincidences") mergeCoinTree(chain,output_handle,*iter_tree);
 119
 120                         delete chain;
 121                 }
 122         }
 123
 124
 125 }
 126
 127
 128 //-----------------------------------------------------------------------------
 129 int main(int argc, char * argv[]) {
 130
 131   gROOT->GetPluginManager()->AddHandler("TVirtualStreamerInfo", "*",
 132       "TStreamerInfo", "RIO", "TStreamerInfo()");
 133
 134   // init command line
 135   GGO(clitkMergeRootFiles, args_info);
 136
 137   // Check parameters
 138   if (args_info.input_given < 2) {
 139     FATAL("Error, please provide at least two inputs files");
 140   }
 141
 142   { // Detect Pet output
 143           bool all_pet_output = true;
 144           Strings input_filenames;
 145           for (uint i=0; i<args_info.input_given; i++)
 146           {
 147                   const char* filename = args_info.input_arg[i];
 148                   input_filenames.push_back(filename);
 149                   TFile* handle = TFile::Open(filename,"READ");
 150                   assert(handle);
 151                   TTree* hits = dynamic_cast<TTree*>(handle->Get("Hits"));
 152                   TTree* singles = dynamic_cast<TTree*>(handle->Get("Singles"));
 153                   const bool is_pet_output = (hits!=NULL) && (singles!=NULL);
 154                   cout << "testing " << filename << " is_pet_output " << is_pet_output << endl;
 155                   handle->Close();
 156                   delete handle;
 157                   all_pet_output &= is_pet_output;
 158           }
 159           cout << "all_pet_output " << all_pet_output << endl;
 160
 161           if (all_pet_output)
 162           {
 163                   mergePetRoot(input_filenames,args_info.output_arg);
 164                   return 0;
 165           }
 166   }
 167
 168
 169   // Merge
 170   TFileMerger * merger = new TFileMerger;
 171   for (uint i=0; i<args_info.input_given; i++) merger->AddFile(args_info.input_arg[i]);
 172   merger->OutputFile(args_info.output_arg);
 173   merger->Merge();
 174
 175   // this is the end my friend
 176   return 0;
 177 }
 178 //-----------------------------------------------------------------------------