algorithms/onnx/CalorimeterParticleIDPostML.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 Dmitry Kalinkin
0003
0004 #include <edm4hep/MCParticle.h>
0005 #include <podio/LinkNavigator.h>
0006 #include <podio/detail/LinkCollectionImpl.h>
0007 #include <cstddef>
0008 #include <format>
0009 #include <gsl/pointers>
0010 #include <stdexcept>
0011 #include <tuple>
0012 #include <vector>
0013
0014 #include "CalorimeterParticleIDPostML.h"
0015 #include "algorithms/interfaces/LinkTruthUtils.h"
0016
0017 namespace eicrecon {
0018
0019 void CalorimeterParticleIDPostML::init() {
0020   // Nothing
0021 }
0022
0023 void CalorimeterParticleIDPostML::process(const CalorimeterParticleIDPostML::Input& input,
0024                                           const CalorimeterParticleIDPostML::Output& output) const {
0025
0026   const auto [in_clusters, in_links, prediction_tensors]       = input;
0027   auto [out_clusters, out_links, out_assocs, out_particle_ids] = output;
0028
0029   const truth::EventLinkNavigator<edm4eic::MCRecoClusterParticleLinkCollection> link_nav(in_links);
0030
0031   if (prediction_tensors->size() != 1) {
0032     error("Expected to find a single tensor, found {}", prediction_tensors->size());
0033     throw std::runtime_error("");
0034   }
0035   edm4eic::Tensor prediction_tensor = (*prediction_tensors)[0];
0036
0037   if (prediction_tensor.shape_size() != 2) {
0038     error("Expected tensor rank to be 2, but it is {}", prediction_tensor.shape_size());
0039     throw std::runtime_error(
0040         std::format("Expected tensor rank to be 2, but it is {}", prediction_tensor.shape_size()));
0041   }
0042
0043   if (prediction_tensor.getShape(0) != static_cast<long>(in_clusters->size())) {
0044     error("Length mismatch between tensor's 0th axis and number of clusters: {} != {}",
0045           prediction_tensor.getShape(0), in_clusters->size());
0046     throw std::runtime_error(
0047         std::format("Length mismatch between tensor's 0th axis and number of clusters: {} != {}",
0048                     prediction_tensor.getShape(0), in_clusters->size()));
0049   }
0050
0051   if (prediction_tensor.getShape(1) != 2) {
0052     error("Expected 2 values per cluster in the output tensor, got {}",
0053           prediction_tensor.getShape(0));
0054     throw std::runtime_error(
0055         std::format("Expected 2 values per cluster in the output tensor, got {}",
0056                     prediction_tensor.getShape(0)));
0057   }
0058
0059   if (prediction_tensor.getElementType() != 1) { // 1 - float
0060     error("Expected a tensor of floats, but element type is {}",
0061           prediction_tensor.getElementType());
0062     throw std::runtime_error(std::format("Expected a tensor of floats, but element type is {}",
0063                                          prediction_tensor.getElementType()));
0064   }
0065
0066   for (std::size_t cluster_ix = 0; cluster_ix < in_clusters->size(); cluster_ix++) {
0067     edm4eic::Cluster in_cluster         = (*in_clusters)[cluster_ix];
0068     edm4eic::MutableCluster out_cluster = in_cluster.clone();
0069     out_clusters->push_back(out_cluster);
0070
0071     float prob_pion =
0072         prediction_tensor.getFloatData(cluster_ix * prediction_tensor.getShape(1) + 0);
0073     float prob_electron =
0074         prediction_tensor.getFloatData(cluster_ix * prediction_tensor.getShape(1) + 1);
0075
0076     out_cluster.addToParticleIDs(out_particle_ids->create(0,        // std::int32_t type
0077                                                           211,      // std::int32_t PDG
0078                                                           0,        // std::int32_t algorithmType
0079                                                           prob_pion // float likelihood
0080                                                           ));
0081     out_cluster.addToParticleIDs(out_particle_ids->create(0,  // std::int32_t type
0082                                                           11, // std::int32_t PDG
0083                                                           0,  // std::int32_t algorithmType
0084                                                           prob_electron // float likelihood
0085                                                           ));
0086
0087     // propagate associations
0088     if (link_nav.enabled() && out_links != nullptr && out_assocs != nullptr) {
0089       for (const auto& [sim_particle, weight] : link_nav.linked(in_cluster)) {
0090         truth::addWeightedRelation(
0091             out_cluster, sim_particle, weight,
0092             gsl::not_null<edm4eic::MCRecoClusterParticleLinkCollection*>{out_links},
0093             gsl::not_null<edm4eic::MCRecoClusterParticleAssociationCollection*>{out_assocs});
0094       }
0095     }
0096   }
0097 }
0098
0099 } // namespace eicrecon