ActsExamples/EventData/NeuralCalibrator.hpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 #include "Acts/EventData/SourceLink.hpp"
0012 #include "Acts/Utilities/CalibrationContext.hpp"
0013 #include <Acts/Plugins/Onnx/OnnxRuntimeBase.hpp>
0014 #include <ActsExamples/EventData/MeasurementCalibration.hpp>
0015
0016 #include <filesystem>
0017
0018 namespace ActsExamples {
0019
0020 class NeuralCalibrator : public MeasurementCalibrator {
0021  public:
0022   /// Measurement position calibration based on mixture density network
0023   /// (MDN) model. The model takes as input:
0024   ///
0025   /// - A 7x7 charge matrix centered on the center pixel of the cluster;
0026   /// - The volume and layer identifiers from
0027   ///   the GeometryIdentifier of the containing surface;
0028   /// - The bound phi and theta angles of the predicted track state;
0029   /// - The initial estimated position
0030   /// - The initial estimated variance
0031   ///
0032   /// Given these inputs, a mixture density network estimates
0033   /// the parameters of a gaussian mixture model:
0034   ///
0035   /// P(Y|X) = \sum_i P(Prior_i) N(Y|Mean_i(X), Variance_i(X))
0036   ///
0037   /// These are translated to single position + variance estimate by
0038   /// taking the most probable value based on the estimated priors.
0039   /// The measurements are assumed to be 2-dimensional.
0040   ///
0041   /// This class implements the MeasurementCalibrator interface, and
0042   /// therefore internally computes the network input and runs the
0043   /// inference engine itself.
0044   ///
0045   /// @param [in] modelPath The path to the .onnx model file
0046   /// @param [in] nComponent The number of components in the gaussian mixture
0047   /// @param [in] volumes The volume ids for which to apply the calibration
0048   NeuralCalibrator(const std::filesystem::path& modelPath,
0049                    std::size_t nComponents = 1,
0050                    std::vector<std::size_t> volumeIds = {7, 8, 9});
0051
0052   /// The MeasurementCalibrator interface methods
0053   void calibrate(
0054       const MeasurementContainer& measurements,
0055       const ClusterContainer* clusters, const Acts::GeometryContext& gctx,
0056       const Acts::CalibrationContext& cctx, const Acts::SourceLink& sourceLink,
0057       Acts::MultiTrajectory<Acts::VectorMultiTrajectory>::TrackStateProxy&
0058           trackState) const override;
0059
0060   bool needsClusters() const override { return true; }
0061
0062  private:
0063   Ort::Env m_env;
0064   Acts::OnnxRuntimeBase m_model;
0065   std::size_t m_nComponents;
0066   std::size_t m_nInputs =
0067       57;  // TODO make this configurable? e.g. for changing matrix size?
0068
0069   // TODO: this should probably be handled outside of the calibrator,
0070   // by setting up a GeometryHierarchyMap<MeasurementCalibrator>
0071   std::vector<std::size_t> m_volumeIds;
0072   PassThroughCalibrator m_fallback;
0073 };
0074
0075 }  // namespace ActsExamples