ActsPlugins/Gnn/EdgeLayerConnector.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/Utilities/Logger.hpp"
0012 #include "ActsPlugins/Gnn/Stages.hpp"
0013
0014 #include <memory>
0015
0016 namespace ActsPlugins {
0017 /// @addtogroup gnn_plugin
0018 /// @{
0019
0020 /// Track building stage based on the EdgeLayerConnector algorithm from
0021 /// the ModuleMapGraph library (CUDA based)
0022 class EdgeLayerConnector final : public TrackBuildingBase {
0023  public:
0024   /// Configuration for the EdgeLayerConnector
0025   struct Config {
0026     /// Number of thread blocks for parallel edge processing
0027     std::size_t blockSize = 512;
0028     /// Maximum number of hits allowed per track candidate
0029     std::size_t maxHitsPerTrack = 30;
0030     /// Minimum number of hits required to keep a track candidate
0031     std::size_t minHits = 3;
0032     /// Edge weight threshold below which edges are discarded
0033     float weightsCut = 0.01;
0034   };
0035
0036   /// @param cfg Configuration struct
0037   /// @param logger Logger instance
0038   EdgeLayerConnector(const Config &cfg,
0039                      std::unique_ptr<const Acts::Logger> logger)
0040       : m_cfg(cfg), m_logger(std::move(logger)) {}
0041
0042   std::vector<std::vector<int>> operator()(
0043       PipelineTensors tensors, std::vector<int> &spacepointIDs,
0044       const ExecutionContext &execContext = {}) override;
0045
0046   /// @return Read-only reference to the configuration
0047   const Config &config() const { return m_cfg; }
0048
0049  private:
0050   Config m_cfg;
0051   std::unique_ptr<const Acts::Logger> m_logger;
0052   const auto &logger() const { return *m_logger; }
0053 };
0054
0055 /// @}
0056 }  // namespace ActsPlugins