Test

Author: vkucera

PWGHF/Core/make_pdg_header.py

@@ -17,11 +17,13 @@
 @date   2023-09-21
 """
 
+import os
 from ctypes import c_bool
 from enum import Enum
 
 import ROOT  # pylint: disable=import-error
 
+name_script = os.path.basename(__file__)
 
 # Enum of PDG_t particles
 class PdgROOT(Enum):
@@ -110,31 +112,44 @@ class Pdg(Enum):
     kXiCCPlusPlus = 4422
     kXiCPlus = 4232
     kXiCZero = 4132
+    kDeuteron = 1000010020
+    kTriton = 1000010030
+    kHelium3 = 1000020030
+    kAlpha = 1000020040
+    kHyperTriton = 1010010030
+    kHyperHydrogen4 = 1010010040
+    kHyperHelium4 = 1010020040
+
+
+dbPdg = ROOT.o2.O2DatabasePDG
 
 
 def mass(code):
-    """Returns particle mass from o2::O2DatabasePDG except for special cases."""
-    # Special cases (present in TDatabasePDG but with wrong values)
+    """Returns particle mass from o2::O2DatabasePDG."""
     # Missing particles should be added in O2DatabasePDG.h.
-    if abs(code) == Pdg.kXiCCPlusPlus.value:
-        return 3.62155  # PDG 2022: https://pdg.lbl.gov/2022/listings/rpp2022-list-xicc-plus-plus.pdf
-    if abs(code) == Pdg.kOmegaC0.value:
-        return 2.69520  # PDG 2022: https://pdg.lbl.gov/2022/listings/rpp2022-list-omegac-zero.pdf
-    # Default case
     success = c_bool(True)
-    return ROOT.o2.O2DatabasePDG.Mass(code, success)
+    return dbPdg.Mass(code, success)
 
 
 def declare_mass(pdg, type="double") -> str:
     """Returns a C++ declaration of a particle mass constant."""
     return f"constexpr {type} Mass{pdg.name[1:]} = {mass(pdg.value)};\n"
 
 
+# Comment at the beginning of the output
+str_block_begin = f"""// BEGINNING OF THE GENERATED BLOCK.
+// DO NOT EDIT THIS BLOCK DIRECTLY!
+// It has been generated by the {name_script} script.
+// For modifications, edit the script and generate this block again.
+"""
+# Comment at the end of the output
+str_block_end = """// END OF THE GENERATED BLOCK
+"""
 # Start of enum declarations of additional particles
 str_enum_head = """/// \\brief Declarations of named PDG codes of particles missing in ROOT PDG_t
 /// \\note Follow kCamelCase naming convention
 /// \\link https://root.cern/doc/master/TPDGCode_8h.html
-enum Code {
+enum Pdg {
 """
 # End of enum declarations of additional particles
 str_enum_foot = "};\n"
@@ -160,5 +175,5 @@ def declare_mass(pdg, type="double") -> str:
     str_mass_root += declare_mass(d)
 
 # Header body
-str_header = "\n".join([str_enum, str_mass_o2, str_mass_root])
+str_header = "\n".join((str_block_begin, str_enum, str_mass_o2, str_mass_root, str_block_end))
 print(str_header)