algorithms/onnx/CalorimeterParticleIDPreML.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 Dmitry Kalinkin
0003
0004 #include <edm4eic/EDM4eicVersion.h>
0005
0006 #if EDM4EIC_VERSION_MAJOR >= 8
0007 #include <cstddef>
0008 #include <edm4hep/MCParticle.h>
0009 #include <edm4hep/Vector3f.h>
0010 #include <edm4hep/utils/vector_utils.h>
0011 #include <cstdint>
0012 #include <stdexcept>
0013 #include <fmt/core.h>
0014 #include <gsl/pointers>
0015
0016 #include "CalorimeterParticleIDPreML.h"
0017
0018 namespace eicrecon {
0019
0020   void CalorimeterParticleIDPreML::init() {
0021     // Nothing
0022   }
0023
0024   void CalorimeterParticleIDPreML::process(
0025       const CalorimeterParticleIDPreML::Input& input,
0026       const CalorimeterParticleIDPreML::Output& output) const {
0027
0028     const auto [clusters, cluster_assocs] = input;
0029     auto [feature_tensors, target_tensors] = output;
0030
0031     edm4eic::MutableTensor feature_tensor = feature_tensors->create();
0032     feature_tensor.addToShape(clusters->size());
0033     feature_tensor.addToShape(11); // p, E/p, azimuthal, polar, 7 shape parameters
0034     feature_tensor.setElementType(1); // 1 - float
0035
0036     edm4eic::MutableTensor target_tensor;
0037     if (cluster_assocs) {
0038       target_tensor = target_tensors->create();
0039       target_tensor.addToShape(clusters->size());
0040       target_tensor.addToShape(2); // is electron, is hadron
0041       target_tensor.setElementType(7); // 7 - int64
0042     }
0043
0044     for (edm4eic::Cluster cluster : *clusters) {
0045       double momentum;
0046       {
0047         // FIXME: use track momentum once matching to tracks becomes available
0048         edm4eic::MCRecoClusterParticleAssociation best_assoc;
0049         for (auto assoc : *cluster_assocs) {
0050           if (assoc.getRec() == cluster) {
0051             if ((not best_assoc.isAvailable()) || (assoc.getWeight() > best_assoc.getWeight())) {
0052               best_assoc = assoc;
0053             }
0054           }
0055         }
0056         if (best_assoc.isAvailable()) {
0057           momentum = edm4hep::utils::magnitude(best_assoc.getSim().getMomentum());
0058         } else {
0059           warning("Can't find association for cluster. Skipping...");
0060           continue;
0061         }
0062       }
0063
0064       feature_tensor.addToFloatData(momentum);
0065       feature_tensor.addToFloatData(cluster.getEnergy() / momentum);
0066       auto pos = cluster.getPosition();
0067       feature_tensor.addToFloatData(edm4hep::utils::anglePolar(pos));
0068       feature_tensor.addToFloatData(edm4hep::utils::angleAzimuthal(pos));
0069       for (int par_ix = 0; par_ix < cluster.shapeParameters_size(); par_ix++) {
0070         feature_tensor.addToFloatData(cluster.getShapeParameters(par_ix));
0071       }
0072
0073       if (cluster_assocs) {
0074         edm4eic::MCRecoClusterParticleAssociation best_assoc;
0075         for (auto assoc : *cluster_assocs) {
0076           if (assoc.getRec() == cluster) {
0077             if ((not best_assoc.isAvailable()) || (assoc.getWeight() > best_assoc.getWeight())) {
0078               best_assoc = assoc;
0079             }
0080           }
0081         }
0082         int64_t is_electron = 0, is_pion = 0;
0083         if (best_assoc.isAvailable()) {
0084           is_electron = best_assoc.getSim().getPDG() == 11;
0085           is_pion = best_assoc.getSim().getPDG() != 11;
0086         }
0087         target_tensor.addToInt64Data(is_pion);
0088         target_tensor.addToInt64Data(is_electron);
0089       }
0090     }
0091
0092     size_t expected_num_entries = feature_tensor.getShape(0) * feature_tensor.getShape(1);
0093     if (feature_tensor.floatData_size() != expected_num_entries) {
0094       error("Inconsistent output tensor shape and element count: {} != {}", feature_tensor.floatData_size(), expected_num_entries);
0095       throw std::runtime_error(fmt::format("Inconsistent output tensor shape and element count: {} != {}", feature_tensor.floatData_size(), expected_num_entries));
0096     }
0097   }
0098
0099 } // namespace eicrecon
0100 #endif