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