services/pid_lut/PIDLookupTable.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024, Nathan Brei, Dmitry Kalinkin
0003
0004 #include "services/pid_lut/PIDLookupTable.h"
0005
0006 #include <boost/histogram.hpp>
0007 #include <boost/iostreams/categories.hpp>
0008 #include <boost/iostreams/close.hpp>
0009 #include <boost/iostreams/filter/gzip.hpp>
0010 #include <boost/iostreams/filtering_stream.hpp>
0011 #include <fmt/core.h>
0012 #include <math.h>
0013 #include <stdlib.h>
0014 #include <algorithm>
0015 #include <cctype>
0016 #include <fstream> // IWYU pragma: keep
0017 #include <iterator>
0018 #include <sstream> // IWYU pragma: keep
0019 #include <stdexcept>
0020 // IWYU pragma: no_include <boost/mp11/detail/mp_defer.hpp>
0021
0022 namespace bh = boost::histogram;
0023
0024 namespace eicrecon {
0025
0026 const PIDLookupTable::Entry* PIDLookupTable::Lookup(int pdg, int charge, double momentum, double theta_deg, double phi_deg) const {
0027     // Our lookup table expects _unsigned_ PDGs. The charge information is passed separately.
0028     pdg = std::abs(pdg);
0029
0030     if (m_symmetrizing_charges) {
0031       charge = std::abs(charge);
0032     }
0033
0034     return &m_hist[
0035       decltype(m_hist)::multi_index_type {
0036         m_hist.axis(0).index(pdg),
0037         m_hist.axis(1).index(charge),
0038         m_hist.axis(2).index(momentum),
0039         m_hist.axis(3).index(theta_deg),
0040         m_hist.axis(4).index(phi_deg)
0041       }
0042     ];
0043 }
0044
0045 void PIDLookupTable::load_file(const std::string& filename, const PIDLookupTable::Binning &binning) {
0046     bool is_compressed = filename.substr(filename.length() - 3) == ".gz";
0047
0048     std::ios_base::openmode mode = std::ios_base::in;
0049     if (is_compressed) {
0050       mode |= std::ios_base::binary;
0051     }
0052
0053     std::ifstream file(filename, mode);
0054     if (!file) {
0055       throw std::runtime_error("Unable to open LUT file!");
0056     }
0057     boost::iostreams::filtering_istream in;
0058     if (is_compressed) {
0059       in.push(boost::iostreams::gzip_decompressor());
0060     }
0061     in.push(file);
0062
0063     std::string line;
0064     std::istringstream iss;
0065     double step;
0066
0067     const double angle_fudge = binning.use_radians ? 180. / M_PI : 1.;
0068
0069     bh::axis::category<int> pdg_bins(binning.pdg_values);
0070     bh::axis::category<int> charge_bins(binning.charge_values);
0071     bh::axis::variable<> momentum_bins(binning.momentum_edges);
0072     std::vector<double> polar_edges = binning.polar_edges;
0073     for (double &edge : polar_edges) {
0074       edge *= angle_fudge;
0075     }
0076     bh::axis::variable<> polar_bins(polar_edges);
0077     bh::axis::circular<> azimuthal_bins(bh::axis::step(binning.azimuthal_binning.at(2) * angle_fudge), binning.azimuthal_binning.at(0) * angle_fudge, binning.azimuthal_binning.at(1) * angle_fudge);
0078
0079     m_hist = bh::make_histogram_with(bh::dense_storage<PIDLookupTable::Entry>(), pdg_bins, charge_bins, momentum_bins, polar_bins, azimuthal_bins);
0080
0081     m_symmetrizing_charges = binning.charge_values.size() == 1;
0082
0083     while (std::getline(in, line)) {
0084         Entry entry;
0085         if (line.empty() || line[0] == '#' || std::all_of(std::begin(line), std::end(line), [](unsigned char c) { return std::isspace(c); })) continue;
0086
0087         iss.str(line);
0088         iss.clear();
0089         double pdg, charge, momentum, eta, phi, prob_electron, prob_pion, prob_kaon, prob_proton;
0090         // Read each field from the line and assign to Entry struct members
0091         if ((bool)(iss >> pdg
0092                 >> charge
0093                 >> momentum
0094                 >> eta
0095                 >> phi) && (binning.missing_electron_prob || (bool)(iss
0096                 >> prob_electron)) && (bool)(iss
0097                 >> prob_pion
0098                 >> prob_kaon
0099                 >> prob_proton)) {
0100
0101             if (m_symmetrizing_charges) {
0102               charge = std::abs(charge);
0103             }
0104
0105             // operator() here allows to lookup mutable entry and increases the access counter
0106             auto &entry = *m_hist(
0107               pdg,
0108               charge,
0109               momentum + (binning.momentum_bin_centers_in_lut ? 0. : (momentum_bins.bin(0).width() / 2)),
0110               eta * angle_fudge + (binning.polar_bin_centers_in_lut ? 0. : (polar_bins.bin(0).width() / 2)),
0111               phi * angle_fudge + (binning.azimuthal_bin_centers_in_lut ? 0. : (azimuthal_bins.bin(0).width() / 2))
0112             ); // N.B. bin(0) may not be of a correct width
0113             entry.prob_electron = prob_electron;
0114             entry.prob_pion = prob_pion;
0115             entry.prob_kaon = prob_kaon;
0116             entry.prob_proton = prob_proton;
0117         }
0118         else {
0119             error("Unable to parse LUT file!");
0120             throw std::runtime_error("Unable to parse LUT file!");
0121         }
0122     }
0123
0124     for (auto&& b : bh::indexed(m_hist)) {
0125       if (b->value() != 1) {
0126         error(
0127           "Bin {} {} {}:{} {}:{} {}:{} is defined {} times in the PID table",
0128           b.bin(0).lower(),
0129           b.bin(1).lower(),
0130           b.bin(2).lower(),
0131           b.bin(2).upper(),
0132           b.bin(3).lower() / angle_fudge,
0133           b.bin(3).upper() / angle_fudge,
0134           b.bin(4).lower() / angle_fudge,
0135           b.bin(4).upper() / angle_fudge,
0136           b->value()
0137         );
0138       }
0139     }
0140
0141     boost::iostreams::close(in);
0142     file.close();
0143 }
0144
0145 }