Python/src/Base.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/Definitions/PdgParticle.hpp"
0011 #include "Acts/Definitions/Units.hpp"
0012 #include "Acts/Geometry/GeometryContext.hpp"
0013 #include "Acts/MagneticField/MagneticFieldContext.hpp"
0014 #include "Acts/Plugins/Python/Utilities.hpp"
0015 #include "Acts/Utilities/Any.hpp"
0016 #include "Acts/Utilities/AxisDefinitions.hpp"
0017 #include "Acts/Utilities/BinningData.hpp"
0018 #include "Acts/Utilities/CalibrationContext.hpp"
0019 #include "Acts/Utilities/Logger.hpp"
0020
0021 #include <array>
0022 #include <exception>
0023 #include <memory>
0024 #include <string>
0025 #include <unordered_map>
0026
0027 #include <pybind11/eval.h>
0028 #include <pybind11/pybind11.h>
0029 #include <pybind11/pytypes.h>
0030 #include <pybind11/stl.h>
0031
0032 namespace py = pybind11;
0033 using namespace pybind11::literals;
0034
0035 namespace Acts::Python {
0036
0037 void addContext(Context& ctx) {
0038   auto& m = ctx.get("main");
0039
0040   py::class_<Acts::GeometryContext>(m, "GeometryContext").def(py::init<>());
0041   py::class_<Acts::MagneticFieldContext>(m, "MagneticFieldContext")
0042       .def(py::init<>());
0043   py::class_<Acts::CalibrationContext>(m, "CalibrationContext")
0044       .def(py::init<>());
0045 }
0046
0047 void addAny(Context& ctx) {
0048   auto& m = ctx.get("main");
0049
0050   py::class_<Acts::AnyBase<512>>(m, "AnyBase512").def(py::init<>());
0051 }
0052
0053 void addUnits(Context& ctx) {
0054   auto& m = ctx.get("main");
0055   auto u = m.def_submodule("UnitConstants");
0056
0057 #define UNIT(x) u.attr(#x) = Acts::UnitConstants::x;
0058
0059   UNIT(fm)
0060   UNIT(pm)
0061   UNIT(um)
0062   UNIT(nm)
0063   UNIT(mm)
0064   UNIT(cm)
0065   UNIT(m)
0066   UNIT(km)
0067   UNIT(mm2)
0068   UNIT(cm2)
0069   UNIT(m2)
0070   UNIT(mm3)
0071   UNIT(cm3)
0072   UNIT(m3)
0073   UNIT(s)
0074   UNIT(fs)
0075   UNIT(ps)
0076   UNIT(ns)
0077   UNIT(us)
0078   UNIT(ms)
0079   UNIT(min)
0080   UNIT(h)
0081   UNIT(mrad)
0082   UNIT(rad)
0083   UNIT(degree)
0084   UNIT(eV)
0085   UNIT(keV)
0086   UNIT(MeV)
0087   UNIT(GeV)
0088   UNIT(TeV)
0089   UNIT(J)
0090   UNIT(u)
0091   UNIT(g)
0092   UNIT(kg)
0093   UNIT(e)
0094   UNIT(T)
0095   UNIT(Gauss)
0096   UNIT(kGauss)
0097   UNIT(mol)
0098
0099 #undef UNIT
0100 }
0101
0102 class PythonLogger {
0103  public:
0104   PythonLogger(const std::string& name, Acts::Logging::Level level)
0105       : m_name{name}, m_logger{Acts::getDefaultLogger(m_name, level)} {}
0106
0107   void log(Acts::Logging::Level level, const std::string& message) const {
0108     m_logger->log(level, message);
0109   }
0110
0111   void setLevel(Acts::Logging::Level level) {
0112     m_logger = Acts::getDefaultLogger(m_name, level);
0113   }
0114
0115  private:
0116   std::string m_name;
0117   std::unique_ptr<const Logger> m_logger;
0118 };
0119
0120 void addLogging(Acts::Python::Context& ctx) {
0121   auto& m = ctx.get("main");
0122   auto logging = m.def_submodule("logging", "");
0123
0124   auto levelEnum = py::enum_<Acts::Logging::Level>(logging, "Level")
0125                        .value("VERBOSE", Acts::Logging::VERBOSE)
0126                        .value("DEBUG", Acts::Logging::DEBUG)
0127                        .value("INFO", Acts::Logging::INFO)
0128                        .value("WARNING", Acts::Logging::WARNING)
0129                        .value("ERROR", Acts::Logging::ERROR)
0130                        .value("FATAL", Acts::Logging::FATAL)
0131                        .value("MAX", Acts::Logging::MAX)
0132                        .export_values();
0133
0134   levelEnum
0135       .def("__lt__", [](Acts::Logging::Level self,
0136                         Acts::Logging::Level other) { return self < other; })
0137       .def("__gt__", [](Acts::Logging::Level self,
0138                         Acts::Logging::Level other) { return self > other; })
0139       .def("__le__", [](Acts::Logging::Level self,
0140                         Acts::Logging::Level other) { return self <= other; })
0141       .def("__ge__", [](Acts::Logging::Level self, Acts::Logging::Level other) {
0142         return self >= other;
0143       });
0144
0145   auto makeLogFunction = [](Acts::Logging::Level level) {
0146     return
0147         [level](PythonLogger& logger, const std::string& fmt, py::args args) {
0148           auto locals = py::dict();
0149           locals["args"] = args;
0150           locals["fmt"] = fmt;
0151           py::exec(R"(
0152         message = fmt % args
0153     )",
0154                    py::globals(), locals);
0155
0156           auto message = locals["message"].cast<std::string>();
0157
0158           logger.log(level, message);
0159         };
0160   };
0161
0162   py::class_<Logger>(m, "Logger");
0163
0164   auto logger =
0165       py::class_<PythonLogger, std::shared_ptr<PythonLogger>>(logging, "Logger")
0166           .def("log", &PythonLogger::log)
0167           .def("verbose", makeLogFunction(Acts::Logging::VERBOSE))
0168           .def("debug", makeLogFunction(Acts::Logging::DEBUG))
0169           .def("info", makeLogFunction(Acts::Logging::INFO))
0170           .def("warning", makeLogFunction(Acts::Logging::WARNING))
0171           .def("error", makeLogFunction(Acts::Logging::ERROR))
0172           .def("fatal", makeLogFunction(Acts::Logging::FATAL))
0173           .def("setLevel", &PythonLogger::setLevel);
0174
0175   static std::unordered_map<std::string, std::shared_ptr<PythonLogger>>
0176       pythonLoggers = {{"root", std::make_shared<PythonLogger>(
0177                                     "Python", Acts::Logging::INFO)}};
0178
0179   logging.def(
0180       "getLogger",
0181       [](const std::string& name) {
0182         if (!pythonLoggers.contains(name)) {
0183           pythonLoggers[name] =
0184               std::make_shared<PythonLogger>(name, Acts::Logging::INFO);
0185         }
0186         return pythonLoggers[name];
0187       },
0188       py::arg("name") = "root");
0189
0190   logging.def("setLevel", [](Acts::Logging::Level level) {
0191     pythonLoggers.at("root")->setLevel(level);
0192   });
0193
0194   auto makeModuleLogFunction = [](Acts::Logging::Level level) {
0195     return [level](const std::string& fmt, py::args args) {
0196       auto locals = py::dict();
0197       locals["args"] = args;
0198       locals["fmt"] = fmt;
0199       py::exec(R"(
0200         message = fmt % args
0201     )",
0202                py::globals(), locals);
0203
0204       auto message = locals["message"].cast<std::string>();
0205
0206       pythonLoggers.at("root")->log(level, message);
0207     };
0208   };
0209
0210   logging.def("setFailureThreshold", &Logging::setFailureThreshold);
0211   logging.def("getFailureThreshold", &Logging::getFailureThreshold);
0212
0213   struct ScopedFailureThresholdContextManager {
0214     std::optional<Logging::ScopedFailureThreshold> m_scopedFailureThreshold =
0215         std::nullopt;
0216     Logging::Level m_level;
0217
0218     explicit ScopedFailureThresholdContextManager(Logging::Level level)
0219         : m_level(level) {}
0220
0221     void enter() { m_scopedFailureThreshold.emplace(m_level); }
0222
0223     void exit(const py::object& /*exc_type*/, const py::object& /*exc_value*/,
0224               const py::object& /*traceback*/) {
0225       m_scopedFailureThreshold.reset();
0226     }
0227   };
0228
0229   py::class_<ScopedFailureThresholdContextManager>(logging,
0230                                                    "ScopedFailureThreshold")
0231       .def(py::init<Logging::Level>(), "level"_a)
0232       .def("__enter__", &ScopedFailureThresholdContextManager::enter)
0233       .def("__exit__", &ScopedFailureThresholdContextManager::exit);
0234
0235   static py::exception<Logging::ThresholdFailure> exc(
0236       logging, "ThresholdFailure", PyExc_RuntimeError);
0237   // NOLINTNEXTLINE(performance-unnecessary-value-param)
0238   py::register_exception_translator([](std::exception_ptr p) {
0239     try {
0240       if (p) {
0241         std::rethrow_exception(p);
0242       }
0243     } catch (const std::exception& e) {
0244       std::string what = e.what();
0245       if (what.find("ACTS_LOG_FAILURE_THRESHOLD") != std::string::npos) {
0246         py::set_error(exc, e.what());
0247       } else {
0248         std::rethrow_exception(p);
0249       }
0250     }
0251   });
0252
0253   logging.def("verbose", makeModuleLogFunction(Acts::Logging::VERBOSE));
0254   logging.def("debug", makeModuleLogFunction(Acts::Logging::DEBUG));
0255   logging.def("info", makeModuleLogFunction(Acts::Logging::INFO));
0256   logging.def("warning", makeModuleLogFunction(Acts::Logging::WARNING));
0257   logging.def("error", makeModuleLogFunction(Acts::Logging::ERROR));
0258   logging.def("fatal", makeModuleLogFunction(Acts::Logging::FATAL));
0259 }
0260
0261 void addPdgParticle(Acts::Python::Context& ctx) {
0262   auto& m = ctx.get("main");
0263   py::enum_<Acts::PdgParticle>(m, "PdgParticle")
0264       .value("eInvalid", Acts::PdgParticle::eInvalid)
0265       .value("eElectron", Acts::PdgParticle::eElectron)
0266       .value("eAntiElectron", Acts::PdgParticle::eAntiElectron)
0267       .value("ePositron", Acts::PdgParticle::ePositron)
0268       .value("eMuon", Acts::PdgParticle::eMuon)
0269       .value("eAntiMuon", Acts::PdgParticle::eAntiMuon)
0270       .value("eTau", Acts::PdgParticle::eTau)
0271       .value("eAntiTau", Acts::PdgParticle::eAntiTau)
0272       .value("eGamma", Acts::PdgParticle::eGamma)
0273       .value("ePionZero", Acts::PdgParticle::ePionZero)
0274       .value("ePionPlus", Acts::PdgParticle::ePionPlus)
0275       .value("ePionMinus", Acts::PdgParticle::ePionMinus)
0276       .value("eKaonPlus", Acts::PdgParticle::eKaonPlus)
0277       .value("eKaonMinus", Acts::PdgParticle::eKaonMinus)
0278       .value("eNeutron", Acts::PdgParticle::eNeutron)
0279       .value("eAntiNeutron", Acts::PdgParticle::eAntiNeutron)
0280       .value("eProton", Acts::PdgParticle::eProton)
0281       .value("eAntiProton", Acts::PdgParticle::eAntiProton)
0282       .value("eLead", Acts::PdgParticle::eLead);
0283 }
0284
0285 void addAlgebra(Acts::Python::Context& ctx) {
0286   auto& m = ctx.get("main");
0287
0288   py::class_<Acts::Vector2>(m, "Vector2")
0289       .def(py::init<double, double>())
0290       .def(py::init([](std::array<double, 2> a) {
0291         Acts::Vector2 v;
0292         v << a[0], a[1];
0293         return v;
0294       }))
0295       .def("__getitem__",
0296            [](const Acts::Vector2& self, Eigen::Index i) { return self[i]; })
0297       .def("__str__", [](const Acts::Vector3& self) {
0298         std::stringstream ss;
0299         ss << self.transpose();
0300         return ss.str();
0301       });
0302
0303   py::class_<Acts::Vector3>(m, "Vector3")
0304       .def(py::init<double, double, double>())
0305       .def(py::init([](std::array<double, 3> a) {
0306         Acts::Vector3 v;
0307         v << a[0], a[1], a[2];
0308         return v;
0309       }))
0310       .def_static("UnitX", []() -> Vector3 { return Acts::Vector3::UnitX(); })
0311       .def_static("UnitY", []() -> Vector3 { return Acts::Vector3::UnitY(); })
0312       .def_static("UnitZ", []() -> Vector3 { return Acts::Vector3::UnitZ(); })
0313
0314       .def("__getitem__",
0315            [](const Acts::Vector3& self, Eigen::Index i) { return self[i]; })
0316       .def("__str__", [](const Acts::Vector3& self) {
0317         std::stringstream ss;
0318         ss << self.transpose();
0319         return ss.str();
0320       });
0321
0322   py::class_<Acts::Vector4>(m, "Vector4")
0323       .def(py::init<double, double, double, double>())
0324       .def(py::init([](std::array<double, 4> a) {
0325         Acts::Vector4 v;
0326         v << a[0], a[1], a[2], a[3];
0327         return v;
0328       }))
0329       .def("__getitem__",
0330            [](const Acts::Vector4& self, Eigen::Index i) { return self[i]; });
0331
0332   py::class_<Acts::Transform3>(m, "Transform3")
0333       .def(py::init<>())
0334       .def(py::init([](const Vector3& translation) -> Transform3 {
0335         return Transform3{Translation3{translation}};
0336       }))
0337       .def_property_readonly("translation",
0338                              [](const Acts::Transform3& self) -> Vector3 {
0339                                return self.translation();
0340                              })
0341       .def_static("Identity", &Acts::Transform3::Identity)
0342       .def("__mul__",
0343            [](const Acts::Transform3& self, const Acts::Transform3& other) {
0344              return self * other;
0345            })
0346       .def("__mul__",
0347            [](const Acts::Transform3& self, const Acts::Translation3& other) {
0348              return self * other;
0349            })
0350       .def("__mul__",
0351            [](const Acts::Transform3& self, const Acts::AngleAxis3& other) {
0352              return self * other;
0353            })
0354       .def("__str__", [](const Acts::Transform3& self) {
0355         std::stringstream ss;
0356         ss << self.matrix();
0357         return ss.str();
0358       });
0359
0360   py::class_<Acts::Translation3>(m, "Translation3")
0361       .def(py::init(
0362           [](const Acts::Vector3& a) { return Acts::Translation3(a); }))
0363       .def(py::init([](std::array<double, 3> a) {
0364         return Acts::Translation3(Acts::Vector3(a[0], a[1], a[2]));
0365       }))
0366       .def("__str__", [](const Acts::Translation3& self) {
0367         std::stringstream ss;
0368         ss << self.translation().transpose();
0369         return ss.str();
0370       });
0371
0372   py::class_<Acts::AngleAxis3>(m, "AngleAxis3")
0373       .def(py::init([](double angle, const Acts::Vector3& axis) {
0374         return Acts::AngleAxis3(angle, axis);
0375       }))
0376       .def("__str__", [](const Acts::Transform3& self) {
0377         std::stringstream ss;
0378         ss << self.matrix();
0379         return ss.str();
0380       });
0381 }
0382
0383 void addBinning(Context& ctx) {
0384   auto& m = ctx.get("main");
0385
0386   auto binningValue = py::enum_<Acts::AxisDirection>(m, "AxisDirection")
0387                           .value("AxisX", Acts::AxisDirection::AxisX)
0388                           .value("AxisY", Acts::AxisDirection::AxisY)
0389                           .value("AxisZ", Acts::AxisDirection::AxisZ)
0390                           .value("AxisR", Acts::AxisDirection::AxisR)
0391                           .value("AxisPhi", Acts::AxisDirection::AxisPhi)
0392                           .value("AxisRPhi", Acts::AxisDirection::AxisRPhi)
0393                           .value("AxisTheta", Acts::AxisDirection::AxisTheta)
0394                           .value("AxisEta", Acts::AxisDirection::AxisEta)
0395                           .value("AxisMag", Acts::AxisDirection::AxisMag);
0396
0397   auto boundaryType = py::enum_<Acts::AxisBoundaryType>(m, "AxisBoundaryType")
0398                           .value("Bound", Acts::AxisBoundaryType::Bound)
0399                           .value("Closed", Acts::AxisBoundaryType::Closed)
0400                           .value("Open", Acts::AxisBoundaryType::Open);
0401
0402   auto axisType = py::enum_<Acts::AxisType>(m, "AxisType")
0403                       .value("equidistant", Acts::AxisType::Equidistant)
0404                       .value("variable", Acts::AxisType::Variable);
0405 }
0406
0407 }  // namespace Acts::Python