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 <cmath>
0008 #include <gsl/pointers>
0009 #include <stdexcept>
0010
0011 #include "FarDetectorTransportationPreML.h"
0012 #include "algorithms/fardetectors/FarDetectorTransportationPreML.h"
0013
0014 namespace eicrecon {
0015
0016 void FarDetectorTransportationPreML::init() { m_beamE = m_cfg.beamE; }
0017
0018 void FarDetectorTransportationPreML::process(
0019     const FarDetectorTransportationPreML::Input& input,
0020     const FarDetectorTransportationPreML::Output& output) const {
0021
0022   const auto [inputTracks, mcAssociation, beamElectrons] = input;
0023   auto [feature_tensors, target_tensors]                 = output;
0024
0025   //Set beam energy from first MCBeamElectron, using std::call_once
0026   if (beamElectrons != nullptr) {
0027     std::call_once(m_initBeamE, [&]() {
0028       // Check if beam electrons are present
0029       if (beamElectrons->empty()) { // NOLINT(clang-analyzer-core.NullDereference)
0030         if (m_cfg.requireBeamElectron) {
0031           critical("No beam electrons found");
0032           throw std::runtime_error("No beam electrons found");
0033         }
0034         return;
0035       }
0036       m_beamE = beamElectrons->at(0).getEnergy();
0037       //Round beam energy to nearest GeV - Should be 5, 10 or 18GeV
0038       m_beamE = round(m_beamE);
0039     });
0040   }
0041
0042   edm4eic::MutableTensor feature_tensor = feature_tensors->create();
0043   feature_tensor.addToShape(inputTracks->size());
0044   feature_tensor.addToShape(6);     // x,y,z,dirx,diry,dirz
0045   feature_tensor.setElementType(1); // 1 - float
0046
0047   edm4eic::MutableTensor target_tensor;
0048   if (mcAssociation != nullptr) {
0049     target_tensor = target_tensors->create();
0050     target_tensor.addToShape(inputTracks->size());
0051     target_tensor.addToShape(3);     // px,py,pz
0052     target_tensor.setElementType(1); // 1 - float
0053   }
0054
0055   // Loop through inputTracks and simultaneously optionally associations if available
0056   // and fill the feature and target tensors
0057   for (const auto& track : *inputTracks) {
0058
0059     auto position = track.getPosition();
0060     auto momentum = track.getMomentum();
0061
0062     feature_tensor.addToFloatData(position.x); // x
0063     feature_tensor.addToFloatData(position.y); // y
0064     feature_tensor.addToFloatData(position.z); // z
0065     feature_tensor.addToFloatData(momentum.x); // dirx
0066     feature_tensor.addToFloatData(momentum.y); // diry
0067     feature_tensor.addToFloatData(momentum.z); // dirz
0068
0069     if ((mcAssociation != nullptr) && (!mcAssociation->empty())) {
0070       //Loop through the MCRecoTrackParticleAssociationCollection finding the first one associated with the current track
0071       for (const auto& assoc : *mcAssociation) {
0072         if (assoc.getRec() == track) {
0073           // Process the association if it exists and is non-empty
0074           const auto& association = assoc.getSim(); // Assuming 1-to-1 mapping
0075           auto MCElectronMomentum = association.getMomentum() / m_beamE;
0076           target_tensor.addToFloatData(MCElectronMomentum.x);
0077           target_tensor.addToFloatData(MCElectronMomentum.y);
0078           target_tensor.addToFloatData(MCElectronMomentum.z);
0079           break; // Exit loop after finding the first association
0080         }
0081       }
0082     }
0083   }
0084 }
0085
0086 } // namespace eicrecon