updated
[u/mrichter/AliRoot.git] / PWGHF / hfe / macros / test / runGridpPb.C
1 #include <fstream>
2 #include <iostream>
3 #include <string>
4
5 TString g_aliroot_version;
6 TString g_root_version;
7 TString g_sample;
8 TString g_plugin_mode;
9 TString g_train_dir;
10 TArrayI g_runlist;
11
12 AliAnalysisAlien *CreateGridHandler(){
13         //
14         // Setup main settings of the Alien plugin
15         //
16         AliAnalysisAlien *plugin = new AliAnalysisAlien();
17         plugin->SetRunMode(g_plugin_mode.Data());
18         if(!g_plugin_mode.CompareTo("Terminate"))
19                 plugin->SetMergeViaJDL(kFALSE);
20         else
21                 plugin->SetMergeViaJDL(kTRUE);
22         plugin->SetOverwriteMode();
23         plugin->SetNtestFiles(1);
24
25         plugin->SetAPIVersion("V1.1x");
26         plugin->SetROOTVersion(g_root_version.Data());
27         plugin->SetAliROOTVersion(g_aliroot_version.Data());
28
29         plugin->SetOutputToRunNo();
30         plugin->AddIncludePath("-I. .I$ALIEN_ROOT/api/lib -I$ROOTSYS/lib -I$ROOTSYS/include -I$ALICE_ROOT/include -I$ALICE_ROOT/PWGHF/ -I$ALICE_ROOT/PWGHF/hfe/macros -I$ALICE_ROOT/PWGHF/hfe -I$ALICE_ROOT/ANALYSIS -I$ALICE_ROOT/ANALYSIS/Tender -I$ALICE_ROOT/ANALYSIS/TenderSupplies -I$ALICE_ROOT/PWG/ -I$ALICE_ROOT/PWG/FLOW -I$ALICE_ROOT/PWG/Base -I$ALICE_ROOT/PWG/Tasks");
31         plugin->SetAdditionalLibs("libGui.so  libXMLParser.so libSTEERBase.so libESD.so libAOD.so libCDB.so libANALYSIS.so libANALYSISalice.so libCORRFW.so  libPWGflowBase.so libPWGflowTasks.so libPWGHFhfe.so libTENDER.so libProof.so libRAWDatabase.so libSTEER.so libTOFbase.so");
32    
33         plugin->SetDefaultOutputs(kFALSE);
34         plugin->SetOutputFiles("AnalysisResults.root"); 
35         plugin->SetExecutableCommand("aliroot -b -q");
36         plugin->SetTTL(30000);
37         plugin->SetInputFormat("xml-single");
38         plugin->SetPrice(1);      
39         plugin->SetSplitMode("se");
40         return plugin;
41 }
42
43 void SplitConfigEntry(const std::string &input, TString &key, TString &value){
44         //
45         // Decode key and value of a config entry
46         //
47         std::istringstream stream(input, istringstream::in);
48         std::string tmp;
49         stream >> tmp;
50         key = tmp.c_str();
51         stream >> tmp;
52         value = tmp.c_str();
53 }
54
55 void DecodeRunlist(const TString &val){
56         //
57         // Tokenize run list
58         //
59         TObjArray *runstrings = val.Tokenize(",");
60         TObjString *os;
61         TString runstr;
62         TIter runIter(runstrings);
63         g_runlist.Set(runstrings->GetEntries());
64         int nruns(0);
65         while((os = dynamic_cast<TObjString *>(runIter()))){
66                 runstr = os->String();
67                 g_runlist[nruns++] = runstr.Atoi();
68         }
69         delete runstrings;
70 }
71
72 bool IsMC(const TString &val){
73         // 
74         // Determine whether sample is MC or Data
75         //
76         if(!val.CompareTo("MC")) return true;
77         return false;
78 }
79
80 bool FindDataSample(const TMap &lookup, TObjArray &sampleinfis){
81         //
82         // Find Data sample in the list of samples
83         //
84         TObjArray *entry = dynamic_cast<TObjArray *>(lookup.GetValue(g_sample.Data()));
85         if(!entry){
86                 printf("Sample %s not found in the list of samples", g_sample.Data());
87                 return false;
88         }
89         // Copy to output container
90         sampleinfis.SetOwner(kFALSE);
91         for(int ival = 0; ival < 4; ival++) sampleinfis.AddAt(entry->At(ival), ival);
92         return true;
93 }
94
95 bool GetData(TObjArray &in, TString &out, int pos){
96         //
97         // Helper function reading data string
98         //
99         TObjString *entry = dynamic_cast<TObjString *>(in.At(pos));
100         if(!entry){
101                 printf("Entry at pos %d not a string\n", pos);
102                 return false;
103         }
104         out = entry->String();
105         return true;
106 }
107
108 void AddSample(TMap &lookup,
109                 const char *key, const char* datadir, const char * pattern, const char *sampletype, const char *dataformat){
110         //
111         // Add sample entry to the lookup table
112         //
113         TObjArray *infos = new TObjArray(); 
114         infos->AddAt(new TObjString(datadir), 0);
115         infos->AddAt(new TObjString(pattern), 1);
116         infos->AddAt(new TObjString(sampletype), 2);
117         infos->AddAt(new TObjString(dataformat), 3);
118         lookup.Add(new TObjString(key), infos);
119 }
120
121 void Generate_Sample_Lookup(TMap &lookup){
122         // 
123         // Create Lookup table for each period
124         // Vector contains 
125         //   - path
126         //   - pattern
127         //   - MC/Data 
128         //   - ESD/AOD
129         //
130         AddSample(lookup, "LHC13b.pass2", "/alice/data/2013/LHC13b", "pass2/*/AliESDs.root", "Data", "ESD");
131         AddSample(lookup, "LHC13b.pass3", "/alice/data/2013/LHC13b", "pass3/*/AliESDs.root", "Data", "ESD");
132         AddSample(lookup, "LHC13b.pass2.AOD", "/alice/data/2013/LHC13b", "pass2/AOD/*/AliAOD.root", "Data", "AOD"); 
133         AddSample(lookup, "LHC13b.pass2.AOD126", "/alice/data/2013/LHC13b", "*/pass2/AOD126/*/AliAOD.root", "Data", "AOD"); 
134         AddSample(lookup, "LHC13c.pass1", "/alice/data/2013/LHC13c/", "pass1/*/AliESDs.root", "Data", "ESD");
135         AddSample(lookup, "LHC13c.pass2", "/alice/data/2013/LHC13c/", "pass2/*/AliESDs.root", "Data", "ESD");
136         AddSample(lookup, "LHC13c.pass1.AOD", "/alice/data/2013/LHC13c/", "*/pass1/AOD/*/AliAOD.root", "Data", "AOD");
137         AddSample(lookup, "LHC13c.pass1.AOD126", "/alice/data/2013/LHC13c/", "*/pass1/AOD126/*/AliAOD.root", "Data", "AOD");
138         AddSample(lookup, "LHC13b2", "/alice/sim/2013/LHC13b2", "*/*/AliESDs.root", "MC", "ESD");
139         AddSample(lookup, "LHC13b2.AOD", "/alice/sim/2013/LHC13b2", "*/AliAOD.root", "MC", "AOD");
140         AddSample(lookup, "LHC13b2plus", "/alice/sim/2013/LHC13b2_plus", "*/*/AliESDs.root", "MC", "ESD");
141         AddSample(lookup, "LHC13b2plus.AOD", "/alice/sim/2013/LHC13b2_plus", "*/AliAOD.root", "MC", "AOD");
142         AddSample(lookup, "LHC13b3", "/alice/sim/2013/LHC13b3", "*/*/AliESDs.root", "MC", "ESD");
143               AddSample(lookup, "LHC13b3.AOD", "/alice/sim/2013/LHC13b3", "*/AliAOD.root", "MC", "AOD");
144               AddSample(lookup, "LHC13d.pass1", "/alice/data/2013/LHC13d/", "pass1/*/AliESDs.root", "Data", "ESD");
145               AddSample(lookup, "LHC13e.pass1", "/alice/data/2013/LHC13e/", "pass1/*/AliESDs.root", "Data", "ESD");
146               AddSample(lookup, "LHC13f.pass1", "/alice/data/2013/LHC13f/", "pass1/*/AliESDs.root", "Data", "ESD");
147         printf("Lookup table with sample information generated\n");
148 }
149
150 void ConfigParser(const char *configname){
151         //
152         // Parse configuration
153         //
154         std::ifstream in(configname);
155         std::string config;
156         TString key, value;
157         while(getline(in, config)){
158                 SplitConfigEntry(config, key, value);
159                 key.ToLower();
160                 if(!key.CompareTo("aliroot")){
161                         // Aliroot version
162                         g_aliroot_version = value;
163                         continue;
164                 }
165                 if(!key.CompareTo("root")){
166                         // root version
167                         g_root_version = value;
168                         continue;
169                 }
170                 if(!key.CompareTo("sample")){
171                         // sample name
172                         g_sample = value; 
173                         continue;
174                 }
175                 if(!key.CompareTo("runlist")){
176                         // Runlist
177                         DecodeRunlist(value); 
178                         continue;
179                 }
180                 if(!key.CompareTo("mode")){
181                         g_plugin_mode = value;
182                         continue;
183                 }
184                 if(!key.CompareTo("traindir")){
185                         g_train_dir = value;
186                         continue;
187                 }
188                 printf("Unknown key: %s\n", key.Data());
189         }
190 }
191
192 bool MakeSample(AliAnalysisAlien *plugin, TMap &lookup){
193         //
194         // Fill Sample information (Data dir, pattern, run list) to the Alien plugin
195         //
196         TObjArray infos;
197         bool found = FindDataSample(lookup, infos);
198         if(!found){
199                 printf("sample %s not found\n", g_sample.Data());
200                 return false;
201         }
202         TString datadir, pattern, type;
203         GetData(infos, datadir, 0);
204         GetData(infos, pattern, 1);
205         GetData(infos, type, 2);
206         plugin->SetGridDataDir(datadir.Data());
207         plugin->SetDataPattern(pattern.Data());
208         if(!IsMC(type)) plugin->SetRunPrefix("000");
209         // Add runs to the sample
210         for(int irun = 0; irun < g_runlist.GetSize(); irun++){
211                 plugin->AddRunNumber(g_runlist[irun]);
212         }
213         return true;
214 }
215
216 bool CreateTrainDir(AliAnalysisAlien *plugin, const TMap &lookup){
217         //
218         // Make train data dir name and JDL, C and sh file names
219         //
220         TObjArray infos;
221         bool found = FindDataSample(lookup, infos);
222         if(!found){
223                 printf("sample %s not found\n", g_sample.Data());
224                 return false;
225         }
226         TString type; GetData(infos, type, 2);
227
228         // check whether the train dir is already provided or needs to be specified
229         if(!g_train_dir.Length()){
230                 // Query number of train runs before
231                 const char *gridhome = gGrid->GetHomeDirectory();
232                 const char gridoutdir[256];
233                 sprintf(gridoutdir, "%sAnalysis_pPb/%s", gridhome, type.Data());
234                 TGridResult *trainruns = gGrid->Ls(gridoutdir);
235                 int nruns = trainruns->GetEntries();
236                 // Get Date and time
237                 TDatime time;
238                 g_train_dir = Form("%d_%d%02d%02d_%02d%02d", nruns, time.GetYear(), time.GetMonth(), time.GetDay(), time.GetHour(), time.GetMinute());
239         }
240         
241         plugin->SetGridWorkingDir(Form("Analysis_pPb/%s/%s", type.Data(), g_train_dir.Data()));
242         plugin->SetJDLName(Form("TPCTOFanalysispPb_%s_%s.jdl", type.Data(), g_train_dir.Data()));
243         plugin->SetExecutable(Form("TPCTOFanalysispPb_%s_%s.sh", type.Data(), g_train_dir.Data()));
244         plugin->SetAnalysisMacro(Form("TPCTOFanalysispPb_%s_%s.C", type.Data(), g_train_dir.Data()));
245         return true;
246 }
247
248 void SetupUtil(bool isMC, bool isAOD){
249         //
250         // Setup utility packages
251         //
252         // 1. Physics Selection (ESD only)
253         // 2. Tender (ESD only)
254         // 3. PID Response (always)
255         // 4. Centrality Task (ESD only)
256         //
257
258         //===== ADD PID RESPONSE: ===
259         gROOT->LoadMacro("$ALICE_ROOT/ANALYSIS/macros/AddTaskPIDResponse.C");
260         AddTaskPIDResponse(isMC);
261
262         //==== Add tender ====
263         if(!isAOD){
264                 gROOT->LoadMacro("$ALICE_ROOT/ANALYSIS/TenderSupplies/AddTaskTender.C");
265                 //AddTaskTender();
266         }
267
268         //==== Physics Selection ====
269         if(!isAOD){
270                 gROOT->LoadMacro("$ALICE_ROOT/ANALYSIS/macros/AddTaskPhysicsSelection.C");
271                 AliPhysicsSelectionTask* physSelTask = AddTaskPhysicsSelection(isMC);
272         }
273
274         //===== ADD CENTRALITY: ===
275         if(!isAOD){
276                 gROOT->LoadMacro("$ALICE_ROOT/ANALYSIS/macros/AddTaskCentrality.C");
277                 AddTaskCentrality();
278         }
279 }
280
281 void SetupHandlers(bool isMC, bool isAOD){
282         //
283         // Setup Handlers
284         //
285         TString macrobase = "$ALICE_ROOT/ANALYSIS/macros/train/";
286         TString macroname = macrobase;
287         if(isAOD)
288                 macroname += "AddAODHandler.C";
289         else
290                 macroname += "AddESDHandler.C";
291         gROOT->Macro(macroname.Data());
292
293         if(isMC && !isAOD){
294                 // Add MC truth event handler, only in case of ESDs
295                 gROOT->LoadMacro(Form("%s/AddMCHandler.C", macrobase.Data()));
296                 AddMCHandler();
297         }
298 }
299
300 void SetupHFEtask(bool isMC, bool isAOD){
301         gROOT->LoadMacro("$ALICE_ROOT/PWGHF/hfe/macros/AddTaskHFEpPb.C");
302         AddTaskHFEpPb(isMC, isAOD);
303         if(!isAOD){
304                 gROOT->LoadMacro("$ALICE_ROOT/PWGHF/hfe/macros/AddTaskHFEnpepPb.C");
305                 AddTaskHFEnpepPb(isMC, isAOD);
306         }
307 }
308
309 void SetupTrain(const TMap &lookup){
310         //
311         // Setup train:
312         //   Determine whether the trains run on MC or Data 
313         //   and ESDs or AODs and Configure Handlers, utils
314         //   and HFE task according to this
315         //
316         bool isMC(false), isAOD(false);
317         TObjArray infos;
318         bool found = FindDataSample(lookup, infos);
319         if(!found) return;
320         TString type, mode;
321         GetData(infos, type, 2);
322         GetData(infos, mode, 3);
323         isMC = IsMC(type);
324         if(!mode.CompareTo("AOD")) isAOD = true;
325       
326         SetupHandlers(isMC, isAOD);
327         SetupUtil(isMC, isAOD);
328         SetupHFEtask(isMC, isAOD);
329 }
330
331 void GenerateMergeConfigs(){
332         //
333         // Generate configurations for merging 
334         // (MergeViaJDL and Terminate)
335         //
336
337         // Write config for MergeViaJDL
338         std::ofstream outMerge("configMerge.txt");
339         outMerge << "aliroot " << g_aliroot_version.Data() << std::endl;
340         outMerge << "root " << g_root_version.Data() << std::endl;
341         outMerge << "sample " << g_sample.Data() << std::endl;
342         outMerge << "mode MergeViaJDL\n";
343         outMerge << "traindir " << g_train_dir.Data() << std::endl; 
344         outMerge << "runlist ";
345         for(int irun = 0; irun < g_runlist.GetSize()-1; irun++) outMerge << g_runlist[irun] << ",";
346         outMerge << g_runlist[g_runlist.GetSize()-1] << std::endl;
347         outMerge.close();
348         // Write config for Terminate
349         std::ofstream outTerminate("configTerminate.txt");
350         outTerminate << "aliroot " << g_aliroot_version.Data() << std::endl;
351         outTerminate << "root " << g_root_version.Data() << std::endl;
352         outTerminate << "sample " << g_sample.Data() << std::endl;
353         outTerminate << "mode Terminate\n";
354         outTerminate << "traindir " << g_train_dir.Data() << std::endl; 
355         outTerminate << "runlist ";
356         for(int irun = 0; irun < g_runlist.GetSize()-1; irun++) outTerminate << g_runlist[irun] << ",";
357         outTerminate << g_runlist[g_runlist.GetSize()-1] << std::endl;
358         outTerminate.close();
359
360         printf("Configurations for MergeViaJDL and terminate generated\n");
361 }
362
363 void runGridpPb(const char *config = "config.txt"){
364         //
365         // run analysis 
366         //
367
368         TGrid::Connect("alien://");
369
370         // Create Lookup with sample information
371         TMap sampleinfos;
372         Generate_Sample_Lookup(sampleinfos);
373
374         ConfigParser(config);
375
376         // Configure alien plugin
377         AliAnalysisAlien *plugin = CreateGridHandler();
378         if(!CreateTrainDir(plugin, sampleinfos)){
379                 printf("Cannot setup output directory\n");
380                 return;
381         }
382         if(!MakeSample(plugin, sampleinfos)){
383                 printf("Cannot create data sample\n");
384                 return;
385         }
386         if(!g_plugin_mode.CompareTo("full")){
387                 // full mode, creating config files for the merging stage
388                 GenerateMergeConfigs();
389         }
390
391         AliAnalysisManager *mgr = new AliAnalysisManager("tpctofanalysis");
392         mgr->SetGridHandler(plugin);
393         
394         SetupTrain(sampleinfos);
395
396         // Run train
397         if (!mgr->InitAnalysis()) return;
398         mgr->PrintStatus();
399         // Start analysis in grid.
400         mgr->StartAnalysis("grid");
401
402