algorithms/pid/ConvertParticleID.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2023, Christopher Dilks
0003
0004 // Converters of ParticleID objects
0005 //
0006
0007 #pragma once
0008
0009 #include <cstddef>
0010 #include <map>
0011
0012 // data model
0013 #include <edm4eic/CherenkovParticleID.h>
0014 #include <edm4eic/CherenkovParticleIDHypothesis.h>
0015 #include <edm4hep/ParticleIDCollection.h>
0016
0017 namespace eicrecon {
0018
0019 class ConvertParticleID {
0020 public:
0021   // index map, used for external access to the converted objects
0022   using IndexMap = std::map<std::size_t, unsigned int>; // <collection index, object ID>
0023
0024   // convert edm4eic::CherenkovParticleID hypotheses to list of edm4hep::ParticleID objects
0025   // - requires input `CherenkovParticleID` object `in_pid`
0026   // - adds output `ParticleID` objects to collection `out_pids`
0027   // - sorted by likelihood, if `sort_by_likelihood==true`
0028   // - returns a map of the new `out_pid` indices to new `ParticleID` IDs, so the caller can
0029   //   access the newly created `ParticleID` objects
0030   static IndexMap ConvertToParticleIDs(const edm4eic::CherenkovParticleID& in_pid,
0031                                        edm4hep::ParticleIDCollection& out_pids,
0032                                        bool sort_by_likelihood = false) {
0033
0034     // output vector of collection indices
0035     IndexMap out_indices;
0036
0037     // build list of (hypothesis index, hypothesis weight)
0038     using HypIndex =
0039         std::pair<std::size_t, decltype(edm4eic::CherenkovParticleIDHypothesis::weight)>;
0040     std::vector<HypIndex> hyp_indices;
0041     for (std::size_t hyp_index = 0; hyp_index < in_pid.hypotheses_size(); hyp_index++)
0042       hyp_indices.push_back(HypIndex{hyp_index, in_pid.getHypotheses(hyp_index).weight});
0043
0044     // sort it by likelihood, if needed
0045     if (sort_by_likelihood)
0046       std::sort(hyp_indices.begin(), hyp_indices.end(),
0047                 [](HypIndex& a, HypIndex& b) { return a.second > b.second; });
0048
0049     // create and fill output objects
0050     for (const auto& [hyp_index, hyp_weight] : hyp_indices) {
0051
0052       // get the input hypothesis
0053       auto in_hyp = in_pid.getHypotheses(hyp_index);
0054
0055       // create output `ParticleID` object
0056       auto out_index = out_pids.size();
0057       auto out_pid   = out_pids.create();
0058       out_indices.insert({out_index, out_pid.getObjectID().index});
0059
0060       // fill scalars
0061       out_pid.setPDG(static_cast<decltype(edm4hep::ParticleIDData::PDG)>(in_hyp.PDG));
0062       out_pid.setLikelihood(
0063           static_cast<decltype(edm4hep::ParticleIDData::likelihood)>(in_hyp.weight));
0064       out_pid.setType(
0065           static_cast<decltype(edm4hep::ParticleIDData::type)>(0)); // FIXME: not used yet
0066       out_pid.setAlgorithmType(
0067           static_cast<decltype(edm4hep::ParticleIDData::algorithmType)>(0)); // FIXME: not used yet
0068
0069       // fill parameters vector
0070       out_pid.addToParameters(static_cast<float>(in_hyp.npe)); // NPE for this hypothesis
0071     }
0072
0073     return out_indices;
0074   }
0075
0076 }; // class ConvertParticleID
0077 } // namespace eicrecon