algorithms/onnx/CalorimeterParticleIDPreML.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 Dmitry Kalinkin
0003
0004 #include <edm4hep/MCParticle.h>
0005 #include <edm4hep/Vector3f.h>
0006 #include <edm4hep/utils/vector_utils.h>
0007 #include <podio/LinkNavigator.h>
0008 #include <cmath>
0009 #include <cstddef>
0010 #include <cstdint>
0011 #include <format>
0012 #include <limits>
0013 #include <stdexcept>
0014 #include <tuple>
0015 #include <vector>
0016
0017 #include "CalorimeterParticleIDPreML.h"
0018 #include "algorithms/interfaces/LinkTruthUtils.h"
0019
0020 namespace eicrecon {
0021
0022 void CalorimeterParticleIDPreML::init() {
0023   // Nothing
0024 }
0025
0026 void CalorimeterParticleIDPreML::process(const CalorimeterParticleIDPreML::Input& input,
0027                                          const CalorimeterParticleIDPreML::Output& output) const {
0028
0029   const auto [clusters, cluster_links]   = input;
0030   auto [feature_tensors, target_tensors] = output;
0031   const truth::EventLinkNavigator<edm4eic::MCRecoClusterParticleLinkCollection> link_nav(
0032       cluster_links);
0033   const bool fill_targets = link_nav.enabled();
0034
0035   edm4eic::MutableTensor feature_tensor = feature_tensors->create();
0036   feature_tensor.addToShape(clusters->size());
0037   feature_tensor.addToShape(11);    // p, E/p, azimuthal, polar, 7 shape parameters
0038   feature_tensor.setElementType(1); // 1 - float
0039
0040   edm4eic::MutableTensor target_tensor;
0041   if (fill_targets) {
0042     target_tensor = target_tensors->create();
0043     target_tensor.addToShape(clusters->size());
0044     target_tensor.addToShape(2);     // is electron, is hadron
0045     target_tensor.setElementType(7); // 7 - int64
0046   }
0047
0048   for (edm4eic::Cluster cluster : *clusters) {
0049     double momentum = NAN;
0050     edm4hep::MCParticle best_sim;
0051     if (fill_targets) {
0052       // FIXME: use track momentum once matching to tracks becomes available
0053       float best_weight = std::numeric_limits<float>::lowest();
0054       bool found_assoc  = false;
0055       for (const auto& [sim_particle, weight] : link_nav.linked(cluster)) {
0056         if (!found_assoc || weight > best_weight) {
0057           best_sim    = sim_particle;
0058           best_weight = weight;
0059           found_assoc = true;
0060         }
0061       }
0062       if (!found_assoc) {
0063         error("Can't find link for cluster. Targets cannot be constructed.");
0064         throw std::runtime_error("Missing cluster-particle link required for target tensor");
0065       }
0066       momentum = edm4hep::utils::magnitude(best_sim.getMomentum());
0067     }
0068
0069     feature_tensor.addToFloatData(momentum);
0070     feature_tensor.addToFloatData(cluster.getEnergy() / momentum);
0071     auto pos = cluster.getPosition();
0072     feature_tensor.addToFloatData(edm4hep::utils::anglePolar(pos));
0073     feature_tensor.addToFloatData(edm4hep::utils::angleAzimuthal(pos));
0074     for (std::size_t par_ix = 0; par_ix < cluster.shapeParameters_size(); par_ix++) {
0075       feature_tensor.addToFloatData(cluster.getShapeParameters(par_ix));
0076     }
0077
0078     if (fill_targets) {
0079       auto is_electron = static_cast<int64_t>(best_sim.getPDG() == 11);
0080       auto is_pion     = static_cast<int64_t>(best_sim.getPDG() != 11);
0081       target_tensor.addToInt64Data(is_pion);
0082       target_tensor.addToInt64Data(is_electron);
0083     }
0084   }
0085
0086   std::size_t expected_num_entries = feature_tensor.getShape(0) * feature_tensor.getShape(1);
0087   if (feature_tensor.floatData_size() != expected_num_entries) {
0088     error("Inconsistent output tensor shape and element count: {} != {}",
0089           feature_tensor.floatData_size(), expected_num_entries);
0090     throw std::runtime_error(
0091         std::format("Inconsistent output tensor shape and element count: {} != {}",
0092                     feature_tensor.floatData_size(), expected_num_entries));
0093   }
0094 }
0095
0096 } // namespace eicrecon