Added ignores
[u/mrichter/AliRoot.git] / PWGLF / FORWARD / trains / TrainSetup.C
1 /**
2  * @defgroup pwglf_forward_trains Trains
3  * 
4  * Train specifications 
5  *
6  * @ingroup pwglf_forward
7  */
8 /**
9  * @file   TrainSetup.C
10  * @author Christian Holm Christensen <cholm@master.hehi.nbi.dk>
11  * @date   Tue Oct 16 17:56:57 2012
12  * 
13  * @brief  Base classs for train specifications 
14  * 
15  * @ingroup pwglf_forward_trains
16  */
17 #ifndef TRAINSETUP_C
18 #define TRAINSETUP_C
19 #ifndef __CINT__
20 # include "Helper.C"
21 # include "Option.C"
22 # include <TDatime.h>
23 # include <TUrl.h>
24 # include <TString.h>
25 # include <TApplication.h>
26 # include <AliAnalysisManager.h>
27 # include <AliVEventHandler.h>
28 # include <AliPhysicsSelection.h>
29 # include <AliPhysicsSelectionTask.h>
30 # include <AliCentralitySelectionTask.h>
31 # include <AliESDInputHandler.h>
32 # include <AliAODInputHandler.h>
33 # include <AliAODHandler.h>
34 # include <AliMCEventHandler.h>
35 #else 
36 struct Helper;
37 struct OptionList;
38 class TDatime;
39 class TUrl;
40 class TString;
41 class AliVEventHandler;
42 class AliAnalysisManager;
43 class AliInputEventHandler;
44 #endif
45
46 //====================================================================
47 /** 
48  * Generic set-up of an analysis train using the grid-handler (AliEn plugin). 
49  *
50  * See also @ref train_setup_doc
51  *
52  * @ingroup pwglf_forward_trains
53  * 
54  */
55 struct TrainSetup
56 {
57   /** 
58    * Constructor 
59    * 
60    * @param name Name of the train
61    */
62   TrainSetup(const TString& name)
63     : fName(name), 
64       fEscapedName(name),
65       fHelper(0)
66   {
67     fOptions.Add("help", "Show help");
68     fOptions.Add("date", "YYYY-MM-DD HH:MM", "Set date", "now");
69     fOptions.Add("mc", "Assume MC input");
70     fOptions.Add("bare-ps", "Use bare physics selection w/o task");
71     fOptions.Add("verbose", "LEVEL", "Set verbosity level", "0");
72     fOptions.Add("url", "URL", "Job location & input URL", "");
73     fOptions.Add("overwrite", "Allow overwrite");
74     fOptions.Add("events", "N", "Number of events to analyse", "-1");
75     fOptions.Add("type", "ESD|AOD|USER", "Input data stype", "");
76     fEscapedName = EscapeName(fName, "");
77   }
78   /** 
79    * Destructor
80    */
81   virtual ~TrainSetup() {}
82   /* @} */
83   //__________________________________________________________________
84   /** 
85    * @{ 
86    * @name Execution 
87    */
88   /** 
89    * Initialize 
90    * 
91    * @return true on success 
92    */
93   Bool_t Init()
94   {
95     // --- Create the helper -----------------------------------------
96     TString  url     = fOptions.Get("url");
97     Int_t    verbose = fOptions.AsInt("verbose");
98     Bool_t   mc      = fOptions.AsBool("mc");
99
100     fHelper = Helper::Create(url.Data(), verbose);
101     if (!fHelper) { 
102       Error("Init", "Failed to make the worker for URL %s", url.Data());
103       return false;
104     }
105
106     UShort_t type    = fHelper->InputType();
107     if (fOptions.Has("type")) { 
108       const TString& it = fOptions.Get("type");
109       if      (it.EqualTo("ESD",TString::kIgnoreCase)) type = Helper::kESD;
110       else if (it.EqualTo("AOD",TString::kIgnoreCase)) type = Helper::kAOD;
111       else if (it.EqualTo("user",TString::kIgnoreCase)) 
112         type = Helper::kUser;
113     }
114
115     // --- Get current directory and set-up sub-directory ------------
116     TString cwd = gSystem->WorkingDirectory();
117     if (!SetupWorkingDirectory()) return false;
118
119     // --- Do initial helper setup -----------------------------------
120     if (!fHelper->PreSetup()) return false;
121
122     // --- Load ROOT libraries ---------------------------------------
123     if (!fHelper->LoadROOT()) return false;
124     
125     // --- Load AliROOT libraries ------------------------------------
126     if (!fHelper->LoadAliROOT()) return false;
127
128     // --- Create analysis manager -----------------------------------
129     AliAnalysisManager *mgr  = CreateAnalysisManager(fName);
130
131     // In test mode, collect system information on every event 
132     // if (oper == kTest)  mgr->SetNSysInfo(1); 
133     if (verbose  >  0)      mgr->SetDebugLevel(verbose);
134     if (fHelper->Mode() == Helper::kLocal) 
135       mgr->SetUseProgressBar(kTRUE, 100);
136    
137     // --- ESD input handler ------------------------------------------
138     AliVEventHandler*  inputHandler = CreateInputHandler(type);
139     if (inputHandler) mgr->SetInputEventHandler(inputHandler);
140     
141     // --- Monte-Carlo ------------------------------------------------
142     AliVEventHandler*  mcHandler = CreateMCHandler(type,mc);
143     if (mcHandler) mgr->SetMCtruthEventHandler(mcHandler);
144     
145     // --- AOD output handler -----------------------------------------
146     AliVEventHandler*  outputHandler = CreateOutputHandler(type);
147     if (outputHandler) mgr->SetOutputEventHandler(outputHandler);
148
149     // --- Include analysis macro path in search path ----------------
150     gROOT->SetMacroPath(Form("%s:%s:$ALICE_ROOT/ANALYSIS/macros",
151                              cwd.Data(), gROOT->GetMacroPath()));
152
153     // --- Physics selction - only for ESD ---------------------------
154     if (type == Helper::kESD) CreatePhysicsSelection(mc, mgr);
155     
156     // --- Create centrality task ------------------------------------
157     CreateCentralitySelection(mc, mgr);
158
159     // --- Create tasks ----------------------------------------------
160     CreateTasks(mgr);
161
162     // --- Post set-up initialization of helper ----------------------
163     if (!fHelper->PostSetup()) return false;
164
165     // --- Set debug level on defined tasks --------------------------
166     if (verbose > 0) {
167       TIter next(mgr->GetTasks());
168       AliAnalysisTask* sub = 0;
169       while ((sub = static_cast<AliAnalysisTask*>(next()))) { 
170         AliAnalysisTaskSE* se = dynamic_cast<AliAnalysisTaskSE*>(sub);
171         if (!se) continue;
172         se->SetDebugLevel(verbose);
173       }
174     }
175
176     // --- Print this setup ------------------------------------------
177     Print();
178
179     // --- Initialise the train --------------------------------------
180     if (!mgr->InitAnalysis())  {
181       gSystem->ChangeDirectory(cwd.Data());
182       Error("Init","Failed to initialise train");
183       return false;
184     }
185
186     // --- Enable progress bar ---------------------------------------
187     if (fHelper->Mode() != Helper::kGrid) 
188       mgr->SetUseProgressBar(true, 100);
189
190     // --- Save setup to disk ----------------------------------------
191     SaveSetup(true);
192
193     // Some information
194     mgr->PrintStatus();
195
196     return true;
197   }
198   Bool_t Run(Bool_t doExit=false)
199   {
200     TString cwd = gSystem->WorkingDirectory();
201     Bool_t status = false;
202     try {
203       if (!Init()) throw TString("Failed to intialize the train");
204
205       // if (r) SaveSetup(*r, nEvents, asShell);
206       
207       Long64_t nEvents = fOptions.AsLong("events", -1);
208       Long64_t ret     = fHelper->Run(nEvents);
209       
210       // Make sure we go back 
211       gSystem->ChangeDirectory(cwd.Data());
212       
213       // Return. 
214       if (ret < 0) throw TString("Analysis failed");
215
216       status = true;
217     }
218     catch (TString& e) {
219       Error("Main", e);
220       status = false;
221     }
222     if (gApplication && doExit) {
223       gSystem->Sleep(3);
224       gApplication->Terminate(status ? 0 : 1);
225     }
226     return true;
227   }
228   /** 
229    * Get the options 
230    * 
231    * 
232    * @return Reference ot the options 
233    */
234   OptionList& Options() { return fOptions; }
235   /** 
236    * Print information to standard output 
237    * 
238    */
239   void Print(Option_t* ="") const
240   {
241     std::cout << "Train: " << fName << " (" << fEscapedName << ")" 
242               << std::endl;
243     fOptions.Show(std::cout);
244     if (fHelper) fHelper->Print();
245   }
246   /** 
247    * Show the help 
248    * 
249    * @param o 
250    * 
251    * @return 
252    */
253   Bool_t Help(std::ostream& o=std::cout, bool asProg=false)
254   {
255     if (!fOptions.Has("help")) return true;
256
257     if (asProg) 
258       o << "Usage: runTrain --name=NAME --class=CLASS [OPTIONS]";
259     else 
260       o << "Usage: RunTrain(NAME, CLASS, OPTIONS)";
261     
262     o << "\n\nOptions:\n\n";
263     if (asProg) {
264       OptionList tmp(fOptions);
265       tmp.Add("name", "STRING", "Name of train", fName);
266       tmp.Add("class", "NAME", "Name of setup class", "");
267       tmp.Help(o,"  --");
268     }
269     else 
270       fOptions.Help(o, "  ");
271     
272     o << "\n";
273
274     if (!fHelper && fOptions.Has("url")) {
275       TString url = fOptions.Get("url");
276       fHelper = Helper::Create(url.Data());
277     }
278     if (fHelper) { 
279       o << fHelper->Desc() << " URL form:\n\n"
280         << "    " << fHelper->UrlHelp() << "\n\n"
281         << "Options:\n";
282       fHelper->Options().Help(o, "    ");
283       o << "\n";
284     }
285     else { 
286       o << "Possible URL forms:\n\n";
287       Helper::ShowUrlHelp("LocalHelper");
288       Helper::ShowUrlHelp("ProofHelper");
289       Helper::ShowUrlHelp("LiteHelper");
290       Helper::ShowUrlHelp("AAFHelper");
291       Helper::ShowUrlHelp("AAFPluginHelper");
292       Helper::ShowUrlHelp("GridHelper");
293       o << "\n";
294     }
295     return false;
296   }
297   /** 
298    * Run train.  This will AcLic compile the setup script, create
299    * an object of that type with the given name, and then pass the 
300    * options to it.  Then, it will run the setup.
301    * 
302    * @param name  Train name 
303    * @param cls   Class name 
304    * @param opts  Comma seperated list of options
305    * 
306    * @return true on success
307    */
308   static Bool_t Main(const TString& name, const TString& cls, 
309                      const TCollection* opts, 
310                      Bool_t asProg=true)
311   {
312     if (cls.IsNull()) { 
313       Error("Main", "No class name specified");
314       return false;
315     }
316     if (name.IsNull()) { 
317       Error("Main", "No train name specified");
318       return false;
319     }
320     Int_t error = 0;
321     Int_t r1 = gROOT->LoadMacro(Form("%s.C++g", cls.Data()), &error);
322     if (r1 < 0 || error) { 
323       Error("Main", "Failed to load setup %s: %d", cls.Data(), error);
324       return false;
325     }
326
327     // Make our object using the interpreter 
328     TString create = TString::Format("new %s(\"%s\")", 
329                                    cls.Data(), name.Data());
330     gROOT->ProcessLine("gSystem->RedirectOutput(\"/dev/null\",\"w\");");
331     Long_t ret = gROOT->ProcessLine(create, &error);
332     gROOT->ProcessLine("gSystem->RedirectOutput(0);");
333     if (!ret || error) { 
334       Error("Main", "Failed to make object of class %s: 0x%08lx/%d\n\t%s", 
335             cls.Data(), ret, error, create.Data());
336       return false;
337     }
338     TrainSetup* train = reinterpret_cast<TrainSetup*>(ret);
339     
340     // Now parse the options 
341     if (!train->Options().Parse(opts)) { 
342       Error("Main", "Failed to parse options");
343       return false;
344     }
345
346     // Check if we got a help request
347     if (train->Options().Has("help")) { 
348       train->Help(std::cout, asProg);
349       return true;
350     }
351     // return train->Init();
352     return train->Run(asProg);
353   }
354 protected:
355   //__________________________________________________________________
356   /** 
357    * @{ 
358    * @Name Overloadable behaviour 
359    */
360   //------------------------------------------------------------------
361   /** 
362    * Create the analysis manager 
363    * 
364    * @param name Name of the analysis 
365    * 
366    * @return Created analysis manager 
367    */
368   virtual AliAnalysisManager* CreateAnalysisManager(const char* name)
369   {
370     return new AliAnalysisManager(name,"Analysis Train");
371   }
372   //------------------------------------------------------------------
373   /** 
374    * Create input handler 
375    * 
376    * @param type 
377    * 
378    * @return 
379    */
380   virtual AliVEventHandler* CreateInputHandler(UShort_t type)
381   {
382     switch (type) {
383     case Helper::kESD:  return new AliESDInputHandler(); 
384     case Helper::kAOD:  return new AliAODInputHandler(); 
385     case Helper::kUser: return 0;
386     }
387     return 0;
388   }
389   //------------------------------------------------------------------
390   /** 
391    * Create MC input handler 
392    * 
393    * @param type  Run type (ESD or AOD)
394    * @param mc    Assume monte-carlo input 
395    * 
396    * @return 
397    */
398   virtual AliVEventHandler* CreateMCHandler(UShort_t /*type*/, bool mc)
399   {
400     if (!mc) return 0;
401     AliMCEventHandler* mcHandler = new AliMCEventHandler();
402     mcHandler->SetReadTR(true); 
403     return mcHandler;
404   }
405   //------------------------------------------------------------------
406   /** 
407    * Create output event handler 
408    * 
409    * @param type 
410    * 
411    * @return 
412    */
413   virtual AliVEventHandler* CreateOutputHandler(UShort_t type)
414   {
415     AliAODHandler* ret = new AliAODHandler();
416     switch (type) { 
417     case Helper::kESD: 
418       ret->SetOutputFileName("AliAOD.root");
419       break;
420     case Helper::kAOD: 
421       ret->SetOutputFileName("AliAOD.pass2.root");
422       break;
423     case Helper::kUser: 
424       break;
425     }
426     
427     return ret;
428   }
429   //------------------------------------------------------------------
430   /** 
431    * Create physics selection, and add to manager
432    * 
433    * @param mc Whether this is for MC 
434    * @param mgr Manager
435    */
436   virtual void CreatePhysicsSelection(Bool_t mc, AliAnalysisManager* mgr)
437   {
438     if (fOptions.Has("bare-ps")) {
439       AliInputEventHandler* input = 
440         dynamic_cast<AliInputEventHandler*> (mgr->GetInputEventHandler());
441       if (!input) return;
442
443       AliPhysicsSelection* ps = new AliPhysicsSelection();
444       if (mc) ps->SetAnalyzeMC();
445
446       input->SetEventSelection(ps);
447
448       return;
449     }
450     gROOT->Macro(Form("AddTaskPhysicsSelection.C(%d)", mc));
451     mgr->RegisterExtraFile("event_stat.root");
452   }
453   //------------------------------------------------------------------
454   /** 
455    * Create centrality selection, and add to manager
456    * 
457    * @param mc Whether this is for MC 
458    * @param mgr Manager
459    */
460   virtual void CreateCentralitySelection(Bool_t mc, AliAnalysisManager* mgr)
461   {
462     gROOT->Macro("AddTaskCentrality.C");
463     const char* name = "CentralitySelection";
464     AliCentralitySelectionTask* ctask = 
465       dynamic_cast<AliCentralitySelectionTask*>(mgr->GetTask(name));
466     if (!ctask) return;
467     if (mc) ctask->SetMCInput();
468   }
469   //------------------------------------------------------------------
470   /** 
471    * Create analysis tasks.  Must be overloaded by sub-class
472    * 
473    * @param mgr  Manager
474    */
475   virtual void CreateTasks(AliAnalysisManager* mgr)=0;
476   virtual const Char_t* ClassName() const = 0;
477   /* @} */
478   //__________________________________________________________________
479   /** 
480    * @{ 
481    * @name Utility functions 
482    */
483   /** 
484    * Escape bad elements of the name 
485    * 
486    * @param name   Name to escape 
487    * @param datime Date and Time 
488    * 
489    * @return escaped name 
490    */  
491   static TString EscapeName(const char* name, const TString& datimeStr)
492   {
493     TString escaped = name;
494     char  c[] = { ' ', '/', '@', 0 };
495     char* p   = c;
496     while (*p) { 
497       escaped.ReplaceAll(Form("%c", *p), "_");
498       p++;
499     }
500     if (!datimeStr.IsNull()) {
501       TDatime datime;
502       if (datimeStr.EqualTo("now", TString::kIgnoreCase)) 
503         datime.Set();
504       else 
505         datime.Set(datimeStr.Data());
506       if (datime.GetYear() <= 1995 ||
507           datime.GetMonth() == 0 || 
508           datime.GetDay() == 0) return escaped;
509       escaped.Append(Form("_%04d%02d%02d_%02d%02d", 
510                           datime.GetYear(), 
511                           datime.GetMonth(), 
512                           datime.GetDay(), 
513                           datime.GetHour(), 
514                           datime.GetMinute()));
515     }
516     return escaped;
517   }    
518   /** 
519    * Make our working directory if so requested 
520    * 
521    * @return true on success
522    */
523   Bool_t SetupWorkingDirectory()
524   {
525     // Get the name of the target directory 
526     TString& nam = fEscapedName;
527
528     // Check if the directory exists already 
529     Bool_t exists = gSystem->AccessPathName(nam.Data()) == 0;
530     if (fHelper->Operation() == Helper::kTerminate && !exists) {
531       Error("SetupWorkingDirectory", "File/directory %s does not exists", 
532             nam.Data());
533       return false;
534     }
535
536     Bool_t overwrite = fOptions.Has("overwrite");
537     // If we're not allowed to overwrite, then complain
538     if (!overwrite && exists) {
539       Error("SetupWorkingDirectory", "File/directory %s already exists", 
540             nam.Data());
541       return false;
542     }
543
544     // Make the target directory if it doesn't exists 
545     if (!exists) {
546       if (gSystem->MakeDirectory(nam.Data())) {
547         Error("SetupWorkingDirectory", "Failed to make directory '%s'", 
548               nam.Data());
549         return false;
550       }
551     }
552       
553     // Change directory to target directory 
554     if (!gSystem->ChangeDirectory(nam.Data())) { 
555       Error("SetupWorkingDirectory", "Failed to change directory to %s", 
556             nam.Data());
557       return false;
558     }
559     Info("SetupWorkingDirectory", "Made subdirectory %s, and cd'ed there", 
560          nam.Data());
561     return true;
562   }
563   /** 
564    * Save the setup as a ROOT script and possibly also a shell script
565    * 
566    * @param asShellScript If true, also save as shell script
567    */
568   virtual void SaveSetup(Bool_t asShellScript)
569   {
570     if (asShellScript) 
571       SaveSetupShell("rerun", ClassName(), fName, fOptions);
572     SaveSetupROOT("ReRun", ClassName(), fName, fOptions);
573   }
574   /** 
575    * Save a setup as a shell script 
576    * 
577    * @param out   Output name of shell script 
578    * @param cls   Class of the train 
579    * @param name  Name of the train
580    * @param opts  Option list
581    */
582   static void SaveSetupShell(const TString& out, const TString& cls,
583                               const TString& name, const OptionList& opts)
584   {
585     std::ofstream o(Form("%s.sh", out.Data()));
586     o << "#!/bin/bash\n\n"
587       << "class=\"" << cls << "\"\n"
588       << "name=\"" << name << "\"\n\n"
589       << "# Available options\n"
590       << "# \n";
591     opts.Help(o, "#    --");
592     o << "#\n"
593       << "opts=(--class=$class \\\n"
594       << "  --name=$name";
595     opts.Store(o, " \\\n  --", "", true);
596     o << ")\n\n"
597       << "echo \"Running runTrain ${opts[@]} $@\"\n"
598       << "runTrain \"${opts[@]}\" $@\n\n"
599       << "# EOF" << std::endl;
600     o.close();
601     gSystem->Exec(Form("chmod a+x %s.sh", out.Data()));
602   }
603   /** 
604    * Save a setup as a ROOT script 
605    * 
606    * @param out   Output name of shell script 
607    * @param cls   Class of the train 
608    * @param name  Name of the train
609    * @param opts  Option list
610    */
611   static void SaveSetupROOT(const TString& out, const TString& cls,
612                             const TString& name, const OptionList& opts)
613   {
614     OptionList tmp(opts);
615     tmp.Remove("url");
616     std::ofstream o(Form("%s.C", out.Data()));
617     o << "/* Available options:\n"
618       << " *\n";
619     tmp.Help(o, " *    ");
620     o << " */\n"
621       << "Bool_t " << out << "()\n"
622       << "{\n"
623       << "  TString name(\"" << name << "\");\n"
624       << "  TString cls(\"" << cls << "\");\n"
625       << "  TString uri(\"" << opts.Get("url") << "\");\n"
626       << "  TString opts(";
627     tmp.Store(o, "\"", ",\"\n               ", false);
628     o << ");\n\n"
629       << "  gROOT->LoadMacro(\"RunTrain.C\");\n\n"
630       << "  return RunTrain(name, cls, uri, opts);\n"
631       << "}\n" 
632       << "//\n"
633       << "// EOF\n" 
634       << "//" << std::endl;
635     o.close();
636   }    
637   /* @} */
638   TString      fName;
639   TString      fEscapedName;
640   OptionList   fOptions;
641   Helper*      fHelper;
642 };
643 #endif