Python/src/Generators.cpp

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