Implement PID lookup table (#1365)

### Briefly, what does this PR introduce?

Implement PID lookup table as per:

-
https://docs.google.com/document/d/1ViIBV0448Vnob0zE_Sm7kQsggsKOJZ9oYdmPp--QPvE/edit#heading=h.tf2uykgho79e
- https://github.com/rdom/fastpid/

### What kind of change does this PR introduce?
- [ ] Bug fix (issue #__)
- [x] New feature (issue #__)
- [ ] Documentation update
- [ ] Other: __

### Please check if this PR fulfills the following:
- [x] Tests for the changes have been added
- [ ] Documentation has been added / updated
- [ ] Changes have been communicated to collaborators

### Does this PR introduce breaking changes? What changes might users
need to make to their code?
No
### Does this PR change default behavior?
No
However, it does download a ~100MB file to the working directory.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Dmitry Kalinkin <[email protected]>
Co-authored-by: Wouter Deconinck <[email protected]>

.github/iwyu.imp

@@ -3,6 +3,7 @@
 
   # Boost
   { ref: '/usr/local/share/include-what-you-use/boost-all.imp' },
+  { include: ['@<boost/histogram/.*\.hpp>', private, '<boost/histogram.hpp>', public] },
 
   # libc++
   { ref: '/usr/local/share/include-what-you-use/libcxx.imp' },

.github/workflows/linux-eic-shell.yml

@@ -144,6 +144,7 @@ jobs:
         platform-release: "${{ env.platform }}:${{ matrix.release }}"
         run: |
           export LD_LIBRARY_PATH=$PWD/install/lib:$LD_LIBRARY_PATH
+          export JANA_PLUGIN_PATH=$PWD/install/lib/EICrecon/plugins${JANA_PLUGIN_PATH:+:${JANA_PLUGIN_PATH}}
           ctest --test-dir build -V
     - name: Compress install directory
       run: tar -caf install.tar.zst install/
@@ -267,8 +268,9 @@ jobs:
         llvm-profdata-15 merge -sparse src/tests/algorithms_test/algorithms_test.profraw \
           -o src/tests/algorithms_test/algorithms_test.profdata
         LIB_PATHS=()
-        for LIB_NAME in algorithms_calorimetry algorithms_digi algorithms_pid algorithms_reco algorithms_tracking; do
-          LIB_PATH="$(find src/ -type f -name "*$LIB_NAME.so")"
+        for LIB_PATH in $(find src/ -type f -name "libalgorithms_*.so"); do
+          LIB_NAME="$(basename ${LIB_PATH%%.so})"
+          LIB_NAME="${LIB_NAME##lib}"
           LIB_PATHS+=("$LIB_PATH")
           llvm-cov-15 report $LIB_PATH \
             -instr-profile=src/tests/algorithms_test/algorithms_test.profdata \

src/algorithms/CMakeLists.txt

@@ -4,6 +4,7 @@ add_subdirectory(interfaces)
 add_subdirectory(calorimetry)
 add_subdirectory(tracking)
 add_subdirectory(pid)
+add_subdirectory(pid_lut)
 add_subdirectory(digi)
 add_subdirectory(reco)
 add_subdirectory(fardetectors)

src/algorithms/calorimetry/CMakeLists.txt

@@ -12,7 +12,7 @@ plugin_add(${PLUGIN_NAME} WITH_SHARED_LIBRARY WITHOUT_PLUGIN)
 plugin_glob_all(${PLUGIN_NAME})
 
 # Find dependencies
-plugin_add_algorithms(${PLUGIN_NAME} algocalorimetry)
+plugin_add_algorithms(${PLUGIN_NAME})
 plugin_add_dd4hep(${PLUGIN_NAME})
 plugin_add_event_model(${PLUGIN_NAME})
 plugin_add_eigen3(${PLUGIN_NAME})

src/algorithms/pid_lut/CMakeLists.txt

@@ -0,0 +1,14 @@
+set(PLUGIN_NAME "algorithms_pid_lut")
+
+# Function creates ${PLUGIN_NAME}_plugin and ${PLUGIN_NAME}_library targets
+# Setting default includes, libraries and installation paths
+plugin_add(${PLUGIN_NAME} WITH_SHARED_LIBRARY WITHOUT_PLUGIN)
+
+# The macro grabs sources as *.cc *.cpp *.c and headers as *.h *.hh *.hpp Then
+# correctly sets sources for ${_name}_plugin and ${_name}_library targets Adds
+# headers to the correct installation directory
+plugin_glob_all(${PLUGIN_NAME})
+
+# Find dependencies
+plugin_add_algorithms(${PLUGIN_NAME})
+plugin_add_event_model(${PLUGIN_NAME})

src/algorithms/pid_lut/PIDLookup.cc

@@ -0,0 +1,142 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024, Nathan Brei, Dmitry Kalinkin
+
+#include <algorithms/service.h>
+#include <edm4eic/MCRecoParticleAssociationCollection.h>
+#include <edm4eic/ReconstructedParticleCollection.h>
+#include <edm4hep/MCParticleCollection.h>
+#include <edm4hep/utils/vector_utils.h>
+#include <fmt/core.h>
+#include <podio/ObjectID.h>
+#include <cmath>
+#include <gsl/pointers>
+#include <stdexcept>
+
+#include "algorithms/pid_lut/PIDLookup.h"
+#include "algorithms/pid_lut/PIDLookupConfig.h"
+#include "services/pid_lut/PIDLookupTableSvc.h"
+
+namespace eicrecon {
+
+void PIDLookup::init() {
+  auto& serviceSvc = algorithms::ServiceSvc::instance();
+  auto lut_svc = serviceSvc.service<PIDLookupTableSvc>("PIDLookupTableSvc");
+
+  m_lut = lut_svc->load(m_cfg.filename, {
+    .pdg_values=m_cfg.pdg_values,
+    .charge_values=m_cfg.charge_values,
+    .momentum_edges=m_cfg.momentum_edges,
+    .polar_edges=m_cfg.polar_edges,
+    .azimuthal_binning=m_cfg.azimuthal_binning,
+    .azimuthal_bin_centers_in_lut=m_cfg.azimuthal_bin_centers_in_lut,
+    .momentum_bin_centers_in_lut=m_cfg.momentum_bin_centers_in_lut,
+    .polar_bin_centers_in_lut=m_cfg.polar_bin_centers_in_lut,
+    .skip_legacy_header=m_cfg.skip_legacy_header,
+    .use_radians=m_cfg.use_radians,
+    .missing_electron_prob=m_cfg.missing_electron_prob,
+  });
+  if (m_lut == nullptr) {
+    throw std::runtime_error("LUT not available");
+  }
+}
+
+void PIDLookup::process(const Input& input, const Output& output) const {
+  const auto [recoparts_in, partassocs_in]          = input;
+  auto [recoparts_out, partassocs_out, partids_out] = output;
+
+  for (const auto& recopart_without_pid : *recoparts_in) {
+    edm4hep::MCParticle mcpart;
+    auto recopart = recopart_without_pid.clone();
+
+    // Find MCParticle from associations and propagate the relevant ones further
+    bool assoc_found = false;
+    for (auto assoc_in : *partassocs_in) {
+      if (assoc_in.getRec() == recopart_without_pid) {
+        if (assoc_found) {
+          warning("Found a duplicate association for ReconstructedParticle at index {}", recopart_without_pid.getObjectID().index);
+          warning("The previous MCParticle was at {} and the duplicate is at {}", mcpart.getObjectID().index, assoc_in.getSim().getObjectID().index);
+        }
+        assoc_found    = true;
+        mcpart         = assoc_in.getSim();
+        auto assoc_out = assoc_in.clone();
+        assoc_out.setRec(recopart);
+        partassocs_out->push_back(assoc_out);
+      }
+    }
+    if (not assoc_found) {
+      recoparts_out->push_back(recopart);
+      continue;
+    }
+
+    int true_pdg    = mcpart.getPDG();
+    int true_charge = mcpart.getCharge();
+    int charge      = recopart.getCharge();
+    double momentum = edm4hep::utils::magnitude(recopart.getMomentum());
+
+    double theta = edm4hep::utils::anglePolar(recopart.getMomentum()) / M_PI * 180.;
+    double phi   = edm4hep::utils::angleAzimuthal(recopart.getMomentum()) / M_PI * 180.;
+
+    trace("lookup for true_pdg={}, true_charge={}, momentum={:.2f} GeV, polar={:.2f}, aziumthal={:.2f}",
+      true_pdg, true_charge, momentum, theta, phi);
+    auto entry = m_lut->Lookup(true_pdg, true_charge, momentum, theta, phi);
+
+    int identified_pdg = 0; // unknown
+
+    if ((entry != nullptr) && ((entry->prob_electron != 0.) || (entry->prob_pion != 0.) || (entry->prob_kaon != 0.) || (entry->prob_electron != 0.))) {
+      double random_unit_interval = m_dist(m_gen);
+
+      trace("entry with e:pi:K:P={}:{}:{}:{}", entry->prob_electron, entry->prob_pion, entry->prob_kaon, entry->prob_proton);
+
+      recopart.addToParticleIDs(partids_out->create(
+        m_cfg.system,                // std::int32_t type
+        std::copysign(11, charge),   // std::int32_t PDG
+        0,                           // std::int32_t algorithmType
+        static_cast<float>(entry->prob_electron) // float likelihood
+      ));
+      recopart.addToParticleIDs(partids_out->create(
+        m_cfg.system,                // std::int32_t type
+        std::copysign(211, charge),  // std::int32_t PDG
+        0,                           // std::int32_t algorithmType
+        static_cast<float>(entry->prob_pion) // float likelihood
+      ));
+      recopart.addToParticleIDs(partids_out->create(
+        m_cfg.system,                // std::int32_t type
+        std::copysign(321, charge),  // std::int32_t PDG
+        0,                           // std::int32_t algorithmType
+        static_cast<float>(entry->prob_kaon) // float likelihood
+      ));
+      recopart.addToParticleIDs(partids_out->create(
+        m_cfg.system,                // std::int32_t type
+        std::copysign(2212, charge), // std::int32_t PDG
+        0,                           // std::int32_t algorithmType
+        static_cast<float>(entry->prob_proton) // float likelihood
+      ));
+
+      if (random_unit_interval < entry->prob_electron) {
+        identified_pdg = 11; // electron
+        recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 4]);
+      } else if (random_unit_interval < (entry->prob_electron + entry->prob_pion)) {
+        identified_pdg = 211; // pion
+        recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 3]);
+      } else if (random_unit_interval <
+                 (entry->prob_electron + entry->prob_pion + entry->prob_kaon)) {
+        identified_pdg = 321; // kaon
+        recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 2]);
+      } else if (random_unit_interval < (entry->prob_electron + entry->prob_pion +
+                                         entry->prob_kaon + entry->prob_electron)) {
+        identified_pdg = 2212; // proton
+        recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 1]);
+      }
+    }
+
+    recopart.setPDG(std::copysign(identified_pdg, charge));
+
+    if (identified_pdg != 0) {
+      trace("randomized PDG is {}", recopart.getPDG());
+    }
+
+    recoparts_out->push_back(recopart);
+  }
+}
+
+} // namespace eicrecon

