algorithms/fardetectors/FarDetectorTransportationPreML.cc

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