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