src/algorithms/pid_lut/PIDLookup.h

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024, Nathan Brei, Dmitry Kalinkin
+
+#include <algorithms/algorithm.h>
+#include <edm4eic/MCRecoParticleAssociationCollection.h>
+#include <edm4eic/ReconstructedParticleCollection.h>
+#include <edm4hep/ParticleIDCollection.h>
+#include <random>
+#include <string>
+#include <string_view>
+
+#include "PIDLookupConfig.h"
+#include "algorithms/interfaces/WithPodConfig.h"
+#include "services/pid_lut/PIDLookupTable.h"
+
+namespace eicrecon {
+
+using PIDLookupAlgorithm =
+    algorithms::Algorithm<algorithms::Input<edm4eic::ReconstructedParticleCollection,
+                                            edm4eic::MCRecoParticleAssociationCollection>,
+                          algorithms::Output<edm4eic::ReconstructedParticleCollection,
+                                             edm4eic::MCRecoParticleAssociationCollection,
+                                             edm4hep::ParticleIDCollection>>;
+
+class PIDLookup : public PIDLookupAlgorithm, public WithPodConfig<PIDLookupConfig> {
+
+public:
+  PIDLookup(std::string_view name)
+      : PIDLookupAlgorithm{name,
+                           {"inputParticlesCollection", "inputParticleAssociationsCollection"},
+                           {"outputParticlesCollection", "outputParticleAssociationsCollection",
+                            "outputParticleIDCollection"},
+                           ""} {}
+
+  void init() final;
+  void process(const Input&, const Output&) const final;
+
+private:
+  mutable std::mt19937 m_gen{};
+  mutable std::uniform_real_distribution<double> m_dist{0, 1};
+  const PIDLookupTable* m_lut;
+};
+
+} // namespace eicrecon

