ML/src/NeuralCalibrator.cpp

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 #include "ActsExamples/EventData/NeuralCalibrator.hpp"
0010
0011 #include "Acts/Definitions/TrackParametrization.hpp"
0012 #include "Acts/EventData/MeasurementHelpers.hpp"
0013 #include "Acts/EventData/SourceLink.hpp"
0014 #include "Acts/Utilities/CalibrationContext.hpp"
0015 #include "Acts/Utilities/Helpers.hpp"
0016 #include "Acts/Utilities/UnitVectors.hpp"
0017 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0018 #include "ActsExamples/EventData/Measurement.hpp"
0019
0020 #include <TFile.h>
0021
0022 namespace detail {
0023
0024 template <typename Array>
0025 std::size_t fillChargeMatrix(Array& arr, const ActsExamples::Cluster& cluster,
0026                              std::size_t size0 = 7u, std::size_t size1 = 7u) {
0027   // First, rescale the activations to sum to unity. This promotes
0028   // numerical stability in the index computation
0029   double totalAct = 0;
0030   for (const ActsExamples::Cluster::Cell& cell : cluster.channels) {
0031     totalAct += cell.activation;
0032   }
0033   std::vector<double> weights;
0034   for (const ActsExamples::Cluster::Cell& cell : cluster.channels) {
0035     weights.push_back(cell.activation / totalAct);
0036   }
0037
0038   double acc0 = 0;
0039   double acc1 = 0;
0040   for (std::size_t i = 0; i < cluster.channels.size(); i++) {
0041     acc0 += cluster.channels.at(i).bin[0] * weights.at(i);
0042     acc1 += cluster.channels.at(i).bin[1] * weights.at(i);
0043   }
0044
0045   // By convention, put the center pixel in the middle cell.
0046   // Achieved by translating the cluster --> compute the offsets
0047   int offset0 = static_cast<int>(acc0) - size0 / 2;
0048   int offset1 = static_cast<int>(acc1) - size1 / 2;
0049
0050   // Zero the charge matrix first, to guard against leftovers
0051   arr = Eigen::ArrayXXf::Zero(1, size0 * size1);
0052   // Fill the matrix
0053   for (const ActsExamples::Cluster::Cell& cell : cluster.channels) {
0054     // Translate each pixel
0055     int iMat = cell.bin[0] - offset0;
0056     int jMat = cell.bin[1] - offset1;
0057     if (iMat >= 0 && iMat < static_cast<int>(size0) && jMat >= 0 &&
0058         jMat < static_cast<int>(size1)) {
0059       typename Array::Index index = iMat * size0 + jMat;
0060       if (index < arr.size()) {
0061         arr(index) = cell.activation;
0062       }
0063     }
0064   }
0065   return size0 * size1;
0066 }
0067
0068 }  // namespace detail
0069
0070 ActsExamples::NeuralCalibrator::NeuralCalibrator(
0071     const std::filesystem::path& modelPath, std::size_t nComponents,
0072     std::vector<std::size_t> volumeIds)
0073     : m_env(ORT_LOGGING_LEVEL_WARNING, "NeuralCalibrator"),
0074       m_model(m_env, modelPath.c_str()),
0075       m_nComponents{nComponents},
0076       m_volumeIds{std::move(volumeIds)} {}
0077
0078 void ActsExamples::NeuralCalibrator::calibrate(
0079     const MeasurementContainer& measurements, const ClusterContainer* clusters,
0080     const Acts::GeometryContext& gctx, const Acts::CalibrationContext& cctx,
0081     const Acts::SourceLink& sourceLink,
0082     Acts::MultiTrajectory<Acts::VectorMultiTrajectory>::TrackStateProxy&
0083         trackState) const {
0084   trackState.setUncalibratedSourceLink(Acts::SourceLink{sourceLink});
0085   const IndexSourceLink& idxSourceLink = sourceLink.get<IndexSourceLink>();
0086   assert((idxSourceLink.index() < measurements.size()) and
0087          "Source link index is outside the container bounds");
0088
0089   if (!Acts::rangeContainsValue(m_volumeIds,
0090                                 idxSourceLink.geometryId().volume())) {
0091     m_fallback.calibrate(measurements, clusters, gctx, cctx, sourceLink,
0092                          trackState);
0093     return;
0094   }
0095
0096   Acts::NetworkBatchInput inputBatch(1, m_nInputs);
0097   auto input = inputBatch(0, Eigen::all);
0098
0099   // TODO: Matrix size should be configurable perhaps?
0100   std::size_t matSize0 = 7u;
0101   std::size_t matSize1 = 7u;
0102   std::size_t iInput = ::detail::fillChargeMatrix(
0103       input, (*clusters)[idxSourceLink.index()], matSize0, matSize1);
0104
0105   input[iInput++] = idxSourceLink.geometryId().volume();
0106   input[iInput++] = idxSourceLink.geometryId().layer();
0107
0108   const Acts::Surface& referenceSurface = trackState.referenceSurface();
0109   auto trackParameters = trackState.parameters();
0110
0111   const ConstVariableBoundMeasurementProxy measurement =
0112       measurements.getMeasurement(idxSourceLink.index());
0113
0114   assert(measurement.contains(Acts::eBoundLoc0) &&
0115          "Measurement does not contain the required bound loc0");
0116   assert(measurement.contains(Acts::eBoundLoc1) &&
0117          "Measurement does not contain the required bound loc1");
0118
0119   auto boundLoc0 = measurement.indexOf(Acts::eBoundLoc0);
0120   auto boundLoc1 = measurement.indexOf(Acts::eBoundLoc1);
0121
0122   Acts::Vector2 localPosition{measurement.parameters()[boundLoc0],
0123                               measurement.parameters()[boundLoc1]};
0124   Acts::Vector2 localCovariance{measurement.covariance()(boundLoc0, boundLoc0),
0125                                 measurement.covariance()(boundLoc1, boundLoc1)};
0126
0127   Acts::Vector3 dir = Acts::makeDirectionFromPhiTheta(
0128       trackParameters[Acts::eBoundPhi], trackParameters[Acts::eBoundTheta]);
0129   Acts::Vector3 globalPosition =
0130       referenceSurface.localToGlobal(gctx, localPosition, dir);
0131
0132   // Rotation matrix. When applied to global coordinates, they
0133   // are rotated into the local reference frame of the
0134   // surface. Note that this such a rotation can be found by
0135   // inverting a matrix whose columns correspond to the
0136   // coordinate axes of the local coordinate system.
0137   Acts::RotationMatrix3 rot =
0138       referenceSurface.referenceFrame(gctx, globalPosition, dir).inverse();
0139   std::pair<double, double> angles =
0140       Acts::VectorHelpers::incidentAngles(dir, rot);
0141
0142   input[iInput++] = angles.first;
0143   input[iInput++] = angles.second;
0144   input[iInput++] = localPosition[0];
0145   input[iInput++] = localPosition[1];
0146   input[iInput++] = localCovariance[0];
0147   input[iInput++] = localCovariance[1];
0148   if (iInput != m_nInputs) {
0149     throw std::runtime_error("Expected input size of " +
0150                              std::to_string(m_nInputs) +
0151                              ", got: " + std::to_string(iInput));
0152   }
0153
0154   // Input is a single row, hence .front()
0155   std::vector<float> output = m_model.runONNXInference(inputBatch).front();
0156   // Assuming 2-D measurements, the expected params structure is:
0157   // [           0,    nComponent[ --> priors
0158   // [  nComponent,  3*nComponent[ --> means
0159   // [3*nComponent,  5*nComponent[ --> variances
0160   std::size_t nParams = 5 * m_nComponents;
0161   if (output.size() != nParams) {
0162     throw std::runtime_error("Got output vector of size " +
0163                              std::to_string(output.size()) +
0164                              ", expected size " + std::to_string(nParams));
0165   }
0166
0167   // Most probable value computation of mixture density
0168   std::size_t iMax = 0;
0169   if (m_nComponents > 1) {
0170     iMax = std::distance(
0171         output.begin(),
0172         std::max_element(output.begin(), output.begin() + m_nComponents));
0173   }
0174   std::size_t iLoc0 = m_nComponents + iMax * 2;
0175   std::size_t iVar0 = 3 * m_nComponents + iMax * 2;
0176
0177   Acts::visit_measurement(measurement.size(), [&](auto N) -> void {
0178     constexpr std::size_t kMeasurementSize = decltype(N)::value;
0179     const ConstFixedBoundMeasurementProxy<kMeasurementSize> fixedMeasurement =
0180         static_cast<ConstFixedBoundMeasurementProxy<kMeasurementSize>>(
0181             measurement);
0182
0183     Acts::ActsVector<kMeasurementSize> calibratedParameters =
0184         fixedMeasurement.parameters();
0185     Acts::ActsSquareMatrix<kMeasurementSize> calibratedCovariance =
0186         fixedMeasurement.covariance();
0187
0188     calibratedParameters[boundLoc0] = output[iLoc0];
0189     calibratedParameters[boundLoc1] = output[iLoc0 + 1];
0190     calibratedCovariance(boundLoc0, boundLoc0) = output[iVar0];
0191     calibratedCovariance(boundLoc1, boundLoc1) = output[iVar0 + 1];
0192
0193     trackState.allocateCalibrated(calibratedParameters, calibratedCovariance);
0194     trackState.setProjectorSubspaceIndices(fixedMeasurement.subspaceIndices());
0195   });
0196 }