Fatras/scripts/generate_particle_data_table.py

0001 #!/usr/bin/env python3
0002 #
0003 # use scikit-hep/particle to generate c++ code for the particle data table.
0004 #
0005
0006 import io
0007 import sys
0008 import subprocess
0009
0010 from particle import Particle
0011
0012
0013 def main(output_file):
0014     """
0015     Generate the code and write it to the given output file.
0016     """
0017     # extract relevant entries into a single table
0018     table = []
0019     for p in Particle.all():
0020         table.append((int(p.pdgid), int(p.three_charge), p.mass, p.name))
0021     # use the extracted table to generate the code
0022     code = generate_code(table)
0023     code = clang_format(code)
0024     output_file.write(code)
0025
0026
0027 CODE_HEADER = """\
0028 // This file is part of the ACTS project.
0029 //
0030 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0031 //
0032 // This Source Code Form is subject to the terms of the Mozilla Public
0033 // License, v. 2.0. If a copy of the MPL was not distributed with this
0034 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0035
0036 // The entries within this file have been automatically created using the
0037 // particle data files from the 2019 edition of the Review of Particle Physics
0038 // by the Berkeley Particle Data Group.
0039
0040 #pragma once
0041
0042 #include <cstdint>
0043 #include <limits>
0044
0045 // Rows within the particle data table are sorted by their signed PDG particle
0046 // number and are then stored column-wise. Since the PDG particle number column
0047 // is sorted it can be used to quickly search for the index of a particle
0048 // within all column arrays.
0049
0050 """
0051
0052
0053 def generate_code(table):
0054     """
0055     Generate
0056     """
0057     # ensure the rows are sorted by the signed pdg number (first column)
0058     table = sorted(table, key=lambda _: _[0])
0059     num_rows = len(table)
0060     # name, c++ type, and output format for each column
0061     columns = [
0062         ("PdgNumber", "int32_t", "{}"),
0063         ("ThreeCharge", "int16_t", "{}"),
0064         ("MassMeV", "float", "{}f"),
0065         ("Name", "char* const   ", '"{}"'),
0066     ]
0067     lines = [
0068         CODE_HEADER,
0069         f"static constexpr uint32_t kParticlesCount = {num_rows}u;",
0070     ]
0071     # build a separate array for each column
0072     for i, (variable_name, type_name, value_format) in enumerate(columns):
0073         lines.append(
0074             f"static const {type_name} kParticles{variable_name}[kParticlesCount] = {{"
0075         )
0076
0077         for row in table:
0078             if i < 3:
0079                 lines.append(f"// {row[-1]}")
0080             if row[i] is None and type_name == "float":
0081                 lines.append("  std::numeric_limits<float>::quiet_NaN(),")
0082             else:
0083                 lines.append("  " + value_format.format(row[i]) + ",")
0084
0085         lines.append("};")
0086     # ensure we end with a newline
0087     lines.append("")
0088     return "\n".join(lines)
0089
0090
0091 def clang_format(content):
0092     """
0093     Format the given content using clang-format and return it.
0094     """
0095     args = [
0096         "clang-format",
0097         "--assume-filename=ParticleData.hpp",
0098         "-",
0099     ]
0100     process = subprocess.run(
0101         args,
0102         input=content,
0103         capture_output=True,
0104         check=True,
0105         encoding="utf-8",
0106         text=True,
0107     )
0108     return process.stdout
0109
0110
0111 if __name__ == "__main__":
0112     if 2 < len(sys.argv):
0113         print("usage: {} [<output_file>]".format(sys.argv[0]))
0114         sys.exit(2)
0115     if len(sys.argv) == 1:
0116         output_file = sys.stdout
0117     else:
0118         # will overwrite existing file
0119         output_file = io.open(sys.argv[1], mode="wt", encoding="utf-8")
0120     main(output_file)