Low-Q2 particle reconstruction using onnx framework

src/algorithms/fardetectors/FarDetectorTransportationPostML.cc

@@ -0,0 +1,107 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 - 2025 Simon Gardner
+
+#include <edm4eic/EDM4eicVersion.h>
+
+#if EDM4EIC_VERSION_MAJOR >= 8
+#include <cmath>
+#include <cstddef>
+#include <fmt/core.h>
+#include <gsl/pointers>
+#include <podio/RelationRange.h>
+#include <stdexcept>
+
+#include "FarDetectorTransportationPostML.h"
+
+namespace eicrecon {
+
+  void FarDetectorTransportationPostML::init() {
+
+    m_beamE = m_cfg.beamE;
+
+  }
+
+  void FarDetectorTransportationPostML::process(
+      const FarDetectorTransportationPostML::Input& input,
+      const FarDetectorTransportationPostML::Output& output) const {
+
+    const auto [prediction_tensors,beamElectrons] = input;
+    auto [out_particles] = output;
+
+    //Set beam energy from first MCBeamElectron, using std::call_once
+    if (beamElectrons)
+    {
+      std::call_once(m_initBeamE,[&](){
+        // Check if beam electrons are present
+        if(beamElectrons->size() == 0){
+          error("No beam electrons found keeping default 10GeV beam energy.");
+          return;
+        }
+        m_beamE = beamElectrons->at(0).getEnergy();
+        //Round beam energy to nearest GeV - Should be 5, 10 or 18GeV
+        m_beamE = round(m_beamE);
+      });
+    }
+
+    if (prediction_tensors->size() != 1) {
+      error("Expected to find a single tensor, found {}", prediction_tensors->size());
+      throw std::runtime_error("");
+    }
+    edm4eic::Tensor prediction_tensor = (*prediction_tensors)[0];
+
+    if (prediction_tensor.shape_size() != 2) {
+      error("Expected tensor rank to be 2, but it is {}", prediction_tensor.shape_size());
+      throw std::runtime_error(fmt::format("Expected tensor rank to be 2, but it is {}", prediction_tensor.shape_size()));
+    }
+
+    if (prediction_tensor.getShape(1) != 3) {
+      error("Expected 2 values per cluster in the output tensor, got {}", prediction_tensor.getShape(0));
+      throw std::runtime_error(fmt::format("Expected 2 values per cluster in the output tensor, got {}", prediction_tensor.getShape(0)));
+    }
+
+    if (prediction_tensor.getElementType() != 1) { // 1 - float
+      error("Expected a tensor of floats, but element type is {}", prediction_tensor.getElementType());
+      throw std::runtime_error(fmt::format("Expected a tensor of floats, but element type is {}", prediction_tensor.getElementType()));
+    }
+
+    auto prediction_tensor_data = prediction_tensor.getFloatData();
+
+    // Ensure the size of prediction_tensor_data is a multiple of its shape
+    if (prediction_tensor_data.size() % 3 != 0) {
+        error("The size of prediction_tensor_data is not a multiple of 3.");
+        throw std::runtime_error("The size of prediction_tensor_data is not a multiple of 3.");
+    }
+
+
+    edm4eic::MutableReconstructedParticle particle;
+
+    // Iterate over the prediction_tensor_data in steps of three
+    for (size_t i = 0; i < prediction_tensor_data.size(); i += 3) {
+        if (i + 2 >= prediction_tensor_data.size()) {
+            error("Incomplete data for a prediction tensor at the end of the vector.");
+            throw std::runtime_error("Incomplete data for a prediction tensor at the end of the vector.");
+        }
+
+        // Extract the current prediction
+        float px = prediction_tensor_data[i] * m_beamE;
+        float py = prediction_tensor_data[i + 1] * m_beamE;
+        float pz = prediction_tensor_data[i + 2] * m_beamE;
+
+        // Calculate reconstructed electron energy
+        double energy = sqrt(px * px + py * py + pz * pz + 0.000511 * 0.000511);
+
+        particle = out_particles->create();
+
+        particle.setEnergy(energy);
+        particle.setMomentum({px, py, pz});
+        particle.setCharge(-1);
+        particle.setMass(0.000511);
+        particle.setPDG(11);
+    }
+
+    // TODO: Implement the association of the reconstructed particles with the tracks
+
+  }
+
+} // namespace eicrecon
+#endif

