Fatras/scripts/generate_particle_data_table.py

0001 #!/usr/bin/env python3
0002
0003 # /// script
0004 # dependencies = [
0005 #   "particle==0.24.0",
0006 # ]
0007 # ///
0008
0009 #
0010 # use scikit-hep/particle to generate c++ code for the particle data table.
0011 #
0012
0013 import io
0014 import sys
0015 import subprocess
0016 import argparse
0017 from pathlib import Path
0018
0019 from particle import Particle
0020
0021 CODE_HEADER = """\
0022 // This file is part of the ACTS project.
0023 //
0024 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0025 //
0026 // This Source Code Form is subject to the terms of the Mozilla Public
0027 // License, v. 2.0. If a copy of the MPL was not distributed with this
0028 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0029
0030 // The entries within this file have been automatically created using the
0031 // particle data files from Review of Particle Physics
0032 // by the Berkeley Particle Data Group.
0033
0034 #pragma once
0035
0036 #include <cstdint>
0037 #include <map>
0038 #include <limits>
0039
0040 """
0041
0042
0043 def convert_mass_to_GeV(mass_MeV):
0044     """
0045     Convert the mass from MeV to GeV.
0046     """
0047     # return mass_MeV
0048     return mass_MeV / 1000.0
0049
0050
0051 def convert_charge_to_e(charge):
0052     """
0053     Convert the charge from three-charge to electron charge.
0054     """
0055     return charge / 3.0
0056
0057
0058 def identity(v):
0059     """
0060     Return the value unchanged.
0061     """
0062     return v
0063
0064
0065 def generate_code():
0066     """
0067     Generate
0068     """
0069
0070     lines = CODE_HEADER.splitlines()
0071
0072     lines.append(
0073         f"static constexpr std::uint32_t kParticlesCount = {len(Particle.all())}u;"
0074     )
0075
0076     # ensure the rows are sorted by the signed pdg number (first column)
0077     # name, c++ type, and output format for each column
0078     columns = [
0079         ("three_charge", "Charge", "float", "{}", convert_charge_to_e),
0080         ("mass", "Mass", "float", "{}f", convert_mass_to_GeV),
0081         ("name", "Name", " const  char* const", '"{}"', identity),
0082     ]
0083     # build a separate array for each column
0084     for variable_name, target_name, type_name, value_format, transform in columns:
0085         cpp_name = f"particlesMap{target_name}"
0086         map_type = f"std::map<std::int32_t, {type_name}>"
0087
0088         lines.append(f"inline const {map_type}& {cpp_name}() {{")
0089         lines.append(f"static const {map_type} map = {{")
0090
0091         for p in Particle.all():
0092             value = getattr(p, variable_name)
0093
0094             if variable_name != "Name":
0095                 lines.append(f"// {p.name}")
0096
0097             if value is None and type_name == "float":
0098                 lines.append(
0099                     f"{{ {int(p.pdgid)} , std::numeric_limits<float>::quiet_NaN() }},"
0100                 )
0101             else:
0102                 lines.append(
0103                     f"{{ {int(p.pdgid)}, {value_format.format(transform(value))} }},"
0104                 )
0105
0106         lines.append("};")
0107         lines.append("return map;")
0108         lines.append("}")
0109
0110     # ensure we end with a newline
0111     lines.append("")
0112     return "\n".join(lines)
0113
0114
0115 def clang_format(content):
0116     """
0117     Format the given content using clang-format and return it.
0118     """
0119     args = [
0120         "clang-format",
0121         "--assume-filename=ParticleData.hpp",
0122         "-",
0123     ]
0124     process = subprocess.run(
0125         args,
0126         input=content,
0127         capture_output=True,
0128         check=True,
0129         encoding="utf-8",
0130         text=True,
0131     )
0132     return process.stdout
0133
0134
0135 if __name__ == "__main__":
0136     p = argparse.ArgumentParser(description="Generate the particle data table.")
0137     p.add_argument("output", type=Path, nargs="?", default=None, help="Output file.")
0138     p.add_argument(
0139         "--format", action="store_true", help="Run clang-format on the output."
0140     )
0141
0142     args = p.parse_args()
0143     if args.output is None:
0144         output_file = sys.stdout
0145     else:
0146         # will overwrite existing file
0147         output_file = io.open(args.output, mode="wt", encoding="utf-8")
0148
0149     code = generate_code()
0150     if args.format:
0151         code = clang_format(code)
0152     output_file.write(code)