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