diff --git a/src/algorithms/fardetectors/FarDetectorTransportationPostML.cc b/src/algorithms/fardetectors/FarDetectorTransportationPostML.cc new file mode 100644 index 0000000000..7bb1e30fc5 --- /dev/null +++ b/src/algorithms/fardetectors/FarDetectorTransportationPostML.cc @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 - 2025 Simon Gardner + +#include + +#if EDM4EIC_VERSION_MAJOR >= 8 +#include +#include +#include +#include +#include +#include + +#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 diff --git a/src/algorithms/fardetectors/FarDetectorTransportationPostML.h b/src/algorithms/fardetectors/FarDetectorTransportationPostML.h new file mode 100644 index 0000000000..3e98476789 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include + +#include "algorithms/fardetectors/FarDetectorTransportationPostMLConfig.h" +#include "algorithms/interfaces/WithPodConfig.h" + +namespace eicrecon { + +using FarDetectorTransportationPostMLAlgorithm = + algorithms::Algorithm< + algorithms::Input< + edm4eic::TensorCollection, + std::optional + >, + algorithms::Output< + edm4eic::ReconstructedParticleCollection + > + >; + +class FarDetectorTransportationPostML +: public FarDetectorTransportationPostMLAlgorithm, + public WithPodConfig { + +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 diff --git a/src/algorithms/fardetectors/FarDetectorTransportationPostMLConfig.h b/src/algorithms/fardetectors/FarDetectorTransportationPostMLConfig.h new file mode 100644 index 0000000000..35b9ba2388 --- /dev/null +++ b/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; + + }; +} diff --git a/src/algorithms/fardetectors/FarDetectorTransportationPreML.cc b/src/algorithms/fardetectors/FarDetectorTransportationPreML.cc new file mode 100644 index 0000000000..763fbb77cf --- /dev/null +++ b/src/algorithms/fardetectors/FarDetectorTransportationPreML.cc @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 - 2025, Simon Gardner + +#include + +#if EDM4EIC_VERSION_MAJOR >= 8 + +#include +#include +#include +#include +#include + +#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 diff --git a/src/algorithms/fardetectors/FarDetectorTransportationPreML.h b/src/algorithms/fardetectors/FarDetectorTransportationPreML.h new file mode 100644 index 0000000000..4f8009364e --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include + +#include "algorithms/fardetectors/FarDetectorTransportationPreMLConfig.h" +#include "algorithms/interfaces/WithPodConfig.h" + +namespace eicrecon { + + using FarDetectorTransportationPreMLAlgorithm = algorithms::Algorithm< + algorithms::Input< + edm4eic::TrackParametersCollection, + std::optional, + std::optional + >, + algorithms::Output< + edm4eic::TensorCollection, + std::optional + > + >; + + class FarDetectorTransportationPreML + : public FarDetectorTransportationPreMLAlgorithm, + public WithPodConfig { + + 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 diff --git a/src/algorithms/fardetectors/FarDetectorTransportationPreMLConfig.h b/src/algorithms/fardetectors/FarDetectorTransportationPreMLConfig.h new file mode 100644 index 0000000000..78c284c3ec --- /dev/null +++ b/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; + + }; +} diff --git a/src/detectors/LOWQ2/CMakeLists.txt b/src/detectors/LOWQ2/CMakeLists.txt index b114bdcbba..10cd39c860 100644 --- a/src/detectors/LOWQ2/CMakeLists.txt +++ b/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() diff --git a/src/detectors/LOWQ2/LOWQ2.cc b/src/detectors/LOWQ2/LOWQ2.cc index 1b5feb86fc..94dbf0d1df 100644 --- a/src/detectors/LOWQ2/LOWQ2.cc +++ b/src/detectors/LOWQ2/LOWQ2.cc @@ -3,6 +3,7 @@ // // +#include #include #include #include @@ -21,10 +22,17 @@ #include "factories/digi/SiliconTrackerDigi_factory.h" #include "factories/fardetectors/FarDetectorLinearProjection_factory.h" #include "factories/fardetectors/FarDetectorLinearTracking_factory.h" +#if EDM4EIC_VERSION_MAJOR >= 8 +#include "factories/fardetectors/FarDetectorTransportationPreML_factory.h" +#include "factories/fardetectors/FarDetectorTransportationPostML_factory.h" +#endif #include "factories/fardetectors/FarDetectorMLReconstruction_factory.h" #include "factories/fardetectors/FarDetectorTrackerCluster_factory.h" #include "factories/meta/CollectionCollector_factory.h" #include "factories/meta/SubDivideCollection_factory.h" +#if EDM4EIC_VERSION_MAJOR >= 8 +#include "factories/meta/ONNXInference_factory.h" +#endif extern "C" { void InitPlugin(JApplication *app) { @@ -143,6 +151,36 @@ extern "C" { app )); +#if EDM4EIC_VERSION_MAJOR >= 8 + app->Add(new JOmniFactoryGeneratorT( + "TaggerTrackerTransportationPreML", + {"TaggerTrackerProjectedTracks","MCScatteredElectrons","MCBeamElectrons"}, + {"TaggerTrackerFeatureTensor","TaggerTrackerTargetTensor"}, + { + .beamE = 10.0, + }, + app + )); + app->Add(new JOmniFactoryGeneratorT( + "TaggerTrackerTransportationInference", + {"TaggerTrackerFeatureTensor"}, + {"TaggerTrackerPredictionTensor"}, + { + .modelPath = "calibrations/onnx/TaggerTrackerTransportation.onnx", + }, + app + )); + app->Add(new JOmniFactoryGeneratorT( + "TaggerTrackerTransportationPostML", + {"TaggerTrackerPredictionTensor","MCBeamElectrons"}, + {"TaggerTrackerReconstructedParticles"}, + { + .beamE = 10.0, + }, + app + )); +#endif + // Vector reconstruction at origin app->Add(new JOmniFactoryGeneratorT( "TaggerTrackerTrajectories", diff --git a/src/factories/fardetectors/FarDetectorTransportationPostML_factory.h b/src/factories/fardetectors/FarDetectorTransportationPostML_factory.h new file mode 100644 index 0000000000..10d68680e2 --- /dev/null +++ b/src/factories/fardetectors/FarDetectorTransportationPostML_factory.h @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 - 2025, Simon Gardner + +#pragma once + +#include "algorithms/fardetectors/FarDetectorTransportationPostML.h" +#include "services/algorithms_init/AlgorithmsInit_service.h" +#include "extensions/jana/JOmniFactory.h" + + +namespace eicrecon { + +class FarDetectorTransportationPostML_factory : public JOmniFactory { + +public: + using AlgoT = eicrecon::FarDetectorTransportationPostML; +private: + std::unique_ptr m_algo; + + PodioInput m_prediction_tensor_input {this}; + PodioInput m_beamelectrons_input {this}; + + PodioOutput m_particle_output {this}; + +public: + void Configure() { + m_algo = std::make_unique(GetPrefix()); + m_algo->level(static_cast(logger()->level())); + m_algo->applyConfig(config()); + m_algo->init(); + } + + void ChangeRun(int64_t run_number) { + } + + void Process(int64_t run_number, uint64_t event_number) { + m_algo->process({m_prediction_tensor_input(), m_beamelectrons_input()}, + {m_particle_output().get()}); + } +}; + +} // eicrecon diff --git a/src/factories/fardetectors/FarDetectorTransportationPreML_factory.h b/src/factories/fardetectors/FarDetectorTransportationPreML_factory.h new file mode 100644 index 0000000000..bcb2518736 --- /dev/null +++ b/src/factories/fardetectors/FarDetectorTransportationPreML_factory.h @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 - 2025, Simon Gardner + +#pragma once + +#include "algorithms/fardetectors/FarDetectorTransportationPreML.h" +#include "services/algorithms_init/AlgorithmsInit_service.h" +#include "extensions/jana/JOmniFactory.h" + + +namespace eicrecon { + +class FarDetectorTransportationPreML_factory : public JOmniFactory { + +public: + using AlgoT = eicrecon::FarDetectorTransportationPreML; +private: + std::unique_ptr m_algo; + + PodioInput m_trackparam_input {this}; + PodioInput m_scatteredelectrons_input {this}; + PodioInput m_beamelectrons_input {this}; + + PodioOutput m_feature_tensor_output {this}; + PodioOutput m_target_tensor_output {this}; + +public: + void Configure() { + m_algo = std::make_unique(GetPrefix()); + m_algo->level(static_cast(logger()->level())); + m_algo->applyConfig(config()); + m_algo->init(); + } + + void ChangeRun(int64_t run_number) { + } + + void Process(int64_t run_number, uint64_t event_number) { + m_algo->process({m_trackparam_input(),m_scatteredelectrons_input(), m_beamelectrons_input()}, + {m_feature_tensor_output().get(), m_target_tensor_output().get()}); + } +}; + +} // eicrecon diff --git a/src/services/io/podio/JEventProcessorPODIO.cc b/src/services/io/podio/JEventProcessorPODIO.cc index 1d6cbf5d4e..1ba796cc51 100644 --- a/src/services/io/podio/JEventProcessorPODIO.cc +++ b/src/services/io/podio/JEventProcessorPODIO.cc @@ -154,6 +154,7 @@ JEventProcessorPODIO::JEventProcessorPODIO() { "TaggerTrackerTracks", "TaggerTrackerTrajectories", "TaggerTrackerTrackParameters", + "TaggerTrackerReconstructedParticles", // Forward & Far forward hits "B0TrackerRecHits",