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/Vector3f.h>
0010 #include <edm4hep/utils/vector_utils.h>
0011 #include <cmath>
0012 #include <gsl/pointers>
0013
0014 #include "FarDetectorTransportationPreML.h"
0015 #include "algorithms/fardetectors/FarDetectorTransportationPreML.h"
0016
0017 namespace eicrecon {
0018
0019   void FarDetectorTransportationPreML::init() {
0020
0021     m_beamE = m_cfg.beamE;
0022
0023   }
0024
0025   void FarDetectorTransportationPreML::process(
0026       const FarDetectorTransportationPreML::Input& input,
0027       const FarDetectorTransportationPreML::Output& output) {
0028
0029     const auto [inputTracks,MCElectrons,beamElectrons] = input;
0030     auto [feature_tensors, target_tensors] = output;
0031
0032     //Set beam energy from first MCBeamElectron, using std::call_once
0033     if (beamElectrons)
0034     {
0035       std::call_once(m_initBeamE,[&](){
0036         // Check if beam electrons are present
0037         if(beamElectrons->size() == 0){
0038           error("No beam electrons found keeping default 10GeV beam energy.");
0039           return;
0040         }
0041         m_beamE = beamElectrons->at(0).getEnergy();
0042         //Round beam energy to nearest GeV - Should be 5, 10 or 18GeV
0043         m_beamE = round(m_beamE);
0044       });
0045     }
0046
0047     edm4eic::MutableTensor feature_tensor = feature_tensors->create();
0048     feature_tensor.addToShape(inputTracks->size());
0049     feature_tensor.addToShape(4); // x,z,dirx,diry
0050     feature_tensor.setElementType(1); // 1 - float
0051
0052     edm4eic::MutableTensor target_tensor;
0053     if (MCElectrons) {
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     for(const auto& track: *inputTracks){
0061
0062       auto pos        = track.getLoc();
0063       auto trackphi   = track.getPhi();
0064       auto tracktheta = track.getTheta();
0065
0066       feature_tensor.addToFloatData(pos.a); // x
0067       feature_tensor.addToFloatData(pos.b); // z
0068       feature_tensor.addToFloatData(sin(trackphi)*sin(tracktheta)); // dirx
0069       feature_tensor.addToFloatData(cos(trackphi)*sin(tracktheta)); // diry
0070
0071       if (MCElectrons) {
0072         // FIXME: use proper MC matching once available again, assume training sample is indexed correctly
0073         // Take the first scattered/simulated electron
0074         auto MCElectron = MCElectrons->at(0);
0075         auto MCElectronMomentum = MCElectron.getMomentum()/m_beamE;
0076         target_tensor.addToFloatData(MCElectronMomentum.x);
0077         target_tensor.addToFloatData(MCElectronMomentum.y);
0078         target_tensor.addToFloatData(MCElectronMomentum.z);
0079       }
0080
0081     }
0082
0083   }
0084
0085 } // namespace eicrecon
0086 #endif