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/Vector2f.h>
0009 #include <edm4hep/Vector3d.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, MCElectrons, 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(4);     // x,z,dirx,diry
0049   feature_tensor.setElementType(1); // 1 - float
0050
0051   edm4eic::MutableTensor target_tensor;
0052   if (MCElectrons != 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   for (const auto& track : *inputTracks) {
0060
0061     auto pos        = track.getLoc();
0062     auto trackphi   = track.getPhi();
0063     auto tracktheta = track.getTheta();
0064
0065     feature_tensor.addToFloatData(pos.a);                           // x
0066     feature_tensor.addToFloatData(pos.b);                           // z
0067     feature_tensor.addToFloatData(sin(trackphi) * sin(tracktheta)); // dirx
0068     feature_tensor.addToFloatData(cos(trackphi) * sin(tracktheta)); // diry
0069
0070     if (MCElectrons != nullptr) {
0071       // FIXME: use proper MC matching once available again, assume training sample is indexed correctly
0072       // Take the first scattered/simulated electron
0073       auto MCElectron         = MCElectrons->at(0);
0074       auto MCElectronMomentum = MCElectron.getMomentum() / m_beamE;
0075       target_tensor.addToFloatData(MCElectronMomentum.x);
0076       target_tensor.addToFloatData(MCElectronMomentum.y);
0077       target_tensor.addToFloatData(MCElectronMomentum.z);
0078     }
0079   }
0080 }
0081
0082 } // namespace eicrecon
0083 #endif