Fix
[u/mrichter/AliRoot.git] / PWG4 / TwoPartCorr / mergeCorr.C
1 // $Id:$
2
3 #if !defined(__CINT__) || defined(__MAKECINT__)
4 #include <vector>
5 #include <map>
6 #include <Riostream.h>
7 #include <TChain.h>
8 #include <TClonesArray.h>
9 #include <TError.h>
10 #include <TFile.h>
11 #include <TDirectory.h>
12 #include <TMath.h>
13 #include <TRandom3.h>
14 #include <TSystem.h>
15 #include <TObjectTable.h>
16 #include "TreeClasses.C"
17 #endif
18
19 #define DUPLICATECHECK
20 #ifdef DUPLICATECHECK
21 #define DUPLICATEVERBOSE
22 #endif
23
24 struct EvInfo
25 {
26   Int_t         fRun;
27   Short_t       fArrId;
28   UInt_t        fEntry; 
29 };
30
31
32 void mergeCorr(Int_t run_from=-1, 
33                Int_t run_to=-1,
34                const char *outPrefix = "merged",
35                const char *inFileNames = "res/output_*.root");
36
37 void mergeCorr(Int_t nEvents,
38                const char *outFileName,
39                const char *inFileNames);
40
41 //-----------------------------------------------------------------------------------------------------
42
43 void mergeCorr(Int_t run_from,
44                Int_t run_to,
45                const char *outPrefix,
46                const char *inFileNames)
47 {
48   TChain *c = new TChain("MyTree");
49   c->Add(inFileNames);
50
51   std::map<Int_t,Int_t> goodruns;
52
53   TObjArray *filenames = c->GetListOfFiles();
54   Int_t nfiles = filenames->GetEntries();
55   for(Int_t i=0; i<nfiles; ++i) {
56     TString fname(filenames->At(i)->GetTitle());
57     TFile *file = TFile::Open(fname);
58     if (!file)
59       continue;
60     TList *output = dynamic_cast<TList*>(file->Get("output"));
61     output->SetOwner(1);
62     TTree *tree = dynamic_cast<TTree*>(output->FindObject("MyTree"));
63     Int_t nent = tree->GetEntries();
64     if (nent<1) {
65       delete output;
66       file->Close();
67       delete file;
68       continue;
69     }
70
71     MyHeader *header = 0;
72     TBranch *br = tree->GetBranch("header");
73     br->SetAddress(&header);
74
75     std::vector<Int_t> multruns;
76
77     Int_t runno = -1;
78     for (Int_t ev=0;ev<nent;++ev) {
79       br->GetEntry(ev);
80       if (header->fRun==0) {
81         continue;
82       }
83       if (runno<0) 
84         runno = header->fRun;
85       else if (runno!=header->fRun) {
86         Bool_t found = 0;
87         for (UInt_t j=0;j<multruns.size();++j) {
88           if (multruns.at(j)==header->fRun) {
89             found = 1;
90             break;
91           }
92         }
93         if (!found)
94           multruns.push_back(header->fRun);
95       }
96     }
97     if (multruns.size()<=0) {
98       if ((runno>run_from) && (run_to<0 || runno<run_to)) {
99         cout << "Found run " << runno << " associated to " << fname << endl;
100         goodruns.insert(pair<Int_t,Int_t>(runno,0));
101         TString dir(gSystem->DirName(inFileNames));
102         dir+="/runs";
103         TString link(Form("%s/%s",gSystem->pwd(),fname.Data()));
104         TString rundir(Form("%s/%d", dir.Data(), runno));
105         TString cmd(Form("mkdir -p %s && cd %s && ln -sf %s", 
106                          rundir.Data(), rundir.Data(), link.Data()));
107         gSystem->Exec(cmd);
108       }
109     } else {
110       cout << "Multiple runs (" << runno;
111       for (UInt_t j=0;j<multruns.size();++j)
112         cout << ", " << multruns.at(j);
113       cout << ") found in " << fname << endl;
114       cout << " Cannot deal with this case yet!!!" << endl;
115     }
116     delete output;
117     file->Close();
118     delete file;
119   }
120
121   for (map<Int_t, Int_t>::iterator it = goodruns.begin();it!=goodruns.end();++it) {
122     TString infiles(Form("%s/runs/%d/*.root",gSystem->DirName(inFileNames),it->first));
123     TString outdir(Form("%s/mergedruns",gSystem->DirName(inFileNames)));
124     TString outFileName(Form("%s/%s_run%d.root",outdir.Data(),outPrefix,it->first));
125     gSystem->Exec(Form("mkdir -p %s",outdir.Data()));
126     gSystem->Exec(Form("root -b -q mergeCorr.C+'(-1,\"%s\",\"%s\")'",
127                        outFileName.Data(),infiles.Data()));
128   }
129 }
130
131 //-----------------------------------------------------------------------------------------------------
132
133 void mergeCorr(Int_t nEvents,
134                const char *outFileName,
135                const char *inFileNames)
136 {
137   TChain *c = new TChain("MyTree");
138   c->Add(inFileNames);
139
140   map<ULong64_t, EvInfo*> einfos;
141 #ifdef DUPLICATEVERBOSE
142   map<ULong64_t, MyHeader*> eheaders;
143 #endif
144
145   Int_t totnent = 0;
146   Int_t totnsus = 0;
147 #ifdef DUPLICATECHECK
148   Int_t totncan = 0;
149 #ifdef DUPLICATEVERBOSE
150   Int_t totndup = 0;
151 #endif
152 #endif
153
154   TObjArray *maps = new TObjArray;
155   maps->SetOwner(1);
156   Int_t lrun = 0;
157   TObjArray *filenames = c->GetListOfFiles();
158   Int_t nfiles = filenames->GetEntries();
159   for(Int_t i=0; i<nfiles; ++i) {
160     TString fname(filenames->At(i)->GetTitle());
161     TFile *file = TFile::Open(fname);
162     if (!file)
163       continue;
164     TList *output = dynamic_cast<TList*>(file->Get("output"));
165     output->SetOwner(1);
166     TTree *tree = dynamic_cast<TTree*>(output->FindObject("MyTree"));
167     Int_t nent = tree->GetEntries();
168     if (nent<1) {
169       delete output;
170       file->Close();
171       delete file;
172       continue;
173     }
174
175     totnent += nent;
176     cout << "Found "<< nent << " entries in " << fname << endl;
177     TObjArray *filesnames = dynamic_cast<TObjArray*>(output->FindObject("filesnames"));
178
179     MyHeader *header = 0;
180     TBranch *br = tree->GetBranch("header");
181     br->SetAddress(&header);
182
183     for (Int_t ev=0;ev<nent;++ev) {
184       br->GetEntry(ev);
185       if (header->fRun==0) {
186 #ifdef DUPLICATEVERBOSE
187         cout << "--> Suspicious found for " << header->fRun << " " 
188              << header->fTime << " " << header->fPeriod << " " << header->fBx << endl;
189 #endif
190         ++totnsus;
191         continue;
192       }
193
194       //ULong64_t evtid = header->fTime;
195       ULong64_t evtid = header->GetEventId();
196 #ifdef DUPLICATECHECK
197       if (einfos.find(evtid)!=einfos.end()) {
198 #ifdef DUPLICATEVERBOSE        
199         cout << "--> Potential duplicate found for " << header->fRun << " " 
200              << header->fTime << " " << header->fPeriod << " " << header->fBx << endl;
201 #endif
202       ++totncan;
203 #ifdef DUPLICATEVERBOSE        
204         map<ULong64_t, MyHeader*>::iterator it = eheaders.find(evtid);
205         if ( (it->second->fRun==header->fRun) && 
206              (it->second->fNTracks==header->fNTracks) &&
207              (it->second->fNSelTracks==header->fNSelTracks) &&
208              (it->second->fNTracklets==header->fNTracklets) &&
209              (it->second->fVz==header->fVz)) {
210           cout << "--> Duplicate found:" << endl;
211           cout << " Run: " << it->second->fRun << " vs " << header->fRun << endl; 
212           cout << " Ntracks: " << it->second->fNTracks << " vs " << header->fNTracks << endl; 
213           cout << " Nseltracks: " << it->second->fNSelTracks << " vs " << header->fNSelTracks << endl; 
214           cout << " Ntracklets: " << it->second->fNTracklets << " vs " << header->fNTracklets << endl; 
215           cout << " Vz: " << it->second->fVz << " vs " << header->fVz << endl;
216           cout << " Duplicate event " << header->fEvNumberInFile << " in file " 
217                << filesnames->At(header->fFileId)->GetName() << endl;
218           ++totndup;
219         }
220 #endif
221         continue;
222       }
223 #ifdef DUPLICATEVERBOSE        
224       MyHeader *cheader = new MyHeader(*header);
225       eheaders.insert( pair<ULong64_t, MyHeader*>(evtid,cheader) );
226 #endif
227 #endif
228       EvInfo *cinfo = new EvInfo;
229       cinfo->fRun = header->fRun;
230       cinfo->fArrId = i;
231       cinfo->fEntry = ev;
232       einfos.insert( pair<ULong64_t, EvInfo*>(evtid,cinfo) );
233       if (lrun != header->fRun) {
234         for (Int_t lx=0;lx<4;++lx) {
235           lrun = header->fRun;
236           TString cname(Form("L%d_run%d_map",lx,lrun));
237           if (lx==3) 
238             cname=Form("TrigClass_run%d_map",lrun);
239           TMap *map = dynamic_cast<TMap*>(output->FindObject(cname));
240           if (map) {
241             maps->Add(map);
242             output->Remove(map);
243           }
244         }
245       }
246     }
247     delete output;
248     file->Close();
249     delete file;
250     if ((nEvents>0) && (einfos.size()>=(UInt_t)nEvents))
251       break;
252   }
253   cout << "Found total: " << totnent << endl;
254   cout << "Found suspicious: " << totnsus << endl;
255 #ifdef DUPLICATECHECK
256   cout << "Found potential duplicates: " << totncan << endl;
257 #ifdef DUPLICATEVERBOSE        
258   cout << "Found duplicates: " << totndup << endl;
259 #endif
260 #endif
261
262   // sort output tree
263   std::vector<TFile*> files(nfiles);
264   std::vector<TTree*> trees(nfiles);
265   std::vector<MyHeader*> headers(nfiles);
266   std::vector<TClonesArray*> partas(nfiles);
267   std::vector<TClonesArray*> trklas(nfiles);
268
269   MyHeader *newHeader = new MyHeader;
270   if (TClass::GetClass("MyPart"))
271     TClass::GetClass("MyPart")->IgnoreTObjectStreamer();
272   TClonesArray *newParts     = new TClonesArray("MyPart",10000);
273   if (TClass::GetClass("MyTracklet"))
274     TClass::GetClass("MyTracklet")->IgnoreTObjectStreamer();
275   TClonesArray *newTracklets = new TClonesArray("MyTracklet",10000);
276
277   TFile *newFile = TFile::Open(outFileName,"recreate","Merged/Sorted MyTree created by mergeCorr.C",5);
278   TDirectory::TContext cd(newFile);
279   for (Int_t m=0;m<maps->GetEntries();++m) {
280     TMap *map = (TMap*)maps->At(m);
281     //map->Print();
282     map->Write(map->GetName(),TObject::kSingleKey);
283   }
284   TTree *newTree = new TTree("MyTree", "MyTree created by MyAliTask.cxx");
285   newTree->Branch("header",&newHeader, 32*1024, 99);
286   newTree->Branch("parts",&newParts, 256*1024, 99);
287   newTree->Branch("tracklets",&newTracklets, 256*1024, 99);
288   newTree->SetDirectory(newFile);
289   Int_t newEntries = 0;
290   for (map<ULong64_t, EvInfo*>::iterator it = einfos.begin();it!=einfos.end();++it) {
291     Short_t id = it->second->fArrId;
292     UInt_t entry = it->second->fEntry;
293     if (!trees.at(id)) { // connect new tree
294       TString fname(filenames->At(id)->GetTitle());
295       TFile *file = TFile::Open(fname);
296       TDirectory::TContext cd2(newFile,file);
297       files[id] = file;
298       TList *output = dynamic_cast<TList*>(file->Get("output"));
299       TTree *tree = dynamic_cast<TTree*>(output->FindObject("MyTree"));
300       trees[id] = tree;
301       output->Remove(tree);
302       output->SetOwner(1);
303       delete output;
304       tree->SetBranchAddress("header",&headers[id]);
305       tree->SetBranchAddress("parts",&partas[id]);
306       tree->SetBranchAddress("tracklets",&trklas[id]);
307     }
308     files[id]->cd();
309     trees.at(id)->GetEntry(entry);
310     *newHeader = *(headers[id]);
311     newParts->Clear();
312     Int_t nparts = partas[id]->GetEntries();
313     for (Int_t p = 0; p < nparts; ++p) {
314       MyPart *orig = dynamic_cast<MyPart*>(partas[id]->At(p));
315       new((*newParts)[p]) MyPart(*orig);
316     }
317     Int_t ntrkls = trklas[id]->GetEntries();
318     for (Int_t p = 0; p < ntrkls; ++p) {
319       MyTracklet *orig = dynamic_cast<MyTracklet*>(trklas[id]->At(p));
320       new((*newTracklets)[p]) MyTracklet(*orig);
321     }
322     newFile->cd();
323     newTree->Fill();
324     ++newEntries;
325     delete it->second;
326     if ((nEvents>0) && (newEntries>=nEvents))
327       break;
328   }
329
330   newTree->Write();
331   newFile->Close();
332   delete newFile;
333
334   for(Int_t i=0;i<nfiles;++i) {
335     delete trees.at(i);
336     delete files.at(i);
337   }
338 }