File indexing completed on 2024-09-27 07:03:08
0001
0002
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
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
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
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
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 );
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 }