generate_particle_data_table.py

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file generate_particle_data_table.py

#!/usr/bin/env python3
#
# use scikit-hep/particle to generate c++ code for the particle data table.
#

import io
import sys
import subprocess

from particle import Particle


def main(output_file):
    """
    Generate the code and write it to the given output file.
    """
    # extract relevant entries into a single table
    table = []
    for p in Particle.all():
        table.append((int(p.pdgid), int(p.three_charge), p.mass, p.name))
    # use the extracted table to generate the code
    code = generate_code(table)
    code = clang_format(code)
    output_file.write(code)


CODE_HEADER = """\
// This file is part of the Acts project.
//
// Copyright (C) 2020 CERN for the benefit of the Acts project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// The entries within this file have been automatically created using the
// particle data files from the 2019 edition of the Review of Particle Physics
// by the Berkeley Particle Data Group.

#pragma once

#include <cstdint>

// Rows within the particle data table are sorted by their signed PDG particle
// number and are then stored column-wise. Since the PDG particle number column
// is sorted it can be used to quickly search for the index of a particle
// within all column arrays.

"""


def generate_code(table):
    """
    Generate
    """
    # ensure the rows are sorted by the signed pdg number (first column)
    table = sorted(table, key=lambda _: _[0])
    num_rows = len(table)
    # name, c++ type, and output format for each column
    columns = [
        ("PdgNumber", "int32_t", "{}"),
        ("ThreeCharge", "int8_t", "{}"),
        ("MassMeV", "float", "{}f"),
        ("Name", "char* const   ", '"{}"'),
    ]
    lines = [
        CODE_HEADER,
        f"static constexpr uint32_t kParticlesCount = {num_rows}u;",
    ]
    # build a separate array for each column
    for i, (variable_name, type_name, value_format) in enumerate(columns):
        lines.append(
            f"static const {type_name} kParticles{variable_name}[kParticlesCount] = {{"
        )
        lines.append(
            "  " + ", ".join(value_format.format(row[i]) for row in table) + ","
        )
        lines.append("};")
    # ensure we end with a newline
    lines.append("")
    return "\n".join(lines)


def clang_format(content):
    """
    Format the given content using clang-format and return it.
    """
    args = [
        "clang-format",
        "--assume-filename=ParticleData.hpp",
        "-",
    ]
    process = subprocess.run(
        args,
        input=content,
        capture_output=True,
        check=True,
        encoding="utf-8",
        text=True,
    )
    return process.stdout


if __name__ == "__main__":
    if 2 < len(sys.argv):
        print("usage: {} [<output_file>]".format(sys.argv[0]))
        sys.exit(2)
    if len(sys.argv) == 1:
        output_file = sys.stdout
    else:
        # will overwrite existing file
        output_file = io.open(sys.argv[1], mode="wt", encoding="utf-8")
    main(output_file)