services/pid_lut/PIDLookupTable.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024, Nathan Brei, Dmitry Kalinkin
0003
0004 #pragma once
0005
0006 #include <algorithms/logger.h>
0007 #include <boost/histogram.hpp>
0008 #include <string>
0009 #include <tuple>
0010 #include <variant>
0011 #include <vector>
0012 // IWYU pragma: no_include <boost/mp11/detail/mp_defer.hpp>
0013
0014 namespace eicrecon {
0015
0016 class PIDLookupTable : public algorithms::LoggerMixin {
0017
0018 public:
0019   /// The histogram entry with access counter, where the probabilities are stored in metadata
0020   struct Entry : public boost::histogram::accumulators::count<unsigned char, false> {
0021     float prob_electron, prob_pion, prob_kaon, prob_proton;
0022   };
0023
0024   struct Binning {
0025     std::vector<int> pdg_values;
0026     std::vector<int> charge_values;
0027     std::vector<double> momentum_edges;
0028     std::vector<double> polar_edges;
0029     std::vector<double> azimuthal_binning;
0030     bool azimuthal_bin_centers_in_lut;
0031     bool momentum_bin_centers_in_lut;
0032     bool polar_bin_centers_in_lut;
0033     bool use_radians;
0034     bool missing_electron_prob;
0035   };
0036
0037 private:
0038   boost::histogram::histogram<
0039       std::tuple<boost::histogram::axis::category<int>, boost::histogram::axis::category<int>,
0040                  boost::histogram::axis::variable<>, boost::histogram::axis::variable<>,
0041                  boost::histogram::axis::circular<>>,
0042       boost::histogram::dense_storage<Entry>>
0043       m_hist;
0044   bool m_symmetrizing_charges;
0045
0046 public:
0047   PIDLookupTable() : algorithms::LoggerMixin("PIDLookupTable") {};
0048
0049   const Entry* Lookup(int pdg, int charge, double momentum, double theta_deg, double phi_deg) const;
0050
0051   void load_file(const std::string& filename, const Binning& binning);
0052 };
0053
0054 } // namespace eicrecon