Python/src/TrackFinding.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2021-2024 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 http://mozilla.org/MPL/2.0/.
0008
0009 #include "Acts/Geometry/GeometryHierarchyMap.hpp"
0010 #include "Acts/Geometry/GeometryIdentifier.hpp"
0011 #include "Acts/Plugins/Python/Utilities.hpp"
0012 #include "Acts/Seeding/SeedConfirmationRangeConfig.hpp"
0013 #include "Acts/Seeding/SeedFilterConfig.hpp"
0014 #include "Acts/Seeding/SeedFinderConfig.hpp"
0015 #include "Acts/Seeding/SeedFinderGbtsConfig.hpp"
0016 #include "Acts/Seeding/SeedFinderOrthogonalConfig.hpp"
0017 #include "Acts/Seeding/SpacePointGrid.hpp"
0018 #include "Acts/TrackFinding/MeasurementSelector.hpp"
0019 #include "Acts/Utilities/Logger.hpp"
0020 #include "Acts/Utilities/TypeTraits.hpp"
0021 #include "ActsExamples/EventData/Track.hpp"
0022 #include "ActsExamples/TrackFinding/GbtsSeedingAlgorithm.hpp"
0023 #include "ActsExamples/TrackFinding/HoughTransformSeeder.hpp"
0024 #include "ActsExamples/TrackFinding/MuonHoughSeeder.hpp"
0025 #include "ActsExamples/TrackFinding/SeedingAlgorithm.hpp"
0026 #include "ActsExamples/TrackFinding/SeedingOrthogonalAlgorithm.hpp"
0027 #include "ActsExamples/TrackFinding/SpacePointMaker.hpp"
0028 #include "ActsExamples/TrackFinding/TrackFindingAlgorithm.hpp"
0029 #include "ActsExamples/TrackFinding/TrackParamsEstimationAlgorithm.hpp"
0030 #include "ActsExamples/Utilities/MeasurementMapSelector.hpp"
0031 #include "ActsExamples/Utilities/PrototracksToSeeds.hpp"
0032 #include "ActsExamples/Utilities/SeedsToPrototracks.hpp"
0033 #include "ActsExamples/Utilities/TracksToParameters.hpp"
0034 #include "ActsExamples/Utilities/TracksToTrajectories.hpp"
0035 #include "ActsExamples/Utilities/TrajectoriesToPrototracks.hpp"
0036
0037 #include <array>
0038 #include <cstddef>
0039 #include <memory>
0040 #include <tuple>
0041 #include <utility>
0042 #include <vector>
0043
0044 #include <pybind11/pybind11.h>
0045 #include <pybind11/stl.h>
0046
0047 namespace Acts {
0048 class MagneticFieldProvider;
0049 class TrackingGeometry;
0050 }  // namespace Acts
0051 namespace ActsExamples {
0052 class IAlgorithm;
0053 class SimSpacePoint;
0054 }  // namespace ActsExamples
0055
0056 namespace py = pybind11;
0057
0058 using namespace ActsExamples;
0059 using namespace Acts;
0060
0061 namespace Acts::Python {
0062
0063 void addTrackFinding(Context& ctx) {
0064   auto [m, mex] = ctx.get("main", "examples");
0065
0066   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::SpacePointMaker, mex,
0067                                 "SpacePointMaker", inputSourceLinks,
0068                                 inputMeasurements, outputSpacePoints,
0069                                 trackingGeometry, geometrySelection);
0070
0071   {
0072     using Config = Acts::SeedFilterConfig;
0073     auto c = py::class_<Config>(m, "SeedFilterConfig").def(py::init<>());
0074     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0075     ACTS_PYTHON_MEMBER(deltaInvHelixDiameter);
0076     ACTS_PYTHON_MEMBER(impactWeightFactor);
0077     ACTS_PYTHON_MEMBER(zOriginWeightFactor);
0078     ACTS_PYTHON_MEMBER(compatSeedWeight);
0079     ACTS_PYTHON_MEMBER(deltaRMin);
0080     ACTS_PYTHON_MEMBER(maxSeedsPerSpM);
0081     ACTS_PYTHON_MEMBER(compatSeedLimit);
0082     ACTS_PYTHON_MEMBER(seedConfirmation);
0083     ACTS_PYTHON_MEMBER(centralSeedConfirmationRange);
0084     ACTS_PYTHON_MEMBER(forwardSeedConfirmationRange);
0085     ACTS_PYTHON_MEMBER(useDeltaRorTopRadius);
0086     ACTS_PYTHON_MEMBER(seedWeightIncrement);
0087     ACTS_PYTHON_MEMBER(numSeedIncrement);
0088     ACTS_PYTHON_MEMBER(maxSeedsPerSpMConf);
0089     ACTS_PYTHON_MEMBER(maxQualitySeedsPerSpMConf);
0090     ACTS_PYTHON_STRUCT_END();
0091     patchKwargsConstructor(c);
0092   }
0093
0094   {
0095     using Config = Acts::SeedFinderConfig<SimSpacePoint>;
0096     auto c = py::class_<Config>(m, "SeedFinderConfig").def(py::init<>());
0097     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0098     ACTS_PYTHON_MEMBER(minPt);
0099     ACTS_PYTHON_MEMBER(cotThetaMax);
0100     ACTS_PYTHON_MEMBER(deltaRMin);
0101     ACTS_PYTHON_MEMBER(deltaRMax);
0102     ACTS_PYTHON_MEMBER(deltaRMinBottomSP);
0103     ACTS_PYTHON_MEMBER(deltaRMaxBottomSP);
0104     ACTS_PYTHON_MEMBER(deltaRMinTopSP);
0105     ACTS_PYTHON_MEMBER(deltaRMaxTopSP);
0106     ACTS_PYTHON_MEMBER(impactMax);
0107     ACTS_PYTHON_MEMBER(sigmaScattering);
0108     ACTS_PYTHON_MEMBER(maxPtScattering);
0109     ACTS_PYTHON_MEMBER(maxSeedsPerSpM);
0110     ACTS_PYTHON_MEMBER(collisionRegionMin);
0111     ACTS_PYTHON_MEMBER(collisionRegionMax);
0112     ACTS_PYTHON_MEMBER(phiMin);
0113     ACTS_PYTHON_MEMBER(phiMax);
0114     ACTS_PYTHON_MEMBER(zMin);
0115     ACTS_PYTHON_MEMBER(zMax);
0116     ACTS_PYTHON_MEMBER(zOutermostLayers);
0117     ACTS_PYTHON_MEMBER(rMax);
0118     ACTS_PYTHON_MEMBER(rMin);
0119     ACTS_PYTHON_MEMBER(radLengthPerSeed);
0120     ACTS_PYTHON_MEMBER(zAlign);
0121     ACTS_PYTHON_MEMBER(rAlign);
0122     ACTS_PYTHON_MEMBER(sigmaError);
0123     ACTS_PYTHON_MEMBER(maxBlockSize);
0124     ACTS_PYTHON_MEMBER(nTrplPerSpBLimit);
0125     ACTS_PYTHON_MEMBER(nAvgTrplPerSpBLimit);
0126     ACTS_PYTHON_MEMBER(impactMax);
0127     ACTS_PYTHON_MEMBER(deltaZMax);
0128     ACTS_PYTHON_MEMBER(zBinEdges);
0129     ACTS_PYTHON_MEMBER(interactionPointCut);
0130     ACTS_PYTHON_MEMBER(zBinsCustomLooping);
0131     ACTS_PYTHON_MEMBER(useVariableMiddleSPRange);
0132     ACTS_PYTHON_MEMBER(deltaRMiddleMinSPRange);
0133     ACTS_PYTHON_MEMBER(deltaRMiddleMaxSPRange);
0134     ACTS_PYTHON_MEMBER(rRangeMiddleSP);
0135     ACTS_PYTHON_MEMBER(rMinMiddle);
0136     ACTS_PYTHON_MEMBER(rMaxMiddle);
0137     ACTS_PYTHON_MEMBER(binSizeR);
0138     ACTS_PYTHON_MEMBER(seedConfirmation);
0139     ACTS_PYTHON_MEMBER(centralSeedConfirmationRange);
0140     ACTS_PYTHON_MEMBER(forwardSeedConfirmationRange);
0141     ACTS_PYTHON_MEMBER(useDetailedDoubleMeasurementInfo);
0142     ACTS_PYTHON_STRUCT_END();
0143     patchKwargsConstructor(c);
0144   }
0145   {
0146     using seedOptions = Acts::SeedFinderOptions;
0147     auto c = py::class_<seedOptions>(m, "SeedFinderOptions").def(py::init<>());
0148     ACTS_PYTHON_STRUCT_BEGIN(c, seedOptions);
0149     ACTS_PYTHON_MEMBER(beamPos);
0150     ACTS_PYTHON_MEMBER(bFieldInZ);
0151     ACTS_PYTHON_STRUCT_END();
0152     patchKwargsConstructor(c);
0153   }
0154   {
0155     using Config = Acts::SeedFinderOrthogonalConfig<SimSpacePoint>;
0156     auto c =
0157         py::class_<Config>(m, "SeedFinderOrthogonalConfig").def(py::init<>());
0158     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0159     ACTS_PYTHON_MEMBER(minPt);
0160     ACTS_PYTHON_MEMBER(cotThetaMax);
0161     ACTS_PYTHON_MEMBER(deltaRMinBottomSP);
0162     ACTS_PYTHON_MEMBER(deltaRMaxBottomSP);
0163     ACTS_PYTHON_MEMBER(deltaRMinTopSP);
0164     ACTS_PYTHON_MEMBER(deltaRMaxTopSP);
0165     ACTS_PYTHON_MEMBER(impactMax);
0166     ACTS_PYTHON_MEMBER(deltaPhiMax);
0167     ACTS_PYTHON_MEMBER(deltaZMax);
0168     ACTS_PYTHON_MEMBER(sigmaScattering);
0169     ACTS_PYTHON_MEMBER(maxPtScattering);
0170     ACTS_PYTHON_MEMBER(maxSeedsPerSpM);
0171     ACTS_PYTHON_MEMBER(collisionRegionMin);
0172     ACTS_PYTHON_MEMBER(collisionRegionMax);
0173     ACTS_PYTHON_MEMBER(phiMin);
0174     ACTS_PYTHON_MEMBER(phiMax);
0175     ACTS_PYTHON_MEMBER(zMin);
0176     ACTS_PYTHON_MEMBER(zMax);
0177     ACTS_PYTHON_MEMBER(zOutermostLayers);
0178     ACTS_PYTHON_MEMBER(rMax);
0179     ACTS_PYTHON_MEMBER(rMin);
0180     ACTS_PYTHON_MEMBER(radLengthPerSeed);
0181     ACTS_PYTHON_MEMBER(deltaZMax);
0182     ACTS_PYTHON_MEMBER(interactionPointCut);
0183     ACTS_PYTHON_MEMBER(deltaPhiMax);
0184     ACTS_PYTHON_MEMBER(highland);
0185     ACTS_PYTHON_MEMBER(maxScatteringAngle2);
0186     ACTS_PYTHON_MEMBER(useVariableMiddleSPRange);
0187     ACTS_PYTHON_MEMBER(deltaRMiddleMinSPRange);
0188     ACTS_PYTHON_MEMBER(deltaRMiddleMaxSPRange);
0189     ACTS_PYTHON_MEMBER(rRangeMiddleSP);
0190     ACTS_PYTHON_MEMBER(rMinMiddle);
0191     ACTS_PYTHON_MEMBER(rMaxMiddle);
0192     ACTS_PYTHON_MEMBER(seedConfirmation);
0193     ACTS_PYTHON_MEMBER(centralSeedConfirmationRange);
0194     ACTS_PYTHON_MEMBER(forwardSeedConfirmationRange);
0195     ACTS_PYTHON_STRUCT_END();
0196     patchKwargsConstructor(c);
0197   }
0198
0199   {
0200     using Config = Acts::SeedFinderGbtsConfig<SimSpacePoint>;
0201     auto c = py::class_<Config>(m, "SeedFinderGbtsConfig").def(py::init<>());
0202     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0203     ACTS_PYTHON_MEMBER(minPt);
0204     ACTS_PYTHON_MEMBER(sigmaScattering);
0205     ACTS_PYTHON_MEMBER(highland);
0206     ACTS_PYTHON_MEMBER(maxScatteringAngle2);
0207     ACTS_PYTHON_MEMBER(connector_input_file);
0208     ACTS_PYTHON_MEMBER(m_phiSliceWidth);
0209     ACTS_PYTHON_MEMBER(m_nMaxPhiSlice);
0210     ACTS_PYTHON_MEMBER(m_useClusterWidth);
0211     ACTS_PYTHON_MEMBER(m_layerGeometry);
0212     ACTS_PYTHON_MEMBER(maxSeedsPerSpM);
0213     ACTS_PYTHON_STRUCT_END();
0214     patchKwargsConstructor(c);
0215   }
0216
0217   {
0218     using seedConf = Acts::SeedConfirmationRangeConfig;
0219     auto c = py::class_<seedConf>(m, "SeedConfirmationRangeConfig")
0220                  .def(py::init<>());
0221     ACTS_PYTHON_STRUCT_BEGIN(c, seedConf);
0222     ACTS_PYTHON_MEMBER(zMinSeedConf);
0223     ACTS_PYTHON_MEMBER(zMaxSeedConf);
0224     ACTS_PYTHON_MEMBER(rMaxSeedConf);
0225     ACTS_PYTHON_MEMBER(nTopForLargeR);
0226     ACTS_PYTHON_MEMBER(nTopForSmallR);
0227     ACTS_PYTHON_MEMBER(seedConfMinBottomRadius);
0228     ACTS_PYTHON_MEMBER(seedConfMaxZOrigin);
0229     ACTS_PYTHON_MEMBER(minImpactSeedConf);
0230     ACTS_PYTHON_STRUCT_END();
0231     patchKwargsConstructor(c);
0232   }
0233
0234   {
0235     using Config = Acts::CylindricalSpacePointGridConfig;
0236     auto c = py::class_<Config>(m, "SpacePointGridConfig").def(py::init<>());
0237
0238     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0239     ACTS_PYTHON_MEMBER(minPt);
0240     ACTS_PYTHON_MEMBER(rMax);
0241     ACTS_PYTHON_MEMBER(zMax);
0242     ACTS_PYTHON_MEMBER(zMin);
0243     ACTS_PYTHON_MEMBER(phiMin);
0244     ACTS_PYTHON_MEMBER(phiMax);
0245     ACTS_PYTHON_MEMBER(deltaRMax);
0246     ACTS_PYTHON_MEMBER(cotThetaMax);
0247     ACTS_PYTHON_MEMBER(phiBinDeflectionCoverage);
0248     ACTS_PYTHON_MEMBER(maxPhiBins);
0249     ACTS_PYTHON_MEMBER(impactMax);
0250     ACTS_PYTHON_MEMBER(zBinEdges);
0251     ACTS_PYTHON_STRUCT_END();
0252     patchKwargsConstructor(c);
0253   }
0254   {
0255     using Options = Acts::CylindricalSpacePointGridOptions;
0256     auto c = py::class_<Options>(m, "SpacePointGridOptions").def(py::init<>());
0257
0258     ACTS_PYTHON_STRUCT_BEGIN(c, Options);
0259     ACTS_PYTHON_MEMBER(bFieldInZ);
0260     ACTS_PYTHON_STRUCT_END();
0261     patchKwargsConstructor(c);
0262   }
0263
0264   ACTS_PYTHON_DECLARE_ALGORITHM(
0265       ActsExamples::SeedingAlgorithm, mex, "SeedingAlgorithm", inputSpacePoints,
0266       outputSeeds, seedFilterConfig, seedFinderConfig, seedFinderOptions,
0267       gridConfig, gridOptions, allowSeparateRMax, zBinNeighborsTop,
0268       zBinNeighborsBottom, numPhiNeighbors);
0269
0270   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::SeedingOrthogonalAlgorithm, mex,
0271                                 "SeedingOrthogonalAlgorithm", inputSpacePoints,
0272                                 outputSeeds, seedFilterConfig, seedFinderConfig,
0273                                 seedFinderOptions);
0274
0275   ACTS_PYTHON_DECLARE_ALGORITHM(
0276       ActsExamples::GbtsSeedingAlgorithm, mex, "GbtsSeedingAlgorithm",
0277       inputSpacePoints, outputSeeds, seedFilterConfig, seedFinderConfig,
0278       seedFinderOptions, layerMappingFile, geometrySelection, inputSourceLinks,
0279       trackingGeometry, ActsGbtsMap, fill_module_csv, inputClusters);
0280
0281   ACTS_PYTHON_DECLARE_ALGORITHM(
0282       ActsExamples::HoughTransformSeeder, mex, "HoughTransformSeeder",
0283       inputSpacePoints, outputProtoTracks, inputSourceLinks, trackingGeometry,
0284       geometrySelection, inputMeasurements, subRegions, nLayers, xMin, xMax,
0285       yMin, yMax, houghHistSize_x, houghHistSize_y, hitExtend_x, threshold,
0286       localMaxWindowSize, kA);
0287
0288   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::MuonHoughSeeder, mex,
0289                                 "MuonHoughSeeder", inSimHits, inDriftCircles);
0290
0291   ACTS_PYTHON_DECLARE_ALGORITHM(
0292       ActsExamples::TrackParamsEstimationAlgorithm, mex,
0293       "TrackParamsEstimationAlgorithm", inputSeeds, inputProtoTracks,
0294       outputTrackParameters, outputSeeds, outputProtoTracks, trackingGeometry,
0295       magneticField, bFieldMin, initialSigmas, initialSimgaQoverPCoefficients,
0296       initialVarInflation, noTimeVarInflation, particleHypothesis);
0297
0298   {
0299     using Alg = ActsExamples::TrackFindingAlgorithm;
0300     using Config = Alg::Config;
0301
0302     auto alg =
0303         py::class_<Alg, ActsExamples::IAlgorithm, std::shared_ptr<Alg>>(
0304             mex, "TrackFindingAlgorithm")
0305             .def(py::init<const Config&, Acts::Logging::Level>(),
0306                  py::arg("config"), py::arg("level"))
0307             .def_property_readonly("config", &Alg::config)
0308             .def_static("makeTrackFinderFunction",
0309                         [](std::shared_ptr<const Acts::TrackingGeometry>
0310                                trackingGeometry,
0311                            std::shared_ptr<const Acts::MagneticFieldProvider>
0312                                magneticField,
0313                            Logging::Level level) {
0314                           return Alg::makeTrackFinderFunction(
0315                               trackingGeometry, magneticField,
0316                               *Acts::getDefaultLogger("TrackFinding", level));
0317                         });
0318
0319     py::class_<Alg::TrackFinderFunction,
0320                std::shared_ptr<Alg::TrackFinderFunction>>(
0321         alg, "TrackFinderFunction");
0322
0323     auto c = py::class_<Config>(alg, "Config").def(py::init<>());
0324     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0325     ACTS_PYTHON_MEMBER(inputMeasurements);
0326     ACTS_PYTHON_MEMBER(inputSourceLinks);
0327     ACTS_PYTHON_MEMBER(inputInitialTrackParameters);
0328     ACTS_PYTHON_MEMBER(inputSeeds);
0329     ACTS_PYTHON_MEMBER(outputTracks);
0330     ACTS_PYTHON_MEMBER(trackingGeometry);
0331     ACTS_PYTHON_MEMBER(magneticField);
0332     ACTS_PYTHON_MEMBER(findTracks);
0333     ACTS_PYTHON_MEMBER(measurementSelectorCfg);
0334     ACTS_PYTHON_MEMBER(trackSelectorCfg);
0335     ACTS_PYTHON_MEMBER(maxSteps);
0336     ACTS_PYTHON_MEMBER(twoWay);
0337     ACTS_PYTHON_MEMBER(seedDeduplication);
0338     ACTS_PYTHON_MEMBER(stayOnSeed);
0339     ACTS_PYTHON_STRUCT_END();
0340   }
0341
0342   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::TrajectoriesToPrototracks, mex,
0343                                 "TrajectoriesToPrototracks", inputTrajectories,
0344                                 outputProtoTracks);
0345
0346   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::TracksToTrajectories, mex,
0347                                 "TracksToTrajectories", inputTracks,
0348                                 outputTrajectories);
0349
0350   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::TracksToParameters, mex,
0351                                 "TracksToParameters", inputTracks,
0352                                 outputTrackParameters);
0353
0354   {
0355     auto constructor = [](const std::vector<std::pair<
0356                               GeometryIdentifier,
0357                               std::tuple<std::vector<double>,
0358                                          std::vector<double>,
0359                                          std::vector<std::size_t>>>>& input) {
0360       std::vector<std::pair<GeometryIdentifier, MeasurementSelectorCuts>>
0361           converted;
0362       converted.reserve(input.size());
0363       for (const auto& [id, cuts] : input) {
0364         const auto& [bins, chi2, num] = cuts;
0365         converted.emplace_back(id, MeasurementSelectorCuts{bins, chi2, num});
0366       }
0367       return std::make_unique<MeasurementSelector::Config>(converted);
0368     };
0369
0370     py::class_<MeasurementSelectorCuts>(m, "MeasurementSelectorCuts")
0371         .def(py::init<>())
0372         .def(py::init<std::vector<double>, std::vector<double>,
0373                       std::vector<std::size_t>>())
0374         .def_readwrite("etaBins", &MeasurementSelectorCuts::etaBins)
0375         .def_readwrite("chi2CutOff", &MeasurementSelectorCuts::chi2CutOff)
0376         .def_readwrite("numMeasurementsCutOff",
0377                        &MeasurementSelectorCuts::numMeasurementsCutOff);
0378
0379     auto ms = py::class_<MeasurementSelector>(m, "MeasurementSelector");
0380     auto c =
0381         py::class_<MeasurementSelector::Config>(ms, "Config")
0382             .def(py::init<std::vector<
0383                      std::pair<GeometryIdentifier, MeasurementSelectorCuts>>>())
0384             .def(py::init(constructor));
0385   }
0386
0387   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::SeedsToPrototracks, mex,
0388                                 "SeedsToPrototracks", inputSeeds,
0389                                 outputProtoTracks);
0390
0391   ACTS_PYTHON_DECLARE_ALGORITHM(
0392       ActsExamples::PrototracksToSeeds, mex, "PrototracksToSeeds",
0393       inputProtoTracks, inputSpacePoints, outputSeeds, outputProtoTracks);
0394
0395   ACTS_PYTHON_DECLARE_ALGORITHM(
0396       ActsExamples::MeasurementMapSelector, mex, "MeasurementMapSelector",
0397       inputMeasurementParticleMap, inputSourceLinks,
0398       outputMeasurementParticleMap, geometrySelection);
0399 }
0400
0401 }  // namespace Acts::Python