Python/src/TruthTracking.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/Python/Utilities.hpp"
0010 #include "Acts/Utilities/Logger.hpp"
0011 #include "ActsExamples/TruthTracking/HitSelector.hpp"
0012 #include "ActsExamples/TruthTracking/ParticleSelector.hpp"
0013 #include "ActsExamples/TruthTracking/ParticleTrackParamExtractor.hpp"
0014 #include "ActsExamples/TruthTracking/TrackModifier.hpp"
0015 #include "ActsExamples/TruthTracking/TrackParameterSelector.hpp"
0016 #include "ActsExamples/TruthTracking/TrackParameterSmearing.hpp"
0017 #include "ActsExamples/TruthTracking/TrackTruthMatcher.hpp"
0018 #include "ActsExamples/TruthTracking/TruthSeedingAlgorithm.hpp"
0019 #include "ActsExamples/TruthTracking/TruthTrackFinder.hpp"
0020 #include "ActsExamples/TruthTracking/TruthVertexFinder.hpp"
0021
0022 #include <memory>
0023
0024 #include <pybind11/pybind11.h>
0025 #include <pybind11/stl.h>
0026
0027 namespace ActsExamples {
0028 class IAlgorithm;
0029 }  // namespace ActsExamples
0030
0031 namespace py = pybind11;
0032
0033 using namespace ActsExamples;
0034 using namespace Acts;
0035
0036 namespace Acts::Python {
0037
0038 void addTruthTracking(Context& ctx) {
0039   auto mex = ctx.get("examples");
0040
0041   ACTS_PYTHON_DECLARE_ALGORITHM(
0042       ActsExamples::TruthTrackFinder, mex, "TruthTrackFinder", inputParticles,
0043       inputParticleMeasurementsMap, outputProtoTracks);
0044
0045   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::ParticleTrackParamExtractor, mex,
0046                                 "ParticleTrackParamExtractor", inputParticles,
0047                                 outputTrackParameters);
0048
0049   ACTS_PYTHON_DECLARE_ALGORITHM(
0050       ActsExamples::TrackParameterSmearing, mex, "TrackParameterSmearing",
0051       inputTrackParameters, outputTrackParameters, sigmaLoc0, sigmaLoc0PtA,
0052       sigmaLoc0PtB, sigmaLoc1, sigmaLoc1PtA, sigmaLoc1PtB, sigmaTime, sigmaPhi,
0053       sigmaTheta, sigmaPtRel, initialSigmas, initialSigmaQoverPt,
0054       initialSigmaPtRel, initialVarInflation, particleHypothesis,
0055       randomNumbers);
0056
0057   {
0058     using Alg = ActsExamples::ParticleSelector;
0059     using Config = Alg::Config;
0060
0061     auto alg = py::class_<Alg, IAlgorithm, std::shared_ptr<Alg>>(
0062                    mex, "ParticleSelector")
0063                    .def(py::init<const Alg::Config&, Acts::Logging::Level>(),
0064                         py::arg("config"), py::arg("level"))
0065                    .def_property_readonly("config", &Alg::config);
0066
0067     {
0068       auto mc = py::class_<Alg::MeasurementCounter>(alg, "MeasurementCounter")
0069                     .def(py::init<>())
0070                     .def("addCounter", &Alg::MeasurementCounter::addCounter);
0071     }
0072
0073     auto c = py::class_<Config>(alg, "Config").def(py::init<>());
0074
0075     ACTS_PYTHON_STRUCT(
0076         c, inputParticles, inputParticleMeasurementsMap, inputMeasurements,
0077         outputParticles, rhoMin, rhoMax, absZMin, absZMax, timeMin, timeMax,
0078         phiMin, phiMax, etaMin, etaMax, absEtaMin, absEtaMax, mMin, mMax, ptMin,
0079         ptMax, hitsMin, hitsMax, measurementsMin, measurementsMax,
0080         removeCharged, removeNeutral, removeSecondaries, excludeAbsPdgs,
0081         minPrimaryVertexId, maxPrimaryVertexId, measurementCounter);
0082
0083     pythonRangeProperty(c, "rho", &Config::rhoMin, &Config::rhoMax);
0084     pythonRangeProperty(c, "absZ", &Config::absZMin, &Config::absZMax);
0085     pythonRangeProperty(c, "time", &Config::timeMin, &Config::timeMax);
0086     pythonRangeProperty(c, "phi", &Config::phiMin, &Config::phiMax);
0087     pythonRangeProperty(c, "eta", &Config::etaMin, &Config::etaMax);
0088     pythonRangeProperty(c, "absEta", &Config::absEtaMin, &Config::absEtaMax);
0089     pythonRangeProperty(c, "m", &Config::mMin, &Config::mMax);
0090     pythonRangeProperty(c, "pt", &Config::ptMin, &Config::ptMax);
0091     pythonRangeProperty(c, "measurements", &Config::measurementsMin,
0092                         &Config::measurementsMax);
0093     pythonRangeProperty(c, "hits", &Config::hitsMin, &Config::hitsMax);
0094     pythonRangeProperty(c, "primaryVertexId", &Config::minPrimaryVertexId,
0095                         &Config::maxPrimaryVertexId);
0096   }
0097
0098   {
0099     using Alg = ActsExamples::TrackParameterSelector;
0100     using Config = Alg::Config;
0101
0102     auto alg = py::class_<Alg, IAlgorithm, std::shared_ptr<Alg>>(
0103                    mex, "TrackParameterSelector")
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(c, inputTrackParameters, outputTrackParameters, loc0Min,
0111                        loc0Max, loc1Min, loc1Max, timeMin, timeMax, phiMin,
0112                        phiMax, etaMin, etaMax, absEtaMin, absEtaMax, ptMin,
0113                        ptMax);
0114
0115     pythonRangeProperty(c, "loc0", &Config::loc0Min, &Config::loc0Max);
0116     pythonRangeProperty(c, "loc1", &Config::loc1Min, &Config::loc1Max);
0117     pythonRangeProperty(c, "time", &Config::timeMin, &Config::timeMax);
0118     pythonRangeProperty(c, "phi", &Config::phiMin, &Config::phiMax);
0119     pythonRangeProperty(c, "eta", &Config::etaMin, &Config::etaMax);
0120     pythonRangeProperty(c, "absEta", &Config::absEtaMin, &Config::absEtaMax);
0121     pythonRangeProperty(c, "pt", &Config::ptMin, &Config::ptMax);
0122   }
0123
0124   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::TruthVertexFinder, mex,
0125                                 "TruthVertexFinder", inputTracks,
0126                                 inputTrackParticleMatching, outputProtoVertices,
0127                                 excludeSecondaries, separateSecondaries);
0128
0129   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::TrackModifier, mex,
0130                                 "TrackModifier", inputTracks, outputTracks,
0131                                 dropCovariance, covScale, killTime);
0132
0133   ACTS_PYTHON_DECLARE_ALGORITHM(
0134       ActsExamples::TruthSeedingAlgorithm, mex, "TruthSeedingAlgorithm",
0135       inputParticles, inputParticleMeasurementsMap, inputSpacePoints,
0136       outputParticles, outputSeeds, outputProtoTracks, deltaRMin, deltaRMax);
0137
0138   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::HitSelector, mex, "HitSelector",
0139                                 inputHits, inputParticlesSelected, outputHits,
0140                                 minX, maxX, minY, maxY, minZ, maxZ, minR, maxR,
0141                                 minTime, maxTime, minEnergyLoss, maxEnergyLoss,
0142                                 minPrimaryVertexId, maxPrimaryVertexId);
0143
0144   ACTS_PYTHON_DECLARE_ALGORITHM(
0145       ActsExamples::TrackTruthMatcher, mex, "TrackTruthMatcher", inputTracks,
0146       inputParticles, inputMeasurementParticlesMap, outputTrackParticleMatching,
0147       outputParticleTrackMatching, matchingRatio, doubleMatching);
0148 }
0149
0150 }  // namespace Acts::Python