Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-01-18 09:12:04

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