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