algorithms/fardetectors/FarDetectorTransportationPreML.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 - 2025, Simon Gardner
0003
0004 #include <edm4hep/Vector3d.h>
0005 #include <edm4hep/Vector3f.h>
0006 #include <edm4hep/utils/vector_utils.h>
0007 #include <podio/LinkNavigator.h>
0008 #include <cmath>
0009 #include <stdexcept>
0010 #include <tuple>
0011 #include <vector>
0012
0013 #include "FarDetectorTransportationPreML.h"
0014 #include "algorithms/fardetectors/FarDetectorTransportationPreML.h"
0015 #include "algorithms/interfaces/LinkTruthUtils.h"
0016
0017 namespace eicrecon {
0018
0019 void FarDetectorTransportationPreML::init() { m_beamE = m_cfg.beamE; }
0020
0021 void FarDetectorTransportationPreML::process(
0022     const FarDetectorTransportationPreML::Input& input,
0023     const FarDetectorTransportationPreML::Output& output) const {
0024
0025   const auto [inputTracks, trackLinks, beamElectrons] = input;
0026   auto [feature_tensors, target_tensors]              = output;
0027
0028   //Set beam energy from first MCBeamElectron, using std::call_once
0029   if (beamElectrons != nullptr) {
0030     std::call_once(m_initBeamE, [&]() {
0031       // Check if beam electrons are present
0032       if (beamElectrons->empty()) { // NOLINT(clang-analyzer-core.NullDereference)
0033         if (m_cfg.requireBeamElectron) {
0034           error("No beam electrons found");
0035           throw std::runtime_error("No beam electrons found");
0036         }
0037         return;
0038       }
0039       m_beamE = beamElectrons->at(0).getEnergy();
0040       //Round beam energy to nearest GeV - Should be 5, 10 or 18GeV
0041       m_beamE = round(m_beamE);
0042     });
0043   }
0044
0045   edm4eic::MutableTensor feature_tensor = feature_tensors->create();
0046   feature_tensor.addToShape(inputTracks->size());
0047   feature_tensor.addToShape(6);     // x,y,z,dirx,diry,dirz
0048   feature_tensor.setElementType(1); // 1 - float
0049
0050   const truth::EventLinkNavigator<edm4eic::MCRecoTrackParticleLinkCollection> link_nav(trackLinks);
0051
0052   edm4eic::MutableTensor target_tensor;
0053   if (link_nav.enabled() && target_tensors != nullptr) {
0054     target_tensor = target_tensors->create();
0055     target_tensor.addToShape(inputTracks->size());
0056     target_tensor.addToShape(3);     // px,py,pz
0057     target_tensor.setElementType(1); // 1 - float
0058   }
0059
0060   // Loop through inputTracks and simultaneously optionally associations if available
0061   // and fill the feature and target tensors
0062   for (const auto& track : *inputTracks) {
0063
0064     auto position = track.getPosition();
0065     auto momentum = track.getMomentum();
0066
0067     feature_tensor.addToFloatData(position.x); // x
0068     feature_tensor.addToFloatData(position.y); // y
0069     feature_tensor.addToFloatData(position.z); // z
0070     feature_tensor.addToFloatData(momentum.x); // dirx
0071     feature_tensor.addToFloatData(momentum.y); // diry
0072     feature_tensor.addToFloatData(momentum.z); // dirz
0073
0074     if (link_nav.enabled() && target_tensors != nullptr) {
0075       // Use the first linked MC particle, matching previous first-association behavior.
0076       const auto linked_particles = link_nav.linked(track);
0077       if (!linked_particles.empty()) {
0078         auto MCElectronMomentum = linked_particles.front().o.getMomentum() / m_beamE;
0079         target_tensor.addToFloatData(MCElectronMomentum.x);
0080         target_tensor.addToFloatData(MCElectronMomentum.y);
0081         target_tensor.addToFloatData(MCElectronMomentum.z);
0082       } else {
0083         error("Track has no linked MC particle. Targets cannot be constructed.");
0084         throw std::runtime_error("Missing track-particle link required for target tensor");
0085       }
0086     }
0087   }
0088 }
0089
0090 } // namespace eicrecon