src/algorithms/fardetectors/FarDetectorTransportationPostML.h

@@ -0,0 +1,52 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 - 2025 Simon Gardner
+
+#pragma once
+
+#include <algorithms/algorithm.h>
+#include <edm4eic/ReconstructedParticleCollection.h>
+#include <edm4eic/TensorCollection.h>
+#include <edm4hep/MCParticleCollection.h>
+#include <mutex>
+#include <optional>
+#include <string>
+#include <string_view>
+
+#include "algorithms/fardetectors/FarDetectorTransportationPostMLConfig.h"
+#include "algorithms/interfaces/WithPodConfig.h"
+
+namespace eicrecon {
+
+using FarDetectorTransportationPostMLAlgorithm =
+  algorithms::Algorithm<
+    algorithms::Input<
+      edm4eic::TensorCollection,
+      std::optional<edm4hep::MCParticleCollection>
+    >,
+    algorithms::Output<
+      edm4eic::ReconstructedParticleCollection
+    >
+  >;
+
+class FarDetectorTransportationPostML
+: public FarDetectorTransportationPostMLAlgorithm,
+  public WithPodConfig<FarDetectorTransportationPostMLConfig> {
+
+public:
+  FarDetectorTransportationPostML(std::string_view name)
+      : FarDetectorTransportationPostMLAlgorithm{name,
+                            {"inputPredictionsTensor"},
+                            {"outputParticles"},
+                            "Convert ML output tensor into reconstructed electron"} {
+  }
+
+  void init() final;
+  void process(const Input&, const Output&) const final;
+
+  private:
+    mutable float m_beamE = 10.0;
+    mutable std::once_flag m_initBeamE;
+
+};
+
+} // namespace eicrecon

src/algorithms/fardetectors/FarDetectorTransportationPostMLConfig.h

@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 - 2025, Simon Gardner
+
+#pragma once
+
+namespace eicrecon {
+  struct FarDetectorTransportationPostMLConfig {
+
+    float beamE = 10.0;
+
+  };
+}

src/algorithms/fardetectors/FarDetectorTransportationPreML.cc

@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 - 2025, Simon Gardner
+
+#include <edm4eic/EDM4eicVersion.h>
+
+#if EDM4EIC_VERSION_MAJOR >= 8
+
+#include <edm4hep/Vector2f.h>
+#include <edm4hep/Vector3f.h>
+#include <edm4hep/utils/vector_utils.h>
+#include <cmath>
+#include <gsl/pointers>
+
+#include "FarDetectorTransportationPreML.h"
+#include "algorithms/fardetectors/FarDetectorTransportationPreML.h"
+
+namespace eicrecon {
+
+  void FarDetectorTransportationPreML::init() {
+
+    m_beamE = m_cfg.beamE;
+
+  }
+
+  void FarDetectorTransportationPreML::process(
+      const FarDetectorTransportationPreML::Input& input,
+      const FarDetectorTransportationPreML::Output& output) {
+
+    const auto [inputTracks,MCElectrons,beamElectrons] = input;
+    auto [feature_tensors, target_tensors] = output;
+
+    //Set beam energy from first MCBeamElectron, using std::call_once
+    if (beamElectrons)
+    {
+      std::call_once(m_initBeamE,[&](){
+        // Check if beam electrons are present
+        if(beamElectrons->size() == 0){
+          error("No beam electrons found keeping default 10GeV beam energy.");
+          return;
+        }
+        m_beamE = beamElectrons->at(0).getEnergy();
+        //Round beam energy to nearest GeV - Should be 5, 10 or 18GeV
+        m_beamE = round(m_beamE);
+      });
+    }
+
+    edm4eic::MutableTensor feature_tensor = feature_tensors->create();
+    feature_tensor.addToShape(inputTracks->size());
+    feature_tensor.addToShape(4); // x,z,dirx,diry
+    feature_tensor.setElementType(1); // 1 - float
+
+    edm4eic::MutableTensor target_tensor;
+    if (MCElectrons) {
+      target_tensor = target_tensors->create();
+      target_tensor.addToShape(inputTracks->size());
+      target_tensor.addToShape(3); // px,py,pz
+      target_tensor.setElementType(1); // 1 - float
+    }
+
+    for(const auto& track: *inputTracks){
+
+      auto pos        = track.getLoc();
+      auto trackphi   = track.getPhi();
+      auto tracktheta = track.getTheta();
+
+      feature_tensor.addToFloatData(pos.a); // x
+      feature_tensor.addToFloatData(pos.b); // z
+      feature_tensor.addToFloatData(sin(trackphi)*sin(tracktheta)); // dirx
+      feature_tensor.addToFloatData(cos(trackphi)*sin(tracktheta)); // diry
+
+      if (MCElectrons) {
+        // FIXME: use proper MC matching once available again, assume training sample is indexed correctly
+        // Take the first scattered/simulated electron
+        auto MCElectron = MCElectrons->at(0);
+        auto MCElectronMomentum = MCElectron.getMomentum()/m_beamE;
+        target_tensor.addToFloatData(MCElectronMomentum.x);
+        target_tensor.addToFloatData(MCElectronMomentum.y);
+        target_tensor.addToFloatData(MCElectronMomentum.z);
+      }
+
+    }
+
+  }
+
+} // namespace eicrecon
+#endif

