Python/src/TruthTracking.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2021 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/Plugins/Python/Utilities.hpp"
0010 #include "Acts/Utilities/Logger.hpp"
0011 #include "Acts/Utilities/TypeTraits.hpp"
0012 #include "ActsExamples/TruthTracking/ParticleSelector.hpp"
0013 #include "ActsExamples/TruthTracking/ParticleSmearing.hpp"
0014 #include "ActsExamples/TruthTracking/ProtoTrackTruthMatcher.hpp"
0015 #include "ActsExamples/TruthTracking/TrackModifier.hpp"
0016 #include "ActsExamples/TruthTracking/TrackParameterSelector.hpp"
0017 #include "ActsExamples/TruthTracking/TrackTruthMatcher.hpp"
0018 #include "ActsExamples/TruthTracking/TruthSeedSelector.hpp"
0019 #include "ActsExamples/TruthTracking/TruthSeedingAlgorithm.hpp"
0020 #include "ActsExamples/TruthTracking/TruthTrackFinder.hpp"
0021 #include "ActsExamples/TruthTracking/TruthVertexFinder.hpp"
0022 #include "ActsExamples/Utilities/HitSelector.hpp"
0023 #include "ActsExamples/Utilities/Range.hpp"
0024
0025 #include <array>
0026 #include <cstddef>
0027 #include <memory>
0028
0029 #include <pybind11/pybind11.h>
0030 #include <pybind11/stl.h>
0031
0032 namespace ActsExamples {
0033 class IAlgorithm;
0034 }  // namespace ActsExamples
0035
0036 namespace py = pybind11;
0037
0038 using namespace ActsExamples;
0039 using namespace Acts;
0040
0041 namespace Acts::Python {
0042
0043 void addTruthTracking(Context& ctx) {
0044   auto mex = ctx.get("examples");
0045
0046   ACTS_PYTHON_DECLARE_ALGORITHM(
0047       ActsExamples::TruthTrackFinder, mex, "TruthTrackFinder", inputParticles,
0048       inputMeasurementParticlesMap, outputProtoTracks);
0049
0050   {
0051     using Alg = ActsExamples::TruthSeedSelector;
0052     using Config = Alg::Config;
0053
0054     auto alg = py::class_<Alg, IAlgorithm, std::shared_ptr<Alg>>(
0055                    mex, "TruthSeedSelector")
0056                    .def(py::init<const Alg::Config&, Acts::Logging::Level>(),
0057                         py::arg("config"), py::arg("level"))
0058                    .def_property_readonly("config", &Alg::config);
0059
0060     auto c = py::class_<Config>(alg, "Config").def(py::init<>());
0061
0062     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0063     ACTS_PYTHON_MEMBER(inputParticles);
0064     ACTS_PYTHON_MEMBER(inputMeasurementParticlesMap);
0065     ACTS_PYTHON_MEMBER(outputParticles);
0066     ACTS_PYTHON_MEMBER(rhoMin);
0067     ACTS_PYTHON_MEMBER(rhoMax);
0068     ACTS_PYTHON_MEMBER(zMin);
0069     ACTS_PYTHON_MEMBER(zMax);
0070     ACTS_PYTHON_MEMBER(phiMin);
0071     ACTS_PYTHON_MEMBER(phiMax);
0072     ACTS_PYTHON_MEMBER(etaMin);
0073     ACTS_PYTHON_MEMBER(etaMax);
0074     ACTS_PYTHON_MEMBER(absEtaMin);
0075     ACTS_PYTHON_MEMBER(absEtaMax);
0076     ACTS_PYTHON_MEMBER(ptMin);
0077     ACTS_PYTHON_MEMBER(ptMax);
0078     ACTS_PYTHON_MEMBER(keepNeutral);
0079     ACTS_PYTHON_MEMBER(nHitsMin);
0080     ACTS_PYTHON_MEMBER(nHitsMax);
0081     ACTS_PYTHON_STRUCT_END();
0082
0083     pythonRangeProperty(c, "rho", &Config::rhoMin, &Config::rhoMax);
0084     pythonRangeProperty(c, "z", &Config::zMin, &Config::zMax);
0085     pythonRangeProperty(c, "phi", &Config::phiMin, &Config::phiMax);
0086     pythonRangeProperty(c, "eta", &Config::etaMin, &Config::etaMax);
0087     pythonRangeProperty(c, "absEta", &Config::absEtaMin, &Config::absEtaMax);
0088     pythonRangeProperty(c, "pt", &Config::ptMin, &Config::ptMax);
0089     pythonRangeProperty(c, "nHits", &Config::nHitsMin, &Config::nHitsMax);
0090   }
0091
0092   ACTS_PYTHON_DECLARE_ALGORITHM(
0093       ActsExamples::ParticleSmearing, mex, "ParticleSmearing", inputParticles,
0094       outputTrackParameters, sigmaD0, sigmaD0PtA, sigmaD0PtB, sigmaZ0,
0095       sigmaZ0PtA, sigmaZ0PtB, sigmaT0, sigmaPhi, sigmaTheta, sigmaPRel,
0096       initialSigmas, initialVarInflation, particleHypothesis, randomNumbers);
0097
0098   {
0099     using Alg = ActsExamples::ParticleSelector;
0100     using Config = Alg::Config;
0101
0102     auto alg = py::class_<Alg, IAlgorithm, std::shared_ptr<Alg>>(
0103                    mex, "ParticleSelector")
0104                    .def(py::init<const Alg::Config&, Acts::Logging::Level>(),
0105                         py::arg("config"), py::arg("level"))
0106                    .def_property_readonly("config", &Alg::config);
0107
0108     auto c = py::class_<Config>(alg, "Config").def(py::init<>());
0109
0110     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0111     ACTS_PYTHON_MEMBER(inputParticles);
0112     ACTS_PYTHON_MEMBER(inputMeasurementParticlesMap);
0113     ACTS_PYTHON_MEMBER(outputParticles);
0114     ACTS_PYTHON_MEMBER(rhoMin);
0115     ACTS_PYTHON_MEMBER(rhoMax);
0116     ACTS_PYTHON_MEMBER(absZMin);
0117     ACTS_PYTHON_MEMBER(absZMax);
0118     ACTS_PYTHON_MEMBER(timeMin);
0119     ACTS_PYTHON_MEMBER(timeMax);
0120     ACTS_PYTHON_MEMBER(phiMin);
0121     ACTS_PYTHON_MEMBER(phiMax);
0122     ACTS_PYTHON_MEMBER(etaMin);
0123     ACTS_PYTHON_MEMBER(etaMax);
0124     ACTS_PYTHON_MEMBER(absEtaMin);
0125     ACTS_PYTHON_MEMBER(absEtaMax);
0126     ACTS_PYTHON_MEMBER(mMin);
0127     ACTS_PYTHON_MEMBER(mMax);
0128     ACTS_PYTHON_MEMBER(ptMin);
0129     ACTS_PYTHON_MEMBER(ptMax);
0130     ACTS_PYTHON_MEMBER(measurementsMin);
0131     ACTS_PYTHON_MEMBER(measurementsMax);
0132     ACTS_PYTHON_MEMBER(removeCharged);
0133     ACTS_PYTHON_MEMBER(removeNeutral);
0134     ACTS_PYTHON_MEMBER(removeSecondaries);
0135     ACTS_PYTHON_STRUCT_END();
0136
0137     pythonRangeProperty(c, "rho", &Config::rhoMin, &Config::rhoMax);
0138     pythonRangeProperty(c, "absZ", &Config::absZMin, &Config::absZMax);
0139     pythonRangeProperty(c, "time", &Config::timeMin, &Config::timeMax);
0140     pythonRangeProperty(c, "phi", &Config::phiMin, &Config::phiMax);
0141     pythonRangeProperty(c, "eta", &Config::etaMin, &Config::etaMax);
0142     pythonRangeProperty(c, "absEta", &Config::absEtaMin, &Config::absEtaMax);
0143     pythonRangeProperty(c, "m", &Config::mMin, &Config::mMax);
0144     pythonRangeProperty(c, "pt", &Config::ptMin, &Config::ptMax);
0145     pythonRangeProperty(c, "measurements", &Config::measurementsMin,
0146                         &Config::measurementsMax);
0147   }
0148
0149   {
0150     using Alg = ActsExamples::TrackParameterSelector;
0151     using Config = Alg::Config;
0152
0153     auto alg = py::class_<Alg, IAlgorithm, std::shared_ptr<Alg>>(
0154                    mex, "TrackParameterSelector")
0155                    .def(py::init<const Alg::Config&, Acts::Logging::Level>(),
0156                         py::arg("config"), py::arg("level"))
0157                    .def_property_readonly("config", &Alg::config);
0158
0159     auto c = py::class_<Config>(alg, "Config").def(py::init<>());
0160
0161     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0162     ACTS_PYTHON_MEMBER(inputTrackParameters);
0163     ACTS_PYTHON_MEMBER(outputTrackParameters);
0164     ACTS_PYTHON_MEMBER(loc0Min);
0165     ACTS_PYTHON_MEMBER(loc0Max);
0166     ACTS_PYTHON_MEMBER(loc1Min);
0167     ACTS_PYTHON_MEMBER(loc1Max);
0168     ACTS_PYTHON_MEMBER(timeMin);
0169     ACTS_PYTHON_MEMBER(timeMax);
0170     ACTS_PYTHON_MEMBER(phiMin);
0171     ACTS_PYTHON_MEMBER(phiMax);
0172     ACTS_PYTHON_MEMBER(etaMin);
0173     ACTS_PYTHON_MEMBER(etaMax);
0174     ACTS_PYTHON_MEMBER(absEtaMin);
0175     ACTS_PYTHON_MEMBER(absEtaMax);
0176     ACTS_PYTHON_MEMBER(ptMin);
0177     ACTS_PYTHON_MEMBER(ptMax);
0178     ACTS_PYTHON_STRUCT_END();
0179
0180     pythonRangeProperty(c, "loc0", &Config::loc0Min, &Config::loc0Max);
0181     pythonRangeProperty(c, "loc1", &Config::loc1Min, &Config::loc1Max);
0182     pythonRangeProperty(c, "time", &Config::timeMin, &Config::timeMax);
0183     pythonRangeProperty(c, "phi", &Config::phiMin, &Config::phiMax);
0184     pythonRangeProperty(c, "eta", &Config::etaMin, &Config::etaMax);
0185     pythonRangeProperty(c, "absEta", &Config::absEtaMin, &Config::absEtaMax);
0186     pythonRangeProperty(c, "pt", &Config::ptMin, &Config::ptMax);
0187   }
0188
0189   ACTS_PYTHON_DECLARE_ALGORITHM(
0190       ActsExamples::TruthVertexFinder, mex, "TruthVertexFinder", inputTracks,
0191       inputParticles, inputMeasurementParticlesMap, outputProtoVertices,
0192       excludeSecondaries, separateSecondaries, trackMatchingRatio);
0193
0194   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::TrackModifier, mex,
0195                                 "TrackModifier", inputTracks, outputTracks,
0196                                 dropCovariance, covScale, killTime);
0197
0198   ACTS_PYTHON_DECLARE_ALGORITHM(
0199       ActsExamples::TruthSeedingAlgorithm, mex, "TruthSeedingAlgorithm",
0200       inputParticles, inputMeasurementParticlesMap, inputSpacePoints,
0201       outputParticles, outputSeeds, outputProtoTracks, deltaRMin, deltaRMax);
0202
0203   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::HitSelector, mex, "HitSelector",
0204                                 inputHits, outputHits, maxTime);
0205
0206   ACTS_PYTHON_DECLARE_ALGORITHM(
0207       ActsExamples::TrackTruthMatcher, mex, "TrackTruthMatcher", inputTracks,
0208       inputParticles, inputMeasurementParticlesMap, outputTrackParticleMatching,
0209       outputParticleTrackMatching, matchingRatio, doubleMatching);
0210
0211   ACTS_PYTHON_DECLARE_ALGORITHM(
0212       ActsExamples::ProtoTrackTruthMatcher, mex, "ProtoTrackTruthMatcher",
0213       inputProtoTracks, inputParticles, inputMeasurementParticlesMap,
0214       outputProtoTrackParticleMatching, outputParticleProtoTrackMatching,
0215       matchingRatio, doubleMatching);
0216 }
0217
0218 }  // namespace Acts::Python