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 <fmt/format.h>
0006 #include <podio/detail/Link.h>
0007 #include <podio/detail/LinkCollectionImpl.h>
0008 #include <cstddef>
0009 #include <memory>
0010 #include <stdexcept>
0011 #include <tuple>
0012
0013 #include "CalorimeterParticleIDPostML.h"
0014
0015 namespace eicrecon {
0016
0017 void CalorimeterParticleIDPostML::init() {
0018   // Nothing
0019 }
0020
0021 void CalorimeterParticleIDPostML::process(const CalorimeterParticleIDPostML::Input& input,
0022                                           const CalorimeterParticleIDPostML::Output& output) const {
0023
0024   const auto [in_clusters, in_assocs, prediction_tensors]      = input;
0025   auto [out_clusters, out_links, out_assocs, out_particle_ids] = output;
0026
0027   if (prediction_tensors->size() != 1) {
0028     error("Expected to find a single tensor, found {}", prediction_tensors->size());
0029     throw std::runtime_error("");
0030   }
0031   edm4eic::Tensor prediction_tensor = (*prediction_tensors)[0];
0032
0033   if (prediction_tensor.shape_size() != 2) {
0034     error("Expected tensor rank to be 2, but it is {}", prediction_tensor.shape_size());
0035     throw std::runtime_error(
0036         fmt::format("Expected tensor rank to be 2, but it is {}", prediction_tensor.shape_size()));
0037   }
0038
0039   if (prediction_tensor.getShape(0) != static_cast<long>(in_clusters->size())) {
0040     error("Length mismatch between tensor's 0th axis and number of clusters: {} != {}",
0041           prediction_tensor.getShape(0), in_clusters->size());
0042     throw std::runtime_error(
0043         fmt::format("Length mismatch between tensor's 0th axis and number of clusters: {} != {}",
0044                     prediction_tensor.getShape(0), in_clusters->size()));
0045   }
0046
0047   if (prediction_tensor.getShape(1) != 2) {
0048     error("Expected 2 values per cluster in the output tensor, got {}",
0049           prediction_tensor.getShape(0));
0050     throw std::runtime_error(
0051         fmt::format("Expected 2 values per cluster in the output tensor, got {}",
0052                     prediction_tensor.getShape(0)));
0053   }
0054
0055   if (prediction_tensor.getElementType() != 1) { // 1 - float
0056     error("Expected a tensor of floats, but element type is {}",
0057           prediction_tensor.getElementType());
0058     throw std::runtime_error(fmt::format("Expected a tensor of floats, but element type is {}",
0059                                          prediction_tensor.getElementType()));
0060   }
0061
0062   for (std::size_t cluster_ix = 0; cluster_ix < in_clusters->size(); cluster_ix++) {
0063     edm4eic::Cluster in_cluster         = (*in_clusters)[cluster_ix];
0064     edm4eic::MutableCluster out_cluster = in_cluster.clone();
0065     out_clusters->push_back(out_cluster);
0066
0067     float prob_pion =
0068         prediction_tensor.getFloatData(cluster_ix * prediction_tensor.getShape(1) + 0);
0069     float prob_electron =
0070         prediction_tensor.getFloatData(cluster_ix * prediction_tensor.getShape(1) + 1);
0071
0072     out_cluster.addToParticleIDs(out_particle_ids->create(0,        // std::int32_t type
0073                                                           211,      // std::int32_t PDG
0074                                                           0,        // std::int32_t algorithmType
0075                                                           prob_pion // float likelihood
0076                                                           ));
0077     out_cluster.addToParticleIDs(out_particle_ids->create(0,  // std::int32_t type
0078                                                           11, // std::int32_t PDG
0079                                                           0,  // std::int32_t algorithmType
0080                                                           prob_electron // float likelihood
0081                                                           ));
0082
0083     // propagate associations
0084     for (auto in_assoc : *in_assocs) {
0085       if (in_assoc.getRec() == in_cluster) {
0086         auto out_link = out_links->create();
0087         out_link.setFrom(out_cluster);
0088         out_link.setTo(in_assoc.getSim());
0089         out_link.setWeight(in_assoc.getWeight());
0090         auto out_assoc = in_assoc.clone();
0091         out_assoc.setRec(out_cluster);
0092         out_assocs->push_back(out_assoc);
0093       }
0094     }
0095   }
0096 }
0097
0098 } // namespace eicrecon