src/algorithms/fardetectors/FarDetectorTransportationPreML.h

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 - 2025, Simon Gardner
+
+#pragma once
+
+#include <algorithms/algorithm.h>
+#include <edm4eic/TensorCollection.h>
+#include <edm4eic/TrackParametersCollection.h>
+#include <edm4hep/MCParticleCollection.h>
+#include <mutex>
+#include <optional>
+#include <string>
+#include <string_view>
+
+#include "algorithms/fardetectors/FarDetectorTransportationPreMLConfig.h"
+#include "algorithms/interfaces/WithPodConfig.h"
+
+namespace eicrecon {
+
+  using FarDetectorTransportationPreMLAlgorithm = algorithms::Algorithm<
+    algorithms::Input<
+      edm4eic::TrackParametersCollection,
+      std::optional<edm4hep::MCParticleCollection>,
+      std::optional<edm4hep::MCParticleCollection>
+    >,
+    algorithms::Output<
+      edm4eic::TensorCollection,
+      std::optional<edm4eic::TensorCollection>
+    >
+  >;
+
+  class FarDetectorTransportationPreML
+  : public FarDetectorTransportationPreMLAlgorithm,
+    public WithPodConfig<FarDetectorTransportationPreMLConfig> {
+
+  public:
+      FarDetectorTransportationPreML(std::string_view name)
+        : FarDetectorTransportationPreMLAlgorithm{name,
+                              {"TrackParameters","ScatteredElectrons","BeamElectrons"},
+                              {"outputFeatureTensor", "outputTargetTensor"},
+                              "Create tensor for input to far-detector magnetic transportation ML."} {}
+
+
+      void init() final;
+      void process(const Input&, const Output&);
+
+  private:
+    mutable float m_beamE = 10.0;
+    mutable std::once_flag m_initBeamE;
+
+  };
+
+} // eicrecon

src/algorithms/fardetectors/FarDetectorTransportationPreMLConfig.h

@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Copyright (C) 2024 - 2025, Simon Gardner
+
+#pragma once
+
+namespace eicrecon {
+  struct FarDetectorTransportationPreMLConfig {
+
+    float beamE = 10.0;
+
+  };
+}

src/detectors/LOWQ2/CMakeLists.txt

@@ -15,6 +15,11 @@ plugin_glob_all(${PLUGIN_NAME})
 # Find dependencies
 plugin_add_dd4hep(${PLUGIN_NAME})
 plugin_add_event_model(${PLUGIN_NAME})
+plugin_add_onnxruntime(${PLUGIN_NAME})
 
 # Add libraries (works same as target_include_directories)
-plugin_link_libraries(${PLUGIN_NAME} algorithms_fardetectors_library)
+plugin_link_libraries(${PLUGIN_NAME} algorithms_digi_library
+                      algorithms_fardetectors_library)
+if(USE_ONNX)
+  plugin_link_libraries(${PLUGIN_NAME} algorithms_onnx_library)
+endif()