Add more tools to check AODs and HEPMCs (#77)

- Add minor fixes to analysis utilities

Author: njacazio

examples/scripts/createO2tables.py

@@ -28,7 +28,7 @@ def run_cmd(cmd, comment="", check_status=True):
             for i in content.strip().split("\n"):
                 verbose_msg("++", i)
         if "Encountered error" in content:
-            warning_msg("Error encountered runtime in", cmd)
+            warning_msg("Error encountered runtime error in", cmd)
         if check_status:
             if "OK" not in content and "root" not in cmd:
                 fatal_msg("Command", cmd, "does not have the OK tag", content)
@@ -293,9 +293,7 @@ def copy_and_link(file_name):
             hepmc_file = None
             mc_seed = random.randint(1, 800000000)
             if custom_gen:  # Using HEPMC
-                gen_log_file = f"gen.{run_number}.log"
                 hepmc_file = f"hepmcfile.{run_number}.hepmc"
-                custom_gen_option = f" --output {hepmc_file} --nevents {nevents} --seed {mc_seed}"
                 if "INPUT_FILES" in custom_gen:
                     input_hepmc_file = custom_gen.replace("INPUT_FILES",
                                                           "").strip().split(" ")
@@ -306,6 +304,8 @@ def copy_and_link(file_name):
                     write_to_runner(f"ln -s {input_hepmc_file}"
                                     f" {hepmc_file} \n")
                 else:
+                    gen_log_file = f"gen.{run_number}.log"
+                    custom_gen_option = f" --output {hepmc_file} --nevents {nevents} --seed {mc_seed}"
                     write_to_runner(custom_gen + custom_gen_option,
                                     log_file=gen_log_file, check_status=True)
                 write_to_runner(f"DelphesHepMC propagate.tcl {delphes_file} {hepmc_file}",

examples/scripts/diagnostic_tools/check_table_indices.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+
+from sys import argv
+from ROOT import TFile
+import ROOT
+
+
+def main(input_name):
+    f = TFile(input_name, "READ")
+    # f.ls()
+    lk = f.GetListOfKeys()
+
+    def check_dir(directory):
+        print(directory, type(directory))
+        directory.ls()
+        t = {}
+        for i in ["O2bc", "O2track", "O2collision", "O2mccollision",
+                  "O2mcparticle", "O2mctracklabel", "O2mccollisionlabel"]:
+            t[i] = directory.Get(i)
+        if False:
+            for i in t:
+                t[i].Print()
+        # t["O2track"].Print()
+        ROOT.gInterpreter.Declare("""
+        int Count(int x) {
+            if(x >= 998)
+                return 1;
+            return 0;
+        }
+        """)
+        df = ROOT.RDataFrame(directory.GetName()+"/O2track", input_name)
+        print("Has", df.Count().GetValue(), "tracks")
+        df = df.Define("HowMany", "Count(fCollisionsID)").Sum("HowMany").GetValue()
+        print(df)
+
+    for i in lk:
+        d = f.Get(i.GetName())
+        if "TDirectoryFile" not in d.ClassName():
+            continue
+        check_dir(d)
+
+
+if __name__ == "__main__":
+    for i in argv[1:]:
+        main(i)

examples/scripts/diagnostic_tools/doanalysis.py

@@ -200,7 +200,7 @@ def build_list_of_files(file_list):
                 with open(run_list[-1], "w") as f:
                     for j in lines:
                         f.write(j.strip() + "\n")
-        msg("Number or runs:", len(run_list))
+        msg("Number of runs:", len(run_list))
         return run_list
 
     if type(input_file) is list:
@@ -259,7 +259,7 @@ def build_list_of_files(file_list):
                 continue
             merged_files.append(merged_file)
             run_command(f"hadd {merged_file} " + " ".join(files_per_type[i]))
-        if len(merged_files) > 0:
+        if len(merged_files) == 0:
             warning_msg("Merged no files")
         else:
             msg("Merging completed, merged:", *merged_files,

examples/scripts/diagnostic_tools/inspect_hepmc.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+
+"""
+Inspector of the HepMC file
+"""
+
+import pyhepmc_ng as hep
+import argparse
+
+
+poi = {}
+poi_names = {2212: "Proton", -2212: "AntiProton",
+             2112: "Neutron", -2112: "AntiNeutron",
+             3122: "Lambda0", 1000010020: "Deuteron",
+             1000020030: "Triton", 1000020040: "Alpha",
+             1000020030: "Helium3"}
+
+for i in poi_names:
+    poi[i] = 0
+
+
+def main(file_name, min_event, max_event, verbose):
+    print("Reading", file_name, "between",
+          min_event, "and", max_event, "events")
+
+    def print_evt(evt):
+        def msg(*m):
+            if verbose:
+                print(*m)
+        msg("event_number:", evt.event_number)
+        msg("Units:", "momentum_unit:", evt.momentum_unit,
+            "length_unit:", evt.length_unit)
+        msg(len(evt.particles), "particles:")
+        for i in enumerate(evt.particles):
+            pdg = i[1].pid
+            if pdg in poi:
+                poi[pdg] = poi[pdg]+1
+                pdg = f"{pdg} is of interest!!!"
+            msg(i, "PDG code", pdg)
+        msg(len(evt.vertices), "vertices:")
+        for i in enumerate(evt.vertices):
+            msg("Vertex:", i)
+            vertex_pdgs = []
+            msg("Input particles")
+            for j in i[1].particles_in:
+                msg("\t", j, "pdg", j.pid)
+            msg("Output particles")
+            for j in i[1].particles_out:
+                msg("\t", j, "pdg", j.pid)
+                vertex_pdgs.append(j.pid)
+            if 2212 in vertex_pdgs and 2112 in vertex_pdgs:
+                print(evt.event_number, "Has both")
+                print(i)
+                for j in i[1].particles_out:
+                    print(j)
+
+    with hep.open(file_name) as f:
+        while True:
+            e = f.read()
+            if not e:
+                break
+            if e.event_number < min_event:
+                continue
+            print_evt(e)
+            if e.event_number >= max_event:
+                break
+    for i in poi:
+        print("Number of", poi_names[i]+"s", poi[i])
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("hepmcfile", type=str,
+                        help="Input hepmc file.")
+    parser.add_argument("--start", type=int, default=0,
+                        help="Start of the event counter.")
+    parser.add_argument("--stop", type=int, default=100,
+                        help="Stop of the event counter.")
+    parser.add_argument("-v", action="store_true",
+                        help="Verbose mode.")
+    args = parser.parse_args()
+    main(args.hepmcfile, min_event=args.start,
+         max_event=args.stop, verbose=args.v)