Core/src/Definitions.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/EventData/ParticleHypothesis.hpp"
0013
0014 #include <format>
0015 #include <type_traits>
0016
0017 #include <pybind11/pybind11.h>
0018 #include <pybind11/stl.h>
0019
0020 namespace py = pybind11;
0021 using namespace pybind11::literals;
0022
0023 namespace ActsPython {
0024
0025 /// @brief This adds the classes from Core/Definitions to the python module
0026 /// @param m the pybind11 core module
0027 void addDefinitions(py::module_& m) {
0028   using namespace Acts;
0029   // Add algebra classes here
0030   py::class_<Vector2>(m, "Vector2")
0031       .def(py::init<double, double>())
0032       .def(py::init([](std::array<double, 2> a) {
0033         Vector2 v;
0034         v << a[0], a[1];
0035         return v;
0036       }))
0037       .def("__getitem__",
0038            [](const Vector2& self, Eigen::Index i) { return self[i]; })
0039       .def("__str__", [](const Vector2& self) {
0040         std::stringstream ss;
0041         ss << self.transpose();
0042         return ss.str();
0043       });
0044
0045   py::class_<Vector3>(m, "Vector3")
0046       .def(py::init<double, double, double>())
0047       .def(py::init([](std::array<double, 3> a) {
0048         Vector3 v;
0049         v << a[0], a[1], a[2];
0050         return v;
0051       }))
0052       .def_static("UnitX", []() -> Vector3 { return Vector3::UnitX(); })
0053       .def_static("UnitY", []() -> Vector3 { return Vector3::UnitY(); })
0054       .def_static("UnitZ", []() -> Vector3 { return Vector3::UnitZ(); })
0055       .def("__add__",
0056            [](const Vector3& self, const Vector3& other) {
0057              return (self + other).eval();
0058            })
0059       .def("__sub__",
0060            [](const Vector3& self, const Vector3& other) {
0061              return (self - other).eval();
0062            })
0063       .def("__mul__",
0064            [](const Vector3& self, const Vector3& other) {
0065              return self.cwiseProduct(other).eval();
0066            })
0067       .def("cross",
0068            [](const Vector3& self, const Vector3& other) {
0069              return self.cross(other).eval();
0070            })
0071       .def("__getitem__",
0072            [](const Vector3& self, Eigen::Index i) { return self[i]; })
0073       .def("__str__", [](const Vector3& self) {
0074         return std::format("({}, {}, {})", self[0], self[1], self[2]);
0075       });
0076
0077   py::class_<Vector4>(m, "Vector4")
0078       .def(py::init<double, double, double, double>())
0079       .def(py::init([](std::array<double, 4> a) {
0080         Vector4 v;
0081         v << a[0], a[1], a[2], a[3];
0082         return v;
0083       }))
0084       .def("__add__",
0085            [](const Vector4& self, const Vector4& other) {
0086              return (self + other).eval();
0087            })
0088       .def("__sub__",
0089            [](const Vector4& self, const Vector4& other) {
0090              return (self - other).eval();
0091            })
0092       .def("__mul__",
0093            [](const Vector4& self, const Vector4& other) {
0094              return self.cwiseProduct(other).eval();
0095            })
0096       .def("__getitem__",
0097            [](const Vector4& self, Eigen::Index i) { return self[i]; })
0098       .def("__str__", [](const Vector4& self) {
0099         return std::format("({}, {}, {}, {})", self[0], self[1], self[2],
0100                            self[3]);
0101       });
0102
0103   py::class_<Translation3>(m, "Translation3")
0104       .def(py::init([](const Vector3& a) { return Translation3(a); }))
0105       .def(py::init([](std::array<double, 3> a) {
0106         return Translation3(Vector3(a[0], a[1], a[2]));
0107       }))
0108       .def("__getitem__", [](const Translation3& self,
0109                              Eigen::Index i) { return self.translation()[i]; })
0110       .def("__str__", [](const Translation3& self) {
0111         std::stringstream ss;
0112         ss << self.translation().transpose();
0113         return ss.str();
0114       });
0115
0116   py::class_<RotationMatrix3>(m, "RotationMatrix3")
0117       .def(py::init([](const Vector3& u, const Vector3& v, const Vector3& w) {
0118         RotationMatrix3 r;
0119         r.col(0) = u;
0120         r.col(1) = v;
0121         r.col(2) = w;
0122         return r;
0123       }))
0124       .def("__getitem__",
0125            [](const RotationMatrix3& self, std::array<Eigen::Index, 2> ij) {
0126              return self.matrix()(ij[0], ij[1]);
0127            })
0128       .def("__str__", [](const RotationMatrix3& self) {
0129         std::stringstream ss;
0130         ss << self.matrix();
0131         return ss.str();
0132       });
0133
0134   py::class_<AngleAxis3>(m, "AngleAxis3")
0135       .def(py::init([](double angle, const Vector3& axis) {
0136         return AngleAxis3(angle, axis);
0137       }))
0138       .def("__str__", [](const AngleAxis3& self) {
0139         std::stringstream ss;
0140         ss << "Angle: " << self.angle()
0141            << ", Axis: " << self.axis().transpose();
0142         return ss.str();
0143       });
0144
0145   py::class_<Transform3>(m, "Transform3")
0146       .def(py::init([](const Vector3& translation) -> Transform3 {
0147         return Transform3{Translation3{translation}};
0148       }))
0149       .def(py::init([](const Vector3& translation,
0150                        const RotationMatrix3& rotation) -> Transform3 {
0151         Transform3 t;
0152         t.prerotate(rotation);
0153         t.pretranslate(translation);
0154         return t;
0155       }))
0156       .def_property_readonly(
0157           "translation",
0158           [](const Transform3& self) -> Vector3 { return self.translation(); })
0159       .def_property_readonly("rotation",
0160                              [](const Transform3& self) -> RotationMatrix3 {
0161                                return self.rotation();
0162                              })
0163       .def_static("Identity", &Transform3::Identity)
0164       .def("__mul__", [](const Transform3& self,
0165                          const Transform3& other) { return self * other; })
0166       .def("__mul__", [](const Transform3& self,
0167                          const Translation3& other) { return self * other; })
0168       .def("__mul__", [](const Transform3& self,
0169                          const AngleAxis3& other) { return self * other; })
0170       .def("__str__", [](const Transform3& self) {
0171         std::stringstream ss;
0172         ss << self.matrix();
0173         return ss.str();
0174       });
0175
0176   // Add the unit constants
0177   auto u = m.def_submodule("UnitConstants");
0178
0179 #define UNIT(x) u.attr(#x) = UnitConstants::x;
0180
0181   UNIT(fm)
0182   UNIT(pm)
0183   UNIT(um)
0184   UNIT(nm)
0185   UNIT(mm)
0186   UNIT(cm)
0187   UNIT(m)
0188   UNIT(km)
0189   UNIT(mm2)
0190   UNIT(cm2)
0191   UNIT(m2)
0192   UNIT(mm3)
0193   UNIT(cm3)
0194   UNIT(m3)
0195   UNIT(s)
0196   UNIT(fs)
0197   UNIT(ps)
0198   UNIT(ns)
0199   UNIT(us)
0200   UNIT(ms)
0201   UNIT(min)
0202   UNIT(h)
0203   UNIT(mrad)
0204   UNIT(rad)
0205   UNIT(degree)
0206   UNIT(eV)
0207   UNIT(keV)
0208   UNIT(MeV)
0209   UNIT(GeV)
0210   UNIT(TeV)
0211   UNIT(J)
0212   UNIT(u)
0213   UNIT(g)
0214   UNIT(kg)
0215   UNIT(e)
0216   UNIT(T)
0217   UNIT(Gauss)
0218   UNIT(kGauss)
0219   UNIT(mol)
0220
0221 #undef UNIT
0222
0223   // Add the PdgParticle enum
0224   py::enum_<PdgParticle>(m, "PdgParticle")
0225       .value("eInvalid", PdgParticle::eInvalid)
0226       .value("eElectron", PdgParticle::eElectron)
0227       .value("eAntiElectron", PdgParticle::eAntiElectron)
0228       .value("ePositron", PdgParticle::ePositron)
0229       .value("eMuon", PdgParticle::eMuon)
0230       .value("eAntiMuon", PdgParticle::eAntiMuon)
0231       .value("eTau", PdgParticle::eTau)
0232       .value("eAntiTau", PdgParticle::eAntiTau)
0233       .value("eGamma", PdgParticle::eGamma)
0234       .value("ePionZero", PdgParticle::ePionZero)
0235       .value("ePionPlus", PdgParticle::ePionPlus)
0236       .value("ePionMinus", PdgParticle::ePionMinus)
0237       .value("eKaonPlus", PdgParticle::eKaonPlus)
0238       .value("eKaonMinus", PdgParticle::eKaonMinus)
0239       .value("eNeutron", PdgParticle::eNeutron)
0240       .value("eAntiNeutron", PdgParticle::eAntiNeutron)
0241       .value("eProton", PdgParticle::eProton)
0242       .value("eAntiProton", PdgParticle::eAntiProton)
0243       .value("eLead", PdgParticle::eLead)
0244       .value("eJPsi", PdgParticle::eJPsi)
0245       .value("eB0", PdgParticle::eB0)
0246       .value("eBPlus", PdgParticle::eBPlus)
0247       .value("eD0", PdgParticle::eD0)
0248       .value("eDPlus", PdgParticle::eDPlus)
0249       .value("eAntiB0", PdgParticle::eAntiB0)
0250       .value("eAntiD0", PdgParticle::eAntiD0)
0251       .value("eNeutrinoE", PdgParticle::eNeutrinoE)
0252       .value("eNeutrinoMu", PdgParticle::eNeutrinoMu)
0253       .value("eNeutrinoTau", PdgParticle::eNeutrinoTau)
0254       .value("eAntiNeutrinoE", PdgParticle::eAntiNeutrinoE)
0255       .value("eAntiNeutrinoMu", PdgParticle::eAntiNeutrinoMu)
0256       .value("eAntiNeutrinoTau", PdgParticle::eAntiNeutrinoTau);
0257
0258   // Add the parsePdgParticle function
0259   m.def("parsePdgParticle", &parsePdgParticle, py::arg("name"));
0260
0261   py::class_<ParticleHypothesis>(m, "ParticleHypothesis")
0262       .def(py::init([](PdgParticle absPdg, float mass, float absCharge) {
0263              return ParticleHypothesis(absPdg, mass, AnyCharge{absCharge});
0264            }),
0265            py::arg("pdg"), py::arg("mass"), py::arg("absCharge"))
0266       .def(py::init([](std::underlying_type_t<PdgParticle> absPdg, float mass,
0267                        float absCharge) {
0268              return ParticleHypothesis(static_cast<PdgParticle>(absPdg), mass,
0269                                        AnyCharge{absCharge});
0270            }),
0271            py::arg("absPdg"), py::arg("mass"), py::arg("absCharge"))
0272       .def("__str__",
0273            [](const ParticleHypothesis& particleHypothesis) {
0274              std::stringstream os;
0275              particleHypothesis.toStream(os);
0276              return os.str();
0277            })
0278       .def("absolutePdg",
0279            [](const ParticleHypothesis& p) { return p.absolutePdg(); })
0280       .def("mass", [](const ParticleHypothesis& p) { return p.mass(); })
0281       .def("absoluteCharge",
0282            [](const ParticleHypothesis& p) { return p.absoluteCharge(); })
0283       .def_property_readonly_static(
0284           "muon",
0285           [](py::object /* self */) { return ParticleHypothesis::muon(); })
0286       .def_property_readonly_static(
0287           "pion",
0288           [](py::object /* self */) { return ParticleHypothesis::pion(); })
0289       .def_property_readonly_static(
0290           "electron",
0291           [](py::object /* self */) { return ParticleHypothesis::electron(); })
0292       .def_property_readonly_static(
0293           "kaon",
0294           [](py::object /* self */) { return ParticleHypothesis::kaon(); })
0295       .def_property_readonly_static(
0296           "proton",
0297           [](py::object /* self */) { return ParticleHypothesis::proton(); })
0298       .def_property_readonly_static(
0299           "geantino",
0300           [](py::object /* self */) { return ParticleHypothesis::geantino(); })
0301       .def_property_readonly_static(
0302           "chargedGeantino", [](py::object /* self */) {
0303             return ParticleHypothesis::chargedGeantino();
0304           });
0305 }
0306
0307 }  // namespace ActsPython