[PWGCF/FemtoUniverse] Add more flags to efficiency calculator script + fix configurable for timestamps (#9798)

Author: davkk

PWGCF/FemtoUniverse/Core/FemtoUniverseEfficiencyCalculator.h

@@ -40,7 +40,7 @@ struct EfficiencyConfigurableGroup : ConfigurableGroup {
   Configurable<bool> confEfficiencyApplyCorrections{"confEfficiencyApplyCorrections", false, "Should apply corrections from efficiency"};
   Configurable<std::vector<std::string>> confEfficiencyCCDBLabels{"confEfficiencyCCDBLabels", {}, "Custom labels for efficiency objects in CCDB"};
   Configurable<int> confCCDBTrainNumber{"confCCDBTrainNumber", -1, "Train number for which to query CCDB objects (set to -1 to ignore)"};
-  ConfigurableAxis confEfficiencyCCDBTimestamps{"confEfficiencyCCDBTimestamps", {-1, -1}, "Timestamps in CCDB, to query for specific objects (default: -1 for both)"};
+  Configurable<std::vector<std::string>> confEfficiencyCCDBTimestamps{"confEfficiencyCCDBTimestamps", {"-1", "-1"}, "Timestamps in CCDB, to query for specific objects (default: -1 for both, the latest valid object)"};
 
   // NOTE: in the future we might move the below configurables to a separate struct, eg. CCDBConfigurableGroup
   Configurable<std::string> confCCDBUrl{"confCCDBUrl", "http://alice-ccdb.cern.ch", "CCDB URL to be used"};
@@ -119,7 +119,7 @@ class EfficiencyCalculator
     }
 
     auto timestamp = partNo - 1 < config->confEfficiencyCCDBTimestamps->size()
-                       ? static_cast<int64_t>(config->confEfficiencyCCDBTimestamps.value[partNo - 1])
+                       ? std::stoll(config->confEfficiencyCCDBTimestamps.value[partNo - 1])
                        : -1;
 
     auto hEff = ccdb.getSpecific<TH1>(config->confCCDBPath, timestamp, metadata);

PWGCF/FemtoUniverse/Scripts/femto_universe_efficiency_calculator.py

@@ -25,52 +25,82 @@
 
 import ROOT  # pylint: disable=import-error
 
-parser = argparse.ArgumentParser(description="A tool to calculate efficiency and upload it to CCDB")
+
+class CustomHelpFormatter(argparse.HelpFormatter):
+    "Add default value to help format"
+
+    def _get_help_string(self, action):
+        help_str = action.help
+        if help_str is not None and action.default not in [argparse.SUPPRESS, None]:
+            help_str += f" (default: {action.default})"
+        return help_str
+
+
+parser = argparse.ArgumentParser(
+    description="A tool to calculate efficiency and upload it to CCDB",
+    formatter_class=CustomHelpFormatter,
+)
 parser.add_argument(
-    "run_dir",
+    "--alien-path",
     type=Path,
-    help="Path to run directory with analysis results",
+    help="path to train run's directory in Alien with analysis results "
+    "[example: /alice/cern.ch/user/a/alihyperloop/outputs/0033/332611/70301]",
+)
+parser.add_argument(
+    "--mc-reco",
+    type=str,
+    nargs="+",
+    help="paths to MC Reco histograms, separated by space [example: task/mcreco_one/hPt task/mcreco_two/hPt]",
+    required=True,
 )
 parser.add_argument(
-    "--task",
-    "-t",
+    "--mc-truth",
     type=str,
-    help="Task name to query histograms from",
+    nargs="+",
+    help="paths to MC Truth histograms, separated by space [example: task/mctruth_one/hPt task/mctruth_one/hPt]",
     required=True,
 )
 parser.add_argument(
     "--ccdb-path",
-    "-P",
     type=str,
-    help="Path where to save in CCDB",
+    help="location in CCDB to where objects will be uploaded",
     required=True,
 )
 parser.add_argument(
     "--ccdb-url",
-    "-U",
     type=str,
-    help="Path where to save in CCDB",
+    help="URL to CCDB",
     default="http://ccdb-test.cern.ch:8080",
 )
 parser.add_argument(
     "--ccdb-labels",
-    "-L",
     type=str,
     nargs="+",
-    help="Labels for objects' metadata in CCDB",
-    default=[""] * 2,
+    help="custom labels to add to objects' metadata in CCDB [example: label1 label2]",
+    default=[],
 )
 parser.add_argument(
     "--ccdb-lifetime",
-    "-T",
     type=int,
-    help="How long should the objects' validity in CCDB last (default: 1 year)",
+    help="how long should objects in CCDB remain valid (milliseconds)",
     default=365 * 24 * 60 * 60 * 1000,  # one year
 )
 args = parser.parse_args()
 
+if len(args.mc_reco) != len(args.mc_truth):
+    print("[!] Provided number of histograms with MC Reco must match MC Truth", file=sys.stderr)
+    sys.exit(1)
+
+if len(args.ccdb_labels) > 0 and len(args.ccdb_labels) != len(args.mc_reco):
+    print("[!] You must provide labels for all particles", file=sys.stderr)
+    sys.exit(1)
+
+if len(args.ccdb_labels) == 0:
+    # if flag is not provided, fill with empty strings to match size
+    args.ccdb_labels = [""] * len(args.mc_reco)
+
 ANALYSIS_RESULTS = "AnalysisResults.root"
-results_path = args.run_dir / ANALYSIS_RESULTS
+results_path = args.alien_path / ANALYSIS_RESULTS
 job_id = results_path.parent.name
 train_number = results_path.parent.parent.name
 
@@ -81,7 +111,7 @@
 
 # get file from alien
 if not res_dest.is_file():
-    print(f"[↓] Downloading analysis results from Alien to {res_dest} ...", file=sys.stderr)
+    print(f"[↓] Downloading analysis results from Alien to '{res_dest}' ...", file=sys.stderr)
     ROOT.TGrid.Connect("alien://")
     try:
         subprocess.run(
@@ -92,57 +122,60 @@
         print("[-] Download complete!", file=sys.stderr)
     except subprocess.CalledProcessError as error:
         print(f"[!] Error while downloading results file: {error.stderr}", file=sys.stderr)
-        sys.exit(0)
+        sys.exit(1)
 else:
     print(
-        f"[-] Skipping download from Alien, since {res_dest} is already present",
+        f"[-] Skipping download from Alien, since '{res_dest}' is already present",
         file=sys.stderr,
     )
 
-particles = {1: "one", 2: "two"}
+print()
+
 histos_to_upload = []
 
 # get reco & truth histos
 with (
     ROOT.TFile.Open(res_dest.as_uri()) as res_file,
     ROOT.TFile.Open(eff_dest.as_uri(), "recreate") as eff_file,
 ):
-    for idx, part_num in particles.items():
-        reco = res_file.Get(f"{args.task}/Tracks_{part_num}_MC/hPt")
-        truth = res_file.Get(f"{args.task}/MCTruthTracks_{part_num}/hPt")
-        if not reco and not truth:
-            print(
-                f"[-] No MC Reco nor MC Truth histogram found for particle {part_num}",
-                file=sys.stderr,
-            )
-            continue
+    for idx, (mc_reco, mc_truth) in enumerate(zip(args.mc_reco, args.mc_truth)):
+        hist_reco = res_file.Get(mc_reco)
+        if not hist_reco:
+            print(f"[!] Cannot find MC Reco histogram in '{mc_reco}', aborting", file=sys.stderr)
+            sys.exit(1)
 
-        num_bins = reco.GetNbinsX()
-        x_max = reco.GetXaxis().GetBinLowEdge(num_bins) + reco.GetXaxis().GetBinWidth(num_bins)
+        hist_truth = res_file.Get(mc_truth)
+        if not hist_truth:
+            print(f"[!] Cannot find MC Truth histogram in '{mc_truth}', aborting", file=sys.stderr)
+            sys.exit(1)
 
-        hist_name = f"Efficiency_part{idx}"
+        num_bins = hist_reco.GetNbinsX()
+        x_max = hist_reco.GetXaxis().GetBinLowEdge(num_bins)
+        x_max += hist_reco.GetXaxis().GetBinWidth(num_bins)
+
+        hist_name = f"Efficiency_part{idx + 1}"
 
         # calculate efficiency
         eff = ROOT.TH1F(hist_name, "", num_bins, 0, x_max)
         for bin_idx in range(len(eff)):
-            denom = truth.GetBinContent(bin_idx)
-            eff.SetBinContent(bin_idx, reco.GetBinContent(bin_idx) / denom if denom > 0 else 0)
+            denom = hist_truth.GetBinContent(bin_idx)
+            eff.SetBinContent(bin_idx, hist_reco.GetBinContent(bin_idx) / denom if denom > 0 else 0)
 
         # save efficiency object to file
         eff_file.WriteObject(eff, hist_name)
         histos_to_upload.append(hist_name)
 
 if len(histos_to_upload) == 0:
     print("[-] Exiting, since there is nothing to upload", file=sys.stderr)
-    sys.exit(0)
+    sys.exit(1)
 
 # upload objects to ccdb
 try:
     for idx, key in enumerate(histos_to_upload):
         timestamp_start = int(time.time() * 1000)
         timestamp_end = timestamp_start + args.ccdb_lifetime
 
-        print(f"[↑] Uploading {key} to CCDB ... ", file=sys.stderr, end="")
+        print(f"[↑] Uploading {key} to {args.ccdb_url} ... ", file=sys.stderr, end="")
         upload_cmd = [
             "o2-ccdb-upload",
             "--file",
@@ -161,14 +194,23 @@
             f"trainNumber={train_number};label={args.ccdb_labels[idx]}",
         ]
         result = subprocess.run(upload_cmd, capture_output=True, check=True)
+
+        if "html" in result.stdout.decode("utf-8"):
+            print(
+                f"\n[!] Something went wrong with upload request: {result.stdout.decode('utf-8')}",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
         if result.stderr:
-            print(f"[!] Error while uploading: {result.stderr.decode('utf-8')}", file=sys.stderr)
-            sys.exit(0)
+            print(f"\n[!] Error while uploading: {result.stderr.decode('utf-8')}", file=sys.stderr)
+            sys.exit(1)
 
         print("complete!", file=sys.stderr)
 
 except subprocess.CalledProcessError as error:
-    print(f"[!] Error while uploading: {error.stderr.decode('utf-8')}", file=sys.stderr)
-    sys.exit(0)
+    print(f"\n[!] Error while uploading: {error.stderr.decode('utf-8')}", file=sys.stderr)
+    sys.exit(1)
 
+print()
 print("[✓] Success!", file=sys.stderr)