Back to home page

EIC code displayed by LXR
 
 

    

File indexing completed on 2025-10-13 08:18:25 

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

This page was automatically generated by the 2.3.7 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!