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