NNClusterCalibratorToolImpl.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef ACTSTRACKRECONSTRUCTION_NNCLUSTERCALIBRATORTOOLIMPL_H
#define ACTSTRACKRECONSTRUCTION_NNCLUSTERCALIBRATORTOOLIMPL_H
 
#include "ActsToolInterfaces/IPixelOnTrackCalibratorTool.h"
#include "AnalogueClusteringToolImpl.h"
#include "AthenaBaseComps/AthAlgTool.h"
#include "InDetCondTools/ISiLorentzAngleTool.h"
#include "InDetIdentifier/PixelID.h"
#include "InDetReadoutGeometry/SiDetectorElement.h"
#include "PixelClusterCalibrationToolBase.h"
#include "PixelConditionsData/ITkPixelOfflineCalibData.h"
#include "SiClusterizationTool/OnnxNNCollection.h"
#include "StoreGate/ReadCondHandleKey.h"
#include "src/detail/NNPixelClusterCalibratorHelpers.h"
 
namespace ActsTrk::detail {
 
template <typename calib_data_t, typename traj_t>
class NNClusterCalibrator;

template <typename calib_data_t, typename traj_t>
struct NNClusterCalibratorOptions {
  std::size_t m_minClusterSizeForNN = 0;  
  double m_minClusterChargeForNN = 15000.;
  const OnnxNNCollection* m_models =
      nullptr;  
  std::unique_ptr<AnalogueClusteringCalibrator<calib_data_t, traj_t>>
      m_fallback;
  bool isValid() const { return m_models != nullptr and m_fallback; }
};

template <typename calib_data_t, typename traj_t>
class NNClusterCalibrator
    : public PixelClusterCalibratorBase<
          NNClusterCalibrator<calib_data_t, traj_t>, traj_t> {
 public:
  using BASE =
      PixelClusterCalibratorBase<NNClusterCalibrator<calib_data_t, traj_t>,
                                 traj_t>;
  friend BASE;
  // using BaseOptions = BASE::Options;
  using Options = NNClusterCalibratorOptions<calib_data_t, traj_t>;
 
  NNClusterCalibrator(PixelClusterCalibratorOptionsBase&& base_options,
                      Options&& options)
      : BASE(std::move(base_options)), m_options(std::move(options)) {}
 
 protected:
  std::tuple<typename NNClusterCalibrator<calib_data_t, traj_t>::BASE::Pos,
             typename NNClusterCalibrator<calib_data_t, traj_t>::BASE::Cov,
             unsigned int>
  calibrate(const EventContext& ctx, const Acts::GeometryContext& gctx,
            const Acts::CalibrationContext& cctx,
            const xAOD::PixelCluster& cluster,
            const InDetDD::SiDetectorElement& detElement,
            const std::pair<float, float>& angles,
            const Acts::Vector2& predicted_local_position) const;
 
 private:
  std::tuple<int, int, int, int> clusterIndexRanges(
      const xAOD::PixelCluster& cluster) const;

  InDetDD::SiCellId weightedClusterCenter(
      const xAOD::PixelCluster& cluster,
      const InDetDD::PixelModuleDesign* design) const;

  void fillClusterData(NNinput& nn, const xAOD::PixelCluster& cluster,
                       const InDetDD::SiDetectorElement& detElement) const;

  std::array<float, 3> predictNumberOfClusters(NNinput& nn) const;

  PositionNNoutput predictPositions(NNinput& nn, int number) const;

  Ort::Session& selectPositionNetwork(int number) const;
 
  NNClusterCalibratorOptions<calib_data_t, traj_t> m_options;
};

template <typename calib_data_t, typename traj_t>
class NNClusterCalibratorToolImpl
    : public AnalogueClusteringToolImpl<calib_data_t, traj_t> {
 public:
  using BASE = AnalogueClusteringToolImpl<calib_data_t, traj_t>;
  using BASE::BASE;
 
  virtual StatusCode initialize() override;
 
  virtual std::unique_ptr<PixelOnBoundStateCalibratorBase> create(
      const EventContext& ctx) const override {
    return createOnTrackCalibrator(ctx);
  }
  /***
   * @brief produces configured calibrator object for this event
   *
   * The structure is such because the calibrator is used many times per event
   * and it needs event context for IOV potentially
   */
 
  virtual std::unique_ptr<PixelOnTrackCalibratorBase<traj_t>>
  createOnTrackCalibrator(const EventContext& ctx) const override {
    return std::make_unique<NNClusterCalibrator<calib_data_t, traj_t>>(
        BASE::createBaseOptions(ctx), createOptions(ctx));
  }
 
 private:

  NNClusterCalibratorOptions<calib_data_t, traj_t> createOptions(
      const EventContext& ctx) const;
 
  Gaudi::Property<std::size_t> m_minClusterSizeForNN{
      this, "minClusterSizeForNN", 0,
      "Use Analogue calibrator if size is smaller that this value"};
  
  Gaudi::Property<double> m_minClusterChargeForNN{
      this, "minClusterChargeForNN", 0,
      "Minimum total charge to enable NN clustering"};
 
  SG::ReadCondHandleKey<OnnxNNCollection> m_readKeyONNX {
      this, "NnCollectionONNXReadKey", "PixelClusterNNONNX",
      "The conditions key for ONNX-based pixel cluster NNs"};
 
  // to bo added here is tool to interface with NN inference
};
 
}  // namespace ActsTrk::detail
 
#include "src/detail/NNClusterCalibratorToolImpl.icc"
 
#endif  // ACTSTRACKRECONSTRUCTION_NNCLUSTERINGTOOLIMPL_H