diff --git a/Examples/Framework/CMakeLists.txt b/Examples/Framework/CMakeLists.txt index ffdd78e2511..9ee5d487ddb 100644 --- a/Examples/Framework/CMakeLists.txt +++ b/Examples/Framework/CMakeLists.txt @@ -1,5 +1,8 @@ include(ActsTargetLinkLibrariesSystem) +set(ActsExamplesFramework_SOURCES) + + add_library( ActsExamplesFramework SHARED src/EventData/MeasurementCalibration.cpp @@ -19,7 +22,8 @@ add_library( src/Validation/ResPlotTool.cpp src/Validation/TrackClassification.cpp src/Validation/TrackSummaryPlotTool.cpp -) + ) + target_include_directories( ActsExamplesFramework PUBLIC $) @@ -65,3 +69,5 @@ install( install( DIRECTORY include/ActsExamples DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) + +add_subdirectory_if(ML ACTS_BUILD_PLUGIN_ONNX) diff --git a/Examples/Framework/ML/CMakeLists.txt b/Examples/Framework/ML/CMakeLists.txt new file mode 100644 index 00000000000..9ba85666fde --- /dev/null +++ b/Examples/Framework/ML/CMakeLists.txt @@ -0,0 +1,21 @@ +add_library( + ActsExamplesFrameworkML SHARED + src/NeuralCalibrator.cpp + ) + +target_include_directories( + ActsExamplesFrameworkML + PUBLIC $) + +target_link_libraries( + ActsExamplesFrameworkML + PUBLIC ActsExamplesFramework ActsPluginOnnx + ) + +install( + TARGETS ActsExamplesFrameworkML + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}) + +install( + DIRECTORY include/ActsExamples + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) diff --git a/Examples/Framework/ML/include/ActsExamples/EventData/NeuralCalibrator.hpp b/Examples/Framework/ML/include/ActsExamples/EventData/NeuralCalibrator.hpp new file mode 100644 index 00000000000..ce2664c6a76 --- /dev/null +++ b/Examples/Framework/ML/include/ActsExamples/EventData/NeuralCalibrator.hpp @@ -0,0 +1,72 @@ +// This file is part of the Acts project. +// +// Copyright (C) 2023 CERN for the benefit of the Acts project +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at http://mozilla.org/MPL/2.0/. + +#pragma once + +#include +#include + +#include + +namespace ActsExamples { + +class NeuralCalibrator : public MeasurementCalibrator { + public: + /// Measurement position calibration based on mixture density network + /// (MDN) model. The model takes as input: + /// + /// - A 7x7 charge matrix centered on the center pixel of the cluster; + /// - The volume and layer identifiers from + /// the GeometryIdentifier of the containing surface; + /// - The bound phi and theta angles of the predicted track state; + /// - The initial estimated position + /// - The initial estimated variance + /// + /// Given these inputs, a mixture density network estimates + /// the parameters of a gaussian mixture model: + /// + /// P(Y|X) = \sum_i P(Prior_i) N(Y|Mean_i(X), Variance_i(X)) + /// + /// These are translated to single position + variance estimate by + /// taking the most probable value based on the estimated priors. + /// The measurements are assumed to be 2-dimensional. + /// + /// This class implements the MeasurementCalibrator interface, and + /// therefore internally computes the network input and runs the + /// inference engine itself. + /// + /// @param [in] modelPath The path to the .onnx model file + /// @param [in] nComponent The number of components in the gaussian mixture + /// @param [in] volumes The volume ids for which to apply the calibration + NeuralCalibrator(const std::filesystem::path& modelPath, + size_t nComponents = 1, + std::vector volumeIds = {7, 8, 9}); + + /// The MeasurementCalibrator interface methods + void calibrate( + const MeasurementContainer& measurements, + const ClusterContainer* clusters, const Acts::GeometryContext& /*gctx*/, + Acts::MultiTrajectory::TrackStateProxy& + trackState) const override; + + bool needsClusters() const override { return true; } + + private: + Ort::Env m_env; + Acts::OnnxRuntimeBase m_model; + size_t m_nComponents; + size_t m_nInputs = + 57; // TODO make this configurable? e.g. for changing matrix size? + + // TODO: this should probably be handled outside of the calibrator, + // by setting up a GeometryHierarchyMap + std::vector m_volumeIds; + PassThroughCalibrator m_fallback; +}; + +} // namespace ActsExamples diff --git a/Examples/Framework/ML/src/NeuralCalibrator.cpp b/Examples/Framework/ML/src/NeuralCalibrator.cpp new file mode 100644 index 00000000000..eacaec1ccf9 --- /dev/null +++ b/Examples/Framework/ML/src/NeuralCalibrator.cpp @@ -0,0 +1,160 @@ +// This file is part of the Acts project. +// +// Copyright (C) 2023 CERN for the benefit of the Acts project +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at http://mozilla.org/MPL/2.0/. + +#include + +#include + +namespace detail { + +template +size_t fillChargeMatrix(Array& arr, const ActsExamples::Cluster& cluster, + size_t size0 = 7u, size_t size1 = 7u) { + // First, rescale the activations to sum to unity. This promotes + // numerical stability in the index computation + double totalAct = 0; + for (const ActsExamples::Cluster::Cell& cell : cluster.channels) { + totalAct += cell.activation; + } + std::vector weights; + for (const ActsExamples::Cluster::Cell& cell : cluster.channels) { + weights.push_back(cell.activation / totalAct); + } + + double acc0 = 0; + double acc1 = 0; + for (size_t i = 0; i < cluster.channels.size(); i++) { + acc0 += cluster.channels.at(i).bin[0] * weights.at(i); + acc1 += cluster.channels.at(i).bin[1] * weights.at(i); + } + + // By convention, put the center pixel in the middle cell. + // Achieved by translating the cluster --> compute the offsets + int offset0 = static_cast(acc0) - size0 / 2; + int offset1 = static_cast(acc1) - size1 / 2; + + // Zero the charge matrix first, to guard against leftovers + arr = Eigen::ArrayXXf::Zero(1, size0 * size1); + // Fill the matrix + for (const ActsExamples::Cluster::Cell& cell : cluster.channels) { + // Translate each pixel + int iMat = cell.bin[0] - offset0; + int jMat = cell.bin[1] - offset1; + if (iMat >= 0 && iMat < (int)size0 && jMat >= 0 && jMat < (int)size1) { + typename Array::Index index = iMat * size0 + jMat; + if (index < arr.size()) { + arr(index) = cell.activation; + } + } + } + return size0 * size1; +} + +} // namespace detail + +ActsExamples::NeuralCalibrator::NeuralCalibrator( + const std::filesystem::path& modelPath, size_t nComponents, + std::vector volumeIds) + : m_env(ORT_LOGGING_LEVEL_WARNING, "NeuralCalibrator"), + m_model(m_env, modelPath.c_str()), + m_nComponents{nComponents}, + m_volumeIds{std::move(volumeIds)} {} + +void ActsExamples::NeuralCalibrator::calibrate( + const MeasurementContainer& measurements, const ClusterContainer* clusters, + const Acts::GeometryContext& gctx, + Acts::MultiTrajectory::TrackStateProxy& + trackState) const { + Acts::SourceLink usl = trackState.getUncalibratedSourceLink(); + const IndexSourceLink& sourceLink = usl.get(); + assert((sourceLink.index() < measurements.size()) and + "Source link index is outside the container bounds"); + + if (std::find(m_volumeIds.begin(), m_volumeIds.end(), + sourceLink.geometryId().volume()) == m_volumeIds.end()) { + m_fallback.calibrate(measurements, clusters, gctx, trackState); + return; + } + + Acts::NetworkBatchInput inputBatch(1, m_nInputs); + auto input = inputBatch(0, Eigen::all); + + // TODO: Matrix size should be configurable perhaps? + size_t matSize0 = 7u; + size_t matSize1 = 7u; + size_t iInput = ::detail::fillChargeMatrix( + input, (*clusters)[sourceLink.index()], matSize0, matSize1); + + input[iInput++] = sourceLink.geometryId().volume(); + input[iInput++] = sourceLink.geometryId().layer(); + input[iInput++] = trackState.parameters()[Acts::eBoundPhi]; + input[iInput++] = trackState.parameters()[Acts::eBoundTheta]; + + std::visit( + [&](const auto& measurement) { + auto E = measurement.expander(); + auto P = measurement.projector(); + Acts::ActsVector fpar = E * measurement.parameters(); + Acts::ActsSymMatrix fcov = + E * measurement.covariance() * E.transpose(); + + input[iInput++] = fpar[Acts::eBoundLoc0]; + input[iInput++] = fpar[Acts::eBoundLoc1]; + input[iInput++] = fcov(Acts::eBoundLoc0, Acts::eBoundLoc0); + input[iInput++] = fcov(Acts::eBoundLoc1, Acts::eBoundLoc1); + if (iInput != m_nInputs) { + throw std::runtime_error("Expected input size of " + + std::to_string(m_nInputs) + + ", got: " + std::to_string(iInput)); + } + + // Input is a single row, hence .front() + std::vector output = + m_model.runONNXInference(inputBatch).front(); + // Assuming 2-D measurements, the expected params structure is: + // [ 0, nComponent[ --> priors + // [ nComponent, 3*nComponent[ --> means + // [3*nComponent, 5*nComponent[ --> variances + size_t nParams = 5 * m_nComponents; + if (output.size() != nParams) { + throw std::runtime_error( + "Got output vector of size " + std::to_string(output.size()) + + ", expected size " + std::to_string(nParams)); + } + + // Most probable value computation of mixture density + size_t iMax = 0; + if (m_nComponents > 1) { + iMax = std::distance( + output.begin(), + std::max_element(output.begin(), output.begin() + m_nComponents)); + } + size_t iLoc0 = m_nComponents + iMax * 2; + size_t iVar0 = 3 * m_nComponents + iMax * 2; + + fpar[Acts::eBoundLoc0] = output[iLoc0]; + fpar[Acts::eBoundLoc1] = output[iLoc0 + 1]; + fcov(Acts::eBoundLoc0, Acts::eBoundLoc0) = output[iVar0]; + fcov(Acts::eBoundLoc1, Acts::eBoundLoc1) = output[iVar0 + 1]; + + constexpr size_t kSize = + std::remove_reference_t::size(); + std::array indices = measurement.indices(); + Acts::ActsVector cpar = P * fpar; + Acts::ActsSymMatrix ccov = P * fcov * P.transpose(); + + Acts::SourceLink sl{sourceLink.geometryId(), sourceLink}; + + Acts::Measurement calibrated( + std::move(sl), indices, cpar, ccov); + + trackState.allocateCalibrated(calibrated.size()); + trackState.setCalibrated(calibrated); + }, + (measurements)[sourceLink.index()]); +} diff --git a/Examples/Python/CMakeLists.txt b/Examples/Python/CMakeLists.txt index 92c5203b8e8..2f8566629c0 100644 --- a/Examples/Python/CMakeLists.txt +++ b/Examples/Python/CMakeLists.txt @@ -167,8 +167,9 @@ else() endif() if(ACTS_BUILD_PLUGIN_ONNX) - target_link_libraries(ActsPythonBindings PUBLIC ActsExamplesTrackFindingML) + target_link_libraries(ActsPythonBindings PUBLIC ActsExamplesFrameworkML ActsExamplesTrackFindingML) target_sources(ActsPythonBindings PRIVATE src/Onnx.cpp) + target_sources(ActsPythonBindings PRIVATE src/OnnxNeuralCalibrator.cpp) list(APPEND py_files examples/onnx/__init__.py) if(ACTS_BUILD_PLUGIN_MLPACK) @@ -181,6 +182,7 @@ if(ACTS_BUILD_PLUGIN_ONNX) else() target_sources(ActsPythonBindings PRIVATE src/OnnxStub.cpp) target_sources(ActsPythonBindings PRIVATE src/OnnxMlpackStub.cpp) + target_sources(ActsPythonBindings PRIVATE src/OnnxNeuralCalibratorStub.cpp) endif() add_custom_target(ActsPythonGlueCode) diff --git a/Examples/Python/python/acts/examples/reconstruction.py b/Examples/Python/python/acts/examples/reconstruction.py index 4396174a0f2..9dc0b9cdfe4 100644 --- a/Examples/Python/python/acts/examples/reconstruction.py +++ b/Examples/Python/python/acts/examples/reconstruction.py @@ -811,8 +811,8 @@ def addKalmanTracks( inputProtoTracks: str = "truth_particle_tracks", multipleScattering: bool = True, energyLoss: bool = True, - calibrationConfigFile: str = None, clusters: str = None, + calibrator: acts.examples.MeasurementCalibrator = acts.examples.makePassThroughCalibrator(), logLevel: Optional[acts.logging.Level] = None, ) -> None: customLogLevel = acts.examples.defaultLogging(s, logLevel) @@ -836,11 +836,6 @@ def addKalmanTracks( "level": customLogLevel(), } - if calibrationConfigFile is None: - calibrator = acts.examples.makePassThroughCalibrator() - else: - calibrator = acts.examples.makeScalingCalibrator(calibrationConfigFile) - fitAlg = acts.examples.TrackFittingAlgorithm( level=customLogLevel(), inputMeasurements="measurements", diff --git a/Examples/Python/src/ModuleEntry.cpp b/Examples/Python/src/ModuleEntry.cpp index 360b9eae525..f8a54bef186 100644 --- a/Examples/Python/src/ModuleEntry.cpp +++ b/Examples/Python/src/ModuleEntry.cpp @@ -152,6 +152,7 @@ void addEDM4hep(Context& ctx); void addSvg(Context& ctx); void addOnnx(Context& ctx); void addOnnxMlpack(Context& ctx); +void addOnnxNeuralCalibrator(Context& ctx); } // namespace Acts::Python @@ -398,4 +399,5 @@ PYBIND11_MODULE(ActsPythonBindings, m) { addSvg(ctx); addOnnx(ctx); addOnnxMlpack(ctx); + addOnnxNeuralCalibrator(ctx); } diff --git a/Examples/Python/src/OnnxNeuralCalibrator.cpp b/Examples/Python/src/OnnxNeuralCalibrator.cpp new file mode 100644 index 00000000000..b6c1ac5af37 --- /dev/null +++ b/Examples/Python/src/OnnxNeuralCalibrator.cpp @@ -0,0 +1,34 @@ +// This file is part of the Acts project. +// +// Copyright (C) 2023 CERN for the benefit of the Acts project +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at http://mozilla.org/MPL/2.0/. + +#include "Acts/Plugins/Python/Utilities.hpp" +#include + +#include +#include + +namespace py = pybind11; + +using namespace ActsExamples; +using namespace Acts; + +namespace Acts::Python { + +void addOnnxNeuralCalibrator(Context &ctx) { + auto [m, mex, onnx] = ctx.get("main", "examples", "onnx"); + + onnx.def( + "makeNeuralCalibrator", + [](const char *modelPath, size_t nComp, std::vector volumeIds) + -> std::shared_ptr { + return std::make_shared(modelPath, nComp, volumeIds); + }, + py::arg("modelPath"), py::arg("nComp") = 1, + py::arg("volumeIds") = std::vector({7, 8, 9})); +} +} // namespace Acts::Python diff --git a/Examples/Python/src/OnnxNeuralCalibratorStub.cpp b/Examples/Python/src/OnnxNeuralCalibratorStub.cpp new file mode 100644 index 00000000000..bb7b38ef3ac --- /dev/null +++ b/Examples/Python/src/OnnxNeuralCalibratorStub.cpp @@ -0,0 +1,19 @@ +// This file is part of the Acts project. +// +// Copyright (C) 2023 CERN for the benefit of the Acts project +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at http://mozilla.org/MPL/2.0/. + +#include "Acts/Plugins/Python/Utilities.hpp" + +#include +#include + +namespace py = pybind11; + +namespace Acts::Python { + +void addOnnxNeuralCalibrator(Context& /*ctx*/) {} +} // namespace Acts::Python