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, initialSigmaPtRel,
0054       initialVarInflation, particleHypothesis, randomNumbers);
0055
0056   {
0057     using Alg = ActsExamples::ParticleSelector;
0058     using Config = Alg::Config;
0059
0060     auto alg = py::class_<Alg, IAlgorithm, std::shared_ptr<Alg>>(
0061                    mex, "ParticleSelector")
0062                    .def(py::init<const Alg::Config&, Acts::Logging::Level>(),
0063                         py::arg("config"), py::arg("level"))
0064                    .def_property_readonly("config", &Alg::config);
0065
0066     auto c = py::class_<Config>(alg, "Config").def(py::init<>());
0067
0068     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0069     ACTS_PYTHON_MEMBER(inputParticles);
0070     ACTS_PYTHON_MEMBER(inputParticleMeasurementsMap);
0071     ACTS_PYTHON_MEMBER(outputParticles);
0072     ACTS_PYTHON_MEMBER(rhoMin);
0073     ACTS_PYTHON_MEMBER(rhoMax);
0074     ACTS_PYTHON_MEMBER(absZMin);
0075     ACTS_PYTHON_MEMBER(absZMax);
0076     ACTS_PYTHON_MEMBER(timeMin);
0077     ACTS_PYTHON_MEMBER(timeMax);
0078     ACTS_PYTHON_MEMBER(phiMin);
0079     ACTS_PYTHON_MEMBER(phiMax);
0080     ACTS_PYTHON_MEMBER(etaMin);
0081     ACTS_PYTHON_MEMBER(etaMax);
0082     ACTS_PYTHON_MEMBER(absEtaMin);
0083     ACTS_PYTHON_MEMBER(absEtaMax);
0084     ACTS_PYTHON_MEMBER(mMin);
0085     ACTS_PYTHON_MEMBER(mMax);
0086     ACTS_PYTHON_MEMBER(ptMin);
0087     ACTS_PYTHON_MEMBER(ptMax);
0088     ACTS_PYTHON_MEMBER(hitsMin);
0089     ACTS_PYTHON_MEMBER(hitsMax);
0090     ACTS_PYTHON_MEMBER(measurementsMin);
0091     ACTS_PYTHON_MEMBER(measurementsMax);
0092     ACTS_PYTHON_MEMBER(removeCharged);
0093     ACTS_PYTHON_MEMBER(removeNeutral);
0094     ACTS_PYTHON_MEMBER(removeSecondaries);
0095     ACTS_PYTHON_MEMBER(excludeAbsPdgs);
0096     ACTS_PYTHON_MEMBER(minPrimaryVertexId);
0097     ACTS_PYTHON_MEMBER(maxPrimaryVertexId);
0098     ACTS_PYTHON_STRUCT_END();
0099
0100     pythonRangeProperty(c, "rho", &Config::rhoMin, &Config::rhoMax);
0101     pythonRangeProperty(c, "absZ", &Config::absZMin, &Config::absZMax);
0102     pythonRangeProperty(c, "time", &Config::timeMin, &Config::timeMax);
0103     pythonRangeProperty(c, "phi", &Config::phiMin, &Config::phiMax);
0104     pythonRangeProperty(c, "eta", &Config::etaMin, &Config::etaMax);
0105     pythonRangeProperty(c, "absEta", &Config::absEtaMin, &Config::absEtaMax);
0106     pythonRangeProperty(c, "m", &Config::mMin, &Config::mMax);
0107     pythonRangeProperty(c, "pt", &Config::ptMin, &Config::ptMax);
0108     pythonRangeProperty(c, "measurements", &Config::measurementsMin,
0109                         &Config::measurementsMax);
0110     pythonRangeProperty(c, "primaryVertexId", &Config::minPrimaryVertexId,
0111                         &Config::maxPrimaryVertexId);
0112   }
0113
0114   {
0115     using Alg = ActsExamples::TrackParameterSelector;
0116     using Config = Alg::Config;
0117
0118     auto alg = py::class_<Alg, IAlgorithm, std::shared_ptr<Alg>>(
0119                    mex, "TrackParameterSelector")
0120                    .def(py::init<const Alg::Config&, Acts::Logging::Level>(),
0121                         py::arg("config"), py::arg("level"))
0122                    .def_property_readonly("config", &Alg::config);
0123
0124     auto c = py::class_<Config>(alg, "Config").def(py::init<>());
0125
0126     ACTS_PYTHON_STRUCT_BEGIN(c, Config);
0127     ACTS_PYTHON_MEMBER(inputTrackParameters);
0128     ACTS_PYTHON_MEMBER(outputTrackParameters);
0129     ACTS_PYTHON_MEMBER(loc0Min);
0130     ACTS_PYTHON_MEMBER(loc0Max);
0131     ACTS_PYTHON_MEMBER(loc1Min);
0132     ACTS_PYTHON_MEMBER(loc1Max);
0133     ACTS_PYTHON_MEMBER(timeMin);
0134     ACTS_PYTHON_MEMBER(timeMax);
0135     ACTS_PYTHON_MEMBER(phiMin);
0136     ACTS_PYTHON_MEMBER(phiMax);
0137     ACTS_PYTHON_MEMBER(etaMin);
0138     ACTS_PYTHON_MEMBER(etaMax);
0139     ACTS_PYTHON_MEMBER(absEtaMin);
0140     ACTS_PYTHON_MEMBER(absEtaMax);
0141     ACTS_PYTHON_MEMBER(ptMin);
0142     ACTS_PYTHON_MEMBER(ptMax);
0143     ACTS_PYTHON_STRUCT_END();
0144
0145     pythonRangeProperty(c, "loc0", &Config::loc0Min, &Config::loc0Max);
0146     pythonRangeProperty(c, "loc1", &Config::loc1Min, &Config::loc1Max);
0147     pythonRangeProperty(c, "time", &Config::timeMin, &Config::timeMax);
0148     pythonRangeProperty(c, "phi", &Config::phiMin, &Config::phiMax);
0149     pythonRangeProperty(c, "eta", &Config::etaMin, &Config::etaMax);
0150     pythonRangeProperty(c, "absEta", &Config::absEtaMin, &Config::absEtaMax);
0151     pythonRangeProperty(c, "pt", &Config::ptMin, &Config::ptMax);
0152   }
0153
0154   ACTS_PYTHON_DECLARE_ALGORITHM(
0155       ActsExamples::TruthVertexFinder, mex, "TruthVertexFinder", inputTracks,
0156       inputParticles, inputMeasurementParticlesMap, outputProtoVertices,
0157       excludeSecondaries, separateSecondaries, trackMatchingRatio);
0158
0159   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::TrackModifier, mex,
0160                                 "TrackModifier", inputTracks, outputTracks,
0161                                 dropCovariance, covScale, killTime);
0162
0163   ACTS_PYTHON_DECLARE_ALGORITHM(
0164       ActsExamples::TruthSeedingAlgorithm, mex, "TruthSeedingAlgorithm",
0165       inputParticles, inputParticleMeasurementsMap, inputSpacePoints,
0166       outputParticles, outputSeeds, outputProtoTracks, deltaRMin, deltaRMax);
0167
0168   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::HitSelector, mex, "HitSelector",
0169                                 inputHits, inputParticlesSelected, outputHits,
0170                                 minX, maxX, minY, maxY, minZ, maxZ, minR, maxR,
0171                                 minTime, maxTime, minEnergyLoss, maxEnergyLoss,
0172                                 minPrimaryVertexId, maxPrimaryVertexId);
0173
0174   ACTS_PYTHON_DECLARE_ALGORITHM(
0175       ActsExamples::TrackTruthMatcher, mex, "TrackTruthMatcher", inputTracks,
0176       inputParticles, inputMeasurementParticlesMap, outputTrackParticleMatching,
0177       outputParticleTrackMatching, matchingRatio, doubleMatching);
0178 }
0179
0180 }  // namespace Acts::Python