src/algorithms/pid_lut/PIDLookupConfig.h

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024, Nathan Brei, Dmitry Kalinkin
+
+#pragma once
+
+#include <vector>
+
+namespace eicrecon {
+
+struct PIDLookupConfig {
+  std::string filename;
+  int system;
+  std::vector<int> pdg_values;
+  std::vector<int> charge_values;
+  std::vector<double> momentum_edges;
+  std::vector<double> polar_edges;
+  std::vector<double> azimuthal_binning;
+  bool azimuthal_bin_centers_in_lut {false};
+  bool momentum_bin_centers_in_lut {false};
+  bool polar_bin_centers_in_lut {false};
+  bool skip_legacy_header {false};
+  bool use_radians {false};
+  bool missing_electron_prob {false};
+};
+
+} // namespace eicrecon

src/algorithms/tracking/TracksToParticles.cc

@@ -127,15 +127,13 @@ namespace eicrecon {
             }
             auto rec_part = parts->create();
             rec_part.addToTracks(track);
-            int32_t best_pid = 0;
             auto referencePoint = rec_part.getReferencePoint();
             // float time          = 0;
             float mass = 0;
             if (best_match >= 0) {
                 trace("Best match is found and is: {}", best_match);
                 mc_prt_is_consumed[best_match] = true;
                 const auto &best_mc_part = (*mc_particles)[best_match];
-                best_pid = best_mc_part.getPDG();
                 referencePoint = {
                         static_cast<float>(best_mc_part.getVertex().x),
                         static_cast<float>(best_mc_part.getVertex().y),
@@ -151,7 +149,6 @@ namespace eicrecon {
             rec_part.setCharge(charge_rec);
             rec_part.setMass(mass);
             rec_part.setGoodnessOfPID(0); // assume no PID until proven otherwise
-            rec_part.setPDG(best_pid);
             // rec_part.covMatrix()  // @TODO: covariance matrix on 4-momentum
 
             // Also write MC <--> truth particle association if match was found

src/detectors/BTOF/BTOF.cc

@@ -1,16 +1,27 @@
 // Copyright 2022, Dmitry Romanov
 // Subject to the terms in the LICENSE file found in the top-level directory.
 //
-//
 
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024, Dmitry Kalinkin
+
+#include <DD4hep/Detector.h>
 #include <Evaluator/DD4hepUnits.h>
 #include <JANA/JApplication.h>
+#include <fmt/core.h>
+#include <algorithm>
+#include <gsl/pointers>
+#include <memory>
+#include <stdexcept>
 #include <string>
+#include <vector>
 
 #include "algorithms/interfaces/WithPodConfig.h"
 #include "extensions/jana/JOmniFactoryGeneratorT.h"
 #include "factories/digi/SiliconTrackerDigi_factory.h"
 #include "factories/tracking/TrackerHitReconstruction_factory.h"
+#include "global/pid_lut/PIDLookup_factory.h"
+#include "services/geometry/dd4hep/DD4hep_service.h"
 
 extern "C" {
 void InitPlugin(JApplication *app) {
@@ -46,5 +57,39 @@ void InitPlugin(JApplication *app) {
         app
     ));         // Hit reco default config for factories
 
+    int BarrelTOF_ID = 0;
+    try {
+        auto detector = app->GetService<DD4hep_service>()->detector();
+        BarrelTOF_ID = detector->constant<int>("BarrelTOF_ID");
+    } catch(const std::runtime_error&) {
+        // Nothing
+    }
+    for (auto qualifier : std::vector<std::string>({"", "Seeded"})) {
+        app->Add(new JOmniFactoryGeneratorT<PIDLookup_factory>(
+              fmt::format("CombinedTOF{}LUTPID", qualifier),
+              {
+              fmt::format("Reconstructed{}ChargedWithPFRICHPIDParticles", qualifier),
+              fmt::format("Reconstructed{}ChargedWithPFRICHPIDParticleAssociations", qualifier),
+              },
+              {
+              fmt::format("Reconstructed{}ChargedWithPFRICHTOFPIDParticles", qualifier),
+              fmt::format("Reconstructed{}ChargedWithPFRICHTOFPIDParticleAssociations", qualifier),
+              fmt::format("CombinedTOF{}ParticleIDs", qualifier),
+              },
+              {
+                .filename="calibrations/tof.lut",
+                .system=BarrelTOF_ID,
+                .pdg_values={11, 211, 321, 2212},
+                .charge_values={1},
+                .momentum_edges={0.0, 0.3, 0.6, 0.9, 1.2, 1.5, 1.8, 2.1, 2.4, 2.7, 3.0, 3.3, 3.6, 3.9, 4.2, 4.5, 4.8, 5.1, 5.4, 5.7, 6.0},
+                .polar_edges={2.50, 10.95, 19.40, 27.85, 36.30, 44.75, 53.20, 61.65, 70.10, 78.55, 87.00, 95.45, 103.90, 112.35, 120.80, 129.25, 137.70, 146.15, 154.60},
+                .azimuthal_binning={0., 360., 360.}, // lower, upper, step
+                .momentum_bin_centers_in_lut=true,
+                .polar_bin_centers_in_lut=true,
+              },
+              app
+              ));
+    }
+
 }
 } // extern "C"

src/detectors/DIRC/CMakeLists.txt

@@ -15,3 +15,7 @@ plugin_glob_all(${PLUGIN_NAME})
 # Find dependencies
 plugin_add_event_model(${PLUGIN_NAME})
 plugin_add_dd4hep(${PLUGIN_NAME})
+
+# Add libraries (same as target_include_directories but for both plugin and
+# library)
+plugin_link_libraries(${PLUGIN_NAME} algorithms_pid_lut_library)