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