Python/src/GnnTrackFinding.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 "Acts/Plugins/Gnn/BoostTrackBuilding.hpp"
0010 #include "Acts/Plugins/Gnn/CudaTrackBuilding.hpp"
0011 #include "Acts/Plugins/Gnn/GnnPipeline.hpp"
0012 #include "Acts/Plugins/Gnn/ModuleMapCuda.hpp"
0013 #include "Acts/Plugins/Gnn/OnnxEdgeClassifier.hpp"
0014 #include "Acts/Plugins/Gnn/TensorRTEdgeClassifier.hpp"
0015 #include "Acts/Plugins/Gnn/TorchEdgeClassifier.hpp"
0016 #include "Acts/Plugins/Gnn/TorchMetricLearning.hpp"
0017 #include "Acts/Plugins/Gnn/TruthGraphMetricsHook.hpp"
0018 #include "Acts/Plugins/Python/Utilities.hpp"
0019 #include "ActsExamples/TrackFindingGnn/PrototracksToParameters.hpp"
0020 #include "ActsExamples/TrackFindingGnn/TrackFindingAlgorithmGnn.hpp"
0021 #include "ActsExamples/TrackFindingGnn/TrackFindingFromPrototrackAlgorithm.hpp"
0022 #include "ActsExamples/TrackFindingGnn/TruthGraphBuilder.hpp"
0023
0024 #include <memory>
0025
0026 #include <boost/preprocessor/if.hpp>
0027 #include <boost/vmd/tuple/size.hpp>
0028 #include <pybind11/functional.h>
0029 #include <pybind11/pybind11.h>
0030 #include <pybind11/stl.h>
0031
0032 #define ACTS_PYTHON_DECLARE_GNN_STAGE(algorithm, base, mod, ...)            \
0033   do {                                                                      \
0034     using namespace Acts;                                                   \
0035                                                                             \
0036     using Alg = algorithm;                                                  \
0037     using Config = Alg::Config;                                             \
0038     auto alg = py::class_<Alg, base, std::shared_ptr<Alg>>(mod, #algorithm) \
0039                    .def(py::init([](const Config &c, Logging::Level lvl) {  \
0040                           return std::make_shared<Alg>(                     \
0041                               c, getDefaultLogger(#algorithm, lvl));        \
0042                         }),                                                 \
0043                         py::arg("config"), py::arg("level"))                \
0044                    .def_property_readonly("config", &Alg::config);          \
0045                                                                             \
0046     auto c = py::class_<Config>(alg, "Config").def(py::init<>());           \
0047     BOOST_PP_IF(BOOST_VMD_IS_EMPTY(__VA_ARGS__), BOOST_PP_EMPTY(),          \
0048                 ACTS_PYTHON_STRUCT(c, __VA_ARGS__));                        \
0049   } while (0)
0050
0051 namespace py = pybind11;
0052
0053 using namespace ActsExamples;
0054 using namespace Acts;
0055 using namespace py::literals;
0056
0057 namespace Acts::Python {
0058
0059 void addGnnTrackFinding(Context &ctx) {
0060   auto [m, mex] = ctx.get("main", "examples");
0061
0062   {
0063     using C = Acts::GraphConstructionBase;
0064     auto c = py::class_<C, std::shared_ptr<C>>(mex, "GraphConstructionBase");
0065   }
0066   {
0067     using C = Acts::EdgeClassificationBase;
0068     auto c = py::class_<C, std::shared_ptr<C>>(mex, "EdgeClassificationBase");
0069   }
0070   {
0071     using C = Acts::TrackBuildingBase;
0072     auto c = py::class_<C, std::shared_ptr<C>>(mex, "TrackBuildingBase");
0073   }
0074
0075   ACTS_PYTHON_DECLARE_GNN_STAGE(BoostTrackBuilding, TrackBuildingBase, mex);
0076
0077 #ifdef ACTS_GNN_TORCH_BACKEND
0078   ACTS_PYTHON_DECLARE_GNN_STAGE(TorchMetricLearning, GraphConstructionBase, mex,
0079                                 modelPath, selectedFeatures, embeddingDim, rVal,
0080                                 knnVal, deviceID);
0081
0082   ACTS_PYTHON_DECLARE_GNN_STAGE(TorchEdgeClassifier, EdgeClassificationBase,
0083                                 mex, modelPath, selectedFeatures, cut, nChunks,
0084                                 undirected, deviceID, useEdgeFeatures);
0085 #endif
0086
0087 #ifdef ACTS_GNN_WITH_TENSORRT
0088   ACTS_PYTHON_DECLARE_GNN_STAGE(TensorRTEdgeClassifier, EdgeClassificationBase,
0089                                 mex, modelPath, selectedFeatures, cut,
0090                                 numExecutionContexts);
0091 #endif
0092
0093 #ifdef ACTS_GNN_WITH_CUDA
0094   ACTS_PYTHON_DECLARE_GNN_STAGE(CudaTrackBuilding, TrackBuildingBase, mex,
0095                                 useOneBlockImplementation, doJunctionRemoval);
0096 #endif
0097
0098 #ifdef ACTS_GNN_ONNX_BACKEND
0099   ACTS_PYTHON_DECLARE_GNN_STAGE(OnnxEdgeClassifier, EdgeClassificationBase, mex,
0100                                 modelPath, cut);
0101 #endif
0102
0103 #ifdef ACTS_GNN_WITH_MODULEMAP
0104   ACTS_PYTHON_DECLARE_GNN_STAGE(
0105       ModuleMapCuda, GraphConstructionBase, mex, moduleMapPath, rScale,
0106       phiScale, zScale, etaScale, moreParallel, gpuDevice, gpuBlocks, epsilon);
0107 #endif
0108
0109   ACTS_PYTHON_DECLARE_ALGORITHM(
0110       ActsExamples::TruthGraphBuilder, mex, "TruthGraphBuilder",
0111       inputSpacePoints, inputSimHits, inputParticles,
0112       inputMeasurementSimHitsMap, inputMeasurementParticlesMap, outputGraph,
0113       targetMinPT, targetMinSize, uniqueModules);
0114
0115   {
0116     auto nodeFeatureEnum =
0117         py::enum_<TrackFindingAlgorithmGnn::NodeFeature>(mex, "NodeFeature")
0118             .value("R", TrackFindingAlgorithmGnn::NodeFeature::eR)
0119             .value("Phi", TrackFindingAlgorithmGnn::NodeFeature::ePhi)
0120             .value("Z", TrackFindingAlgorithmGnn::NodeFeature::eZ)
0121             .value("X", TrackFindingAlgorithmGnn::NodeFeature::eX)
0122             .value("Y", TrackFindingAlgorithmGnn::NodeFeature::eY)
0123             .value("Eta", TrackFindingAlgorithmGnn::NodeFeature::eEta)
0124             .value("ClusterX",
0125                    TrackFindingAlgorithmGnn::NodeFeature::eClusterLoc0)
0126             .value("ClusterY",
0127                    TrackFindingAlgorithmGnn::NodeFeature::eClusterLoc1)
0128             .value("CellCount",
0129                    TrackFindingAlgorithmGnn::NodeFeature::eCellCount)
0130             .value("ChargeSum",
0131                    TrackFindingAlgorithmGnn::NodeFeature::eChargeSum);
0132
0133     // clang-format off
0134 #define ADD_FEATURE_ENUMS(n) \
0135   nodeFeatureEnum \
0136     .value("Cluster" #n "X", TrackFindingAlgorithmGnn::NodeFeature::eCluster##n##X) \
0137     .value("Cluster" #n "Y", TrackFindingAlgorithmGnn::NodeFeature::eCluster##n##Y) \
0138     .value("Cluster" #n "Z", TrackFindingAlgorithmGnn::NodeFeature::eCluster##n##Z) \
0139     .value("Cluster" #n "R", TrackFindingAlgorithmGnn::NodeFeature::eCluster##n##R) \
0140     .value("Cluster" #n "Phi", TrackFindingAlgorithmGnn::NodeFeature::eCluster##n##Phi) \
0141     .value("Cluster" #n "Eta", TrackFindingAlgorithmGnn::NodeFeature::eCluster##n##Eta) \
0142     .value("CellCount" #n, TrackFindingAlgorithmGnn::NodeFeature::eCellCount##n) \
0143     .value("ChargeSum" #n, TrackFindingAlgorithmGnn::NodeFeature::eChargeSum##n) \
0144     .value("LocEta" #n, TrackFindingAlgorithmGnn::NodeFeature::eLocEta##n) \
0145     .value("LocPhi" #n, TrackFindingAlgorithmGnn::NodeFeature::eLocPhi##n) \
0146     .value("LocDir0" #n, TrackFindingAlgorithmGnn::NodeFeature::eLocDir0##n) \
0147     .value("LocDir1" #n, TrackFindingAlgorithmGnn::NodeFeature::eLocDir1##n) \
0148     .value("LocDir2" #n, TrackFindingAlgorithmGnn::NodeFeature::eLocDir2##n) \
0149     .value("LengthDir0" #n, TrackFindingAlgorithmGnn::NodeFeature::eLengthDir0##n) \
0150     .value("LengthDir1" #n, TrackFindingAlgorithmGnn::NodeFeature::eLengthDir1##n) \
0151     .value("LengthDir2" #n, TrackFindingAlgorithmGnn::NodeFeature::eLengthDir2##n) \
0152     .value("GlobEta" #n, TrackFindingAlgorithmGnn::NodeFeature::eGlobEta##n) \
0153     .value("GlobPhi" #n, TrackFindingAlgorithmGnn::NodeFeature::eGlobPhi##n) \
0154     .value("EtaAngle" #n, TrackFindingAlgorithmGnn::NodeFeature::eEtaAngle##n) \
0155     .value("PhiAngle" #n, TrackFindingAlgorithmGnn::NodeFeature::ePhiAngle##n)
0156     // clang-format on
0157
0158     ADD_FEATURE_ENUMS(1);
0159     ADD_FEATURE_ENUMS(2);
0160
0161 #undef ADD_FEATURE_ENUMS
0162   }
0163
0164   ACTS_PYTHON_DECLARE_ALGORITHM(
0165       ActsExamples::TrackFindingAlgorithmGnn, mex, "TrackFindingAlgorithmGnn",
0166       inputSpacePoints, inputClusters, inputTruthGraph, outputProtoTracks,
0167       outputGraph, graphConstructor, edgeClassifiers, trackBuilder,
0168       nodeFeatures, featureScales, minMeasurementsPerTrack, geometryIdMap);
0169
0170   {
0171     auto cls = py::class_<Acts::GnnHook, std::shared_ptr<Acts::GnnHook>>(
0172         mex, "GnnHook");
0173   }
0174
0175   {
0176     using Class = Acts::TruthGraphMetricsHook;
0177
0178     auto cls = py::class_<Class, Acts::GnnHook, std::shared_ptr<Class>>(
0179                    mex, "TruthGraphMetricsHook")
0180                    .def(py::init([](const std::vector<std::int64_t> &g,
0181                                     Logging::Level lvl) {
0182                      return std::make_shared<Class>(
0183                          g, getDefaultLogger("TruthGraphHook", lvl));
0184                    }));
0185   }
0186
0187   {
0188     auto cls =
0189         py::class_<Acts::Device>(mex, "Device")
0190             .def_static("Cpu", &Acts::Device::Cpu)
0191             .def_static("Cuda", &Acts::Device::Cuda, py::arg("index") = 0);
0192   }
0193
0194   {
0195     using Class = Acts::GnnPipeline;
0196
0197     auto cls =
0198         py::class_<Class, std::shared_ptr<Class>>(mex, "GnnPipeline")
0199             .def(py::init(
0200                      [](std::shared_ptr<GraphConstructionBase> g,
0201                         std::vector<std::shared_ptr<EdgeClassificationBase>> e,
0202                         std::shared_ptr<TrackBuildingBase> t,
0203                         Logging::Level lvl) {
0204                        return std::make_shared<Class>(
0205                            g, e, t, getDefaultLogger("MetricLearning", lvl));
0206                      }),
0207                  py::arg("graphConstructor"), py::arg("edgeClassifiers"),
0208                  py::arg("trackBuilder"), py::arg("level"))
0209             .def("run", &GnnPipeline::run, py::arg("features"),
0210                  py::arg("moduleIds"), py::arg("spacepoints"),
0211                  py::arg("device") = Acts::Device::Cuda(0),
0212                  py::arg("hook") = Acts::GnnHook{},
0213                  py::arg("timing") = nullptr);
0214   }
0215
0216   ACTS_PYTHON_DECLARE_ALGORITHM(
0217       ActsExamples::PrototracksToParameters, mex, "PrototracksToParameters",
0218       inputProtoTracks, inputSpacePoints, outputSeeds, outputParameters,
0219       outputProtoTracks, geometry, magneticField, buildTightSeeds);
0220
0221   ACTS_PYTHON_DECLARE_ALGORITHM(
0222       ActsExamples::TrackFindingFromPrototrackAlgorithm, mex,
0223       "TrackFindingFromPrototrackAlgorithm", inputProtoTracks,
0224       inputMeasurements, inputInitialTrackParameters, outputTracks,
0225       measurementSelectorCfg, trackingGeometry, magneticField, findTracks, tag);
0226 }
0227
0228 }  // namespace Acts::Python