Python/src/Generators.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/Definitions/Algebra.hpp"
0010 #include "Acts/Utilities/AngleHelpers.hpp"
0011 #include "Acts/Utilities/Logger.hpp"
0012 #include "ActsExamples/EventData/SimParticle.hpp"
0013 #include "ActsExamples/Generators/EventGenerator.hpp"
0014 #include "ActsExamples/Utilities/ParametricParticleGenerator.hpp"
0015 #include "ActsExamples/Utilities/VertexGenerators.hpp"
0016 #include "ActsPython/Utilities/Helpers.hpp"
0017 #include "ActsPython/Utilities/Macros.hpp"
0018
0019 #include <cstddef>
0020 #include <memory>
0021 #include <string>
0022 #include <utility>
0023
0024 #include <pybind11/pybind11.h>
0025 #include <pybind11/stl.h>
0026
0027 namespace ActsExamples {
0028 class IReader;
0029 }  // namespace ActsExamples
0030
0031 namespace py = pybind11;
0032
0033 using namespace Acts;
0034 using namespace ActsExamples;
0035
0036 namespace ActsPython {
0037
0038 void addGenerators(Context& ctx) {
0039   auto mex = ctx.get("examples");
0040
0041   {
0042     using Config = EventGenerator::Config;
0043     auto gen =
0044         py::class_<EventGenerator, IReader, std::shared_ptr<EventGenerator>>(
0045             mex, "EventGenerator")
0046             .def(py::init<const Config&, Logging::Level>(), py::arg("config"),
0047                  py::arg("level"))
0048             .def_property_readonly("config", &EventGenerator::config);
0049
0050     py::class_<PrimaryVertexPositionGenerator,
0051                std::shared_ptr<PrimaryVertexPositionGenerator>>(
0052         gen, "VertexGenerator");
0053     py::class_<ParticlesGenerator, std::shared_ptr<ParticlesGenerator>>(
0054         gen, "ParticlesGenerator");
0055     py::class_<MultiplicityGenerator, std::shared_ptr<MultiplicityGenerator>>(
0056         gen, "MultiplicityGenerator");
0057
0058     using EventGenerator = EventGenerator;
0059     using Generator = EventGenerator::Generator;
0060     py::class_<Generator>(gen, "Generator")
0061         .def(py::init<>())
0062         .def(py::init<std::shared_ptr<MultiplicityGenerator>,
0063                       std::shared_ptr<PrimaryVertexPositionGenerator>,
0064                       std::shared_ptr<ParticlesGenerator>>(),
0065              py::arg("multiplicity"), py::arg("vertex"), py::arg("particles"))
0066         .def_readwrite("multiplicity", &Generator::multiplicity)
0067         .def_readwrite("vertex", &Generator::vertex)
0068         .def_readwrite("particles", &Generator::particles);
0069
0070     auto config = py::class_<Config>(gen, "Config").def(py::init<>());
0071
0072     ACTS_PYTHON_STRUCT(config, outputEvent, generators, randomNumbers,
0073                        printListing);
0074   }
0075
0076   py::class_<GaussianPrimaryVertexPositionGenerator,
0077              PrimaryVertexPositionGenerator,
0078              std::shared_ptr<GaussianPrimaryVertexPositionGenerator>>(
0079       mex, "GaussianVertexGenerator")
0080       .def(py::init<>())
0081       .def(py::init([](const Vector4& stddev, const Vector4& mean) {
0082              GaussianPrimaryVertexPositionGenerator g;
0083              g.stddev = stddev;
0084              g.mean = mean;
0085              return g;
0086            }),
0087            py::arg("stddev"), py::arg("mean"))
0088       .def_readwrite("stddev", &GaussianPrimaryVertexPositionGenerator::stddev)
0089       .def_readwrite("mean", &GaussianPrimaryVertexPositionGenerator::mean);
0090   py::class_<GaussianDisplacedVertexPositionGenerator,
0091              PrimaryVertexPositionGenerator,
0092              std::shared_ptr<GaussianDisplacedVertexPositionGenerator>>(
0093       mex, "GaussianDisplacedVertexPositionGenerator")
0094       .def(py::init<>())
0095       .def(py::init([](double rMean, double rStdDev, double zMean,
0096                        double zStdDev, double tMean, double tStdDev) {
0097              GaussianDisplacedVertexPositionGenerator g;
0098              g.rMean = rMean;
0099              g.rStdDev = rStdDev;
0100              g.zMean = zMean;
0101              g.zStdDev = zStdDev;
0102              g.tMean = tMean;
0103              g.tStdDev = tStdDev;
0104              return g;
0105            }),
0106            py::arg("rMean"), py::arg("rStdDev"), py::arg("zMean"),
0107            py::arg("zStdDev"), py::arg("tMean"), py::arg("tStdDev"))
0108       .def_readwrite("rMean", &GaussianDisplacedVertexPositionGenerator::rMean)
0109       .def_readwrite("rStdDev",
0110                      &GaussianDisplacedVertexPositionGenerator::rStdDev)
0111       .def_readwrite("zMean", &GaussianDisplacedVertexPositionGenerator::zMean)
0112       .def_readwrite("zStdDev",
0113                      &GaussianDisplacedVertexPositionGenerator::zStdDev)
0114       .def_readwrite("tMean", &GaussianDisplacedVertexPositionGenerator::tMean)
0115       .def_readwrite("tStdDev",
0116                      &GaussianDisplacedVertexPositionGenerator::tStdDev);
0117
0118   py::class_<
0119       ActsExamples::UniformPrimaryVertexPositionGenerator,
0120       ActsExamples::PrimaryVertexPositionGenerator,
0121       std::shared_ptr<ActsExamples::UniformPrimaryVertexPositionGenerator>>(
0122       mex, "UniformVertexGenerator")
0123       .def(py::init<>())
0124       .def(py::init([](const Acts::Vector4& min, const Acts::Vector4& max) {
0125              ActsExamples::UniformPrimaryVertexPositionGenerator g;
0126              g.min = min;
0127              g.max = max;
0128              return g;
0129            }),
0130            py::arg("min"), py::arg("max"))
0131       .def_readwrite("min",
0132                      &ActsExamples::UniformPrimaryVertexPositionGenerator::min)
0133       .def_readwrite("max",
0134                      &ActsExamples::UniformPrimaryVertexPositionGenerator::max);
0135
0136   py::class_<
0137       ActsExamples::FixedPrimaryVertexPositionGenerator,
0138       ActsExamples::PrimaryVertexPositionGenerator,
0139       std::shared_ptr<ActsExamples::FixedPrimaryVertexPositionGenerator>>(
0140       mex, "FixedVertexGenerator")
0141       .def(py::init<>())
0142       .def(py::init([](const Vector4& v) {
0143              FixedPrimaryVertexPositionGenerator g;
0144              g.fixed = v;
0145              return g;
0146            }),
0147            py::arg("fixed"))
0148       .def_readwrite("fixed", &FixedPrimaryVertexPositionGenerator::fixed);
0149
0150   py::class_<SimParticle>(mex, "SimParticle");
0151   py::class_<SimParticleContainer>(mex, "SimParticleContainer");
0152
0153   {
0154     using Config = ParametricParticleGenerator::Config;
0155     auto gen = py::class_<ParametricParticleGenerator, ParticlesGenerator,
0156                           std::shared_ptr<ParametricParticleGenerator>>(
0157                    mex, "ParametricParticleGenerator")
0158                    .def(py::init<const Config&>());
0159
0160     py::class_<Config>(gen, "Config")
0161         .def(py::init<>())
0162         .def_readwrite("phiMin", &Config::phiMin)
0163         .def_readwrite("phiMax", &Config::phiMax)
0164         .def_readwrite("thetaMin", &Config::thetaMin)
0165         .def_readwrite("thetaMax", &Config::thetaMax)
0166         .def_readwrite("etaUniform", &Config::etaUniform)
0167         .def_readwrite("pMin", &Config::pMin)
0168         .def_readwrite("pMax", &Config::pMax)
0169         .def_readwrite("pTransverse", &Config::pTransverse)
0170         .def_readwrite("pLogUniform", &Config::pLogUniform)
0171         .def_readwrite("pdg", &Config::pdg)
0172         .def_readwrite("randomizeCharge", &Config::randomizeCharge)
0173         .def_readwrite("numParticles", &Config::numParticles)
0174         .def_readwrite("mass", &Config::mass)
0175         .def_readwrite("charge", &Config::charge)
0176         .def_property(
0177             "p", [](Config& cfg) { return std::pair{cfg.pMin, cfg.pMax}; },
0178             [](Config& cfg, std::pair<double, double> value) {
0179               cfg.pMin = value.first;
0180               cfg.pMax = value.second;
0181             })
0182         .def_property(
0183             "phi",
0184             [](Config& cfg) { return std::pair{cfg.phiMin, cfg.phiMax}; },
0185             [](Config& cfg, std::pair<double, double> value) {
0186               cfg.phiMin = value.first;
0187               cfg.phiMax = value.second;
0188             })
0189         .def_property(
0190             "theta",
0191             [](Config& cfg) { return std::pair{cfg.thetaMin, cfg.thetaMax}; },
0192             [](Config& cfg, std::pair<double, double> value) {
0193               cfg.thetaMin = value.first;
0194               cfg.thetaMax = value.second;
0195             })
0196         .def_property(
0197             "eta",
0198             [](Config& cfg) {
0199               return std::pair{AngleHelpers::etaFromTheta(cfg.thetaMin),
0200                                AngleHelpers::etaFromTheta(cfg.thetaMax)};
0201             },
0202             [](Config& cfg, std::pair<double, double> value) {
0203               cfg.thetaMin = AngleHelpers::thetaFromEta(value.first);
0204               cfg.thetaMax = AngleHelpers::thetaFromEta(value.second);
0205             });
0206   }
0207
0208   py::class_<FixedMultiplicityGenerator, MultiplicityGenerator,
0209              std::shared_ptr<FixedMultiplicityGenerator>>(
0210       mex, "FixedMultiplicityGenerator")
0211       .def(py::init<>())
0212       .def(py::init([](std::size_t n) {
0213              FixedMultiplicityGenerator g;
0214              g.n = n;
0215              return g;
0216            }),
0217            py::arg("n"))
0218       .def_readwrite("n", &FixedMultiplicityGenerator::n);
0219
0220   py::class_<PoissonMultiplicityGenerator, MultiplicityGenerator,
0221              std::shared_ptr<PoissonMultiplicityGenerator>>(
0222       mex, "PoissonMultiplicityGenerator")
0223       .def(py::init<>())
0224       .def(py::init([](double mean) {
0225              PoissonMultiplicityGenerator g;
0226              g.mean = mean;
0227              return g;
0228            }),
0229            py::arg("mean"))
0230       .def_readwrite("mean", &PoissonMultiplicityGenerator::mean);
0231 }
0232
0233 }  // namespace ActsPython