Core/src/EventData.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/EventData/BoundTrackParameters.hpp"
0010 #include "Acts/EventData/SeedContainer2.hpp"
0011 #include "Acts/EventData/SeedProxy2.hpp"
0012 #include "Acts/EventData/SourceLink.hpp"
0013 #include "Acts/EventData/SpacePointColumns.hpp"
0014 #include "Acts/EventData/SpacePointContainer2.hpp"
0015 #include "Acts/EventData/SpacePointProxy2.hpp"
0016 #include "Acts/EventData/Types.hpp"
0017 #include "Acts/Surfaces/CurvilinearSurface.hpp"
0018 #include "Acts/Surfaces/Surface.hpp"
0019 #include "ActsPython/Utilities/WhiteBoardRegistry.hpp"
0020
0021 #include <array>
0022 #include <memory>
0023 #include <optional>
0024 #include <span>
0025 #include <vector>
0026
0027 #include <pybind11/numpy.h>
0028 #include <pybind11/pybind11.h>
0029 #include <pybind11/stl.h>
0030
0031 namespace py = pybind11;
0032 using namespace pybind11::literals;
0033
0034 using namespace Acts;
0035
0036 namespace ActsPython {
0037
0038 template <typename T>
0039 auto spanToNumpy1d(std::span<T> s, const py::object& base) {
0040   using type = std::remove_cvref_t<T>;
0041   return py::array_t<type>(
0042       {static_cast<py::ssize_t>(s.size())},      // shape
0043       {static_cast<py::ssize_t>(sizeof(type))},  // strides (bytes)
0044       s.data(),                                  // data ptr
0045       base                                       // base/owner
0046   );
0047 }
0048
0049 /// Zero-copy 2D view over a column of std::array<float, Cols>.
0050 /// Returns shape (N, Cols), dtype float32, read-only.
0051 /// Throws if the container does not have the required column.
0052 template <std::size_t Cols>
0053 using ArrayColumnGetter = ConstSpacePointColumnProxy<std::array<float, Cols>> (
0054     SpacePointContainer2::*)() const;
0055
0056 template <std::size_t Cols>
0057 auto arrayColumn(ArrayColumnGetter<Cols> getColumn,
0058                  SpacePointColumns requiredColumn,
0059                  const std::string_view& columnName) {
0060   return [getColumn, requiredColumn,
0061           columnName](const SpacePointContainer2& self) {
0062     if (!self.hasColumns(requiredColumn)) {
0063       throw py::attribute_error(
0064           std::format("SpacePointContainer2 does not have "
0065                       "the {} column",
0066                       columnName));
0067     }
0068     const auto nRows = static_cast<py::ssize_t>(self.size());
0069     if (nRows == 0) {
0070       auto arr = py::array_t<float>(
0071           std::vector<py::ssize_t>{0, static_cast<py::ssize_t>(Cols)});
0072       arr.attr("flags").attr("writeable") = py::bool_(false);
0073       return arr;
0074     }
0075     const auto col = (self.*getColumn)();
0076     const auto& data = col.data();
0077     constexpr py::ssize_t rowStride =
0078         static_cast<py::ssize_t>(Cols * sizeof(float));
0079     constexpr py::ssize_t colStride = sizeof(float);
0080     auto arr = py::array_t<float>(
0081         {nRows, static_cast<py::ssize_t>(Cols)}, {rowStride, colStride},
0082         reinterpret_cast<const float*>(data.data()), py::cast(self));
0083     arr.attr("flags").attr("writeable") = py::bool_(false);
0084     return arr;
0085   };
0086 }
0087
0088 /// @brief This adds the classes from Core/EventData to the python module
0089 /// @param m the pybind11 core module
0090 void addEventData(py::module_& m) {
0091   // SpacePointColumns enum
0092   py::enum_<SpacePointColumns>(m, "SpacePointColumns")
0093       .value("None", SpacePointColumns::None)
0094       .value("SourceLinks", SpacePointColumns::SourceLinks)
0095       .value("X", SpacePointColumns::X)
0096       .value("Y", SpacePointColumns::Y)
0097       .value("Z", SpacePointColumns::Z)
0098       .value("R", SpacePointColumns::R)
0099       .value("Phi", SpacePointColumns::Phi)
0100       .value("Time", SpacePointColumns::Time)
0101       .value("VarianceZ", SpacePointColumns::VarianceZ)
0102       .value("VarianceR", SpacePointColumns::VarianceR)
0103       .value("TopStripVector", SpacePointColumns::TopStripVector)
0104       .value("BottomStripVector", SpacePointColumns::BottomStripVector)
0105       .value("StripCenterDistance", SpacePointColumns::StripCenterDistance)
0106       .value("TopStripCenter", SpacePointColumns::TopStripCenter)
0107       .value("CopyFromIndex", SpacePointColumns::CopyFromIndex)
0108       .value("PackedXY", SpacePointColumns::PackedXY)
0109       .value("PackedZR", SpacePointColumns::PackedZR)
0110       .value("PackedXYZ", SpacePointColumns::PackedXYZ)
0111       .value("PackedXYZR", SpacePointColumns::PackedXYZR)
0112       .value("PackedVarianceZR", SpacePointColumns::PackedVarianceZR)
0113       .value("Strip", SpacePointColumns::Strip)
0114       .value("All", SpacePointColumns::All)
0115       .def("__or__",
0116            [](SpacePointColumns a, SpacePointColumns b) {
0117              return static_cast<SpacePointColumns>(
0118                  static_cast<std::uint32_t>(a) | static_cast<std::uint32_t>(b));
0119            })
0120       .def("__and__", [](SpacePointColumns a, SpacePointColumns b) {
0121         return static_cast<SpacePointColumns>(static_cast<std::uint32_t>(a) &
0122                                               static_cast<std::uint32_t>(b));
0123       });
0124
0125   using FloatColumnGetter =
0126       ConstSpacePointColumnProxy<float> (SpacePointContainer2::*)() const;
0127   auto floatColumn = [](FloatColumnGetter column) {
0128     return [column](const SpacePointContainer2& self) {
0129       return spanToNumpy1d((self.*column)().data(), py::cast(self));
0130     };
0131   };
0132
0133   // SpacePointContainer2
0134   auto spc2 =
0135       py::classh<SpacePointContainer2>(m, "SpacePointContainer2")
0136           .def(py::init<SpacePointColumns>(),
0137                py::arg("columns") = SpacePointColumns::None)
0138           .def_property_readonly("size", &SpacePointContainer2::size)
0139           .def_property_readonly("empty", &SpacePointContainer2::empty)
0140           .def_property_readonly("hasColumns",
0141                                  &SpacePointContainer2::hasColumns)
0142           .def("reserve", &SpacePointContainer2::reserve, py::arg("size"),
0143                py::arg("averageSourceLinks") = 1)
0144           .def("clear", &SpacePointContainer2::clear)
0145           .def("__len__", &SpacePointContainer2::size)
0146           .def("__getitem__",
0147                [](const SpacePointContainer2& self, SpacePointIndex2 idx) {
0148                  return ConstSpacePointProxy2(self, idx);
0149                })
0150           .def("__iter__",
0151                [](const SpacePointContainer2& self) {
0152                  return py::make_iterator(self.begin(), self.end());
0153                })
0154           .def_property_readonly("x",
0155                                  floatColumn(&SpacePointContainer2::xColumn))
0156           .def_property_readonly("y",
0157                                  floatColumn(&SpacePointContainer2::yColumn))
0158           .def_property_readonly("z",
0159                                  floatColumn(&SpacePointContainer2::zColumn))
0160           .def_property_readonly("r",
0161                                  floatColumn(&SpacePointContainer2::rColumn))
0162           .def_property_readonly("phi",
0163                                  floatColumn(&SpacePointContainer2::phiColumn))
0164           .def_property_readonly("time",
0165                                  floatColumn(&SpacePointContainer2::timeColumn))
0166           .def_property_readonly(
0167               "varianceZ", floatColumn(&SpacePointContainer2::varianceZColumn))
0168           .def_property_readonly(
0169               "varianceR", floatColumn(&SpacePointContainer2::varianceRColumn))
0170           .def_property_readonly(
0171               "xyColumn", arrayColumn<2>(&SpacePointContainer2::xyColumn,
0172                                          SpacePointColumns::PackedXY, "xy"))
0173           .def_property_readonly(
0174               "zrColumn", arrayColumn<2>(&SpacePointContainer2::zrColumn,
0175                                          SpacePointColumns::PackedZR, "zr"))
0176           .def_property_readonly(
0177               "xy", arrayColumn<2>(&SpacePointContainer2::xyColumn,
0178                                    SpacePointColumns::PackedXY, "xy"))
0179           .def_property_readonly(
0180               "zr", arrayColumn<2>(&SpacePointContainer2::zrColumn,
0181                                    SpacePointColumns::PackedZR, "zr"))
0182           .def_property_readonly(
0183               "xyz", arrayColumn<3>(&SpacePointContainer2::xyzColumn,
0184                                     SpacePointColumns::PackedXYZ, "xyz"))
0185           .def_property_readonly(
0186               "xyzr", arrayColumn<4>(&SpacePointContainer2::xyzrColumn,
0187                                      SpacePointColumns::PackedXYZR, "xyzr"))
0188           .def_property_readonly(
0189               "varianceZRColumn",
0190               arrayColumn<2>(&SpacePointContainer2::varianceZRColumn,
0191                              SpacePointColumns::PackedVarianceZR, "varianceZR"))
0192           .def_property_readonly(
0193               "topStripVectorColumn",
0194               arrayColumn<3>(&SpacePointContainer2::topStripVectorColumn,
0195                              SpacePointColumns::TopStripVector,
0196                              "topStripVector"));
0197
0198   WhiteBoardRegistry::registerClass(spc2);
0199
0200   // ConstSpacePointProxy2
0201   // Note: SpacePointProxy2 has both mutable (requires(!ReadOnly)) and const
0202   // overloads of x/y/z etc. with different return types. GCC 13 includes the
0203   // constrained-out overloads in pointer-to-member resolution, so we must
0204   // static_cast to the exact return type to disambiguate.
0205   using FloatGetter = float (ConstSpacePointProxy2::*)() const noexcept;
0206   py::class_<ConstSpacePointProxy2>(m, "ConstSpacePointProxy2")
0207       .def_property_readonly("index", &ConstSpacePointProxy2::index)
0208       .def_property_readonly(
0209           "x", static_cast<FloatGetter>(&ConstSpacePointProxy2::x))
0210       .def_property_readonly(
0211           "y", static_cast<FloatGetter>(&ConstSpacePointProxy2::y))
0212       .def_property_readonly(
0213           "z", static_cast<FloatGetter>(&ConstSpacePointProxy2::z))
0214       .def_property_readonly(
0215           "r", static_cast<FloatGetter>(&ConstSpacePointProxy2::r))
0216       .def_property_readonly(
0217           "phi", static_cast<FloatGetter>(&ConstSpacePointProxy2::phi))
0218       .def_property_readonly(
0219           "time", static_cast<FloatGetter>(&ConstSpacePointProxy2::time))
0220       .def_property_readonly(
0221           "varianceZ",
0222           static_cast<FloatGetter>(&ConstSpacePointProxy2::varianceZ))
0223       .def_property_readonly(
0224           "varianceR",
0225           static_cast<FloatGetter>(&ConstSpacePointProxy2::varianceR))
0226       .def_property_readonly(
0227           "sourceLinks", py::cpp_function(
0228                              [](const ConstSpacePointProxy2& self) {
0229                                auto sls = self.sourceLinks();
0230                                return py::make_iterator(sls.begin(), sls.end());
0231                              },
0232                              py::keep_alive<0, 1>()));
0233
0234   // SeedContainer2
0235   auto seedContainer2 =
0236       py::classh<SeedContainer2>(m, "SeedContainer2")
0237           .def(py::init<>())
0238           .def_property_readonly("size", &SeedContainer2::size)
0239           .def_property_readonly("empty", &SeedContainer2::empty)
0240           .def("__len__", &SeedContainer2::size)
0241           .def("__getitem__",
0242                [](const SeedContainer2& self, SeedIndex2 idx) {
0243                  return ConstSeedProxy2(self, idx);
0244                })
0245           .def("__iter__", [](const SeedContainer2& self) {
0246             return py::make_iterator(self.begin(), self.end());
0247           });
0248
0249   WhiteBoardRegistry::registerClass(seedContainer2);
0250
0251   // ConstSeedProxy2
0252   py::class_<ConstSeedProxy2>(m, "ConstSeedProxy2")
0253       .def_property_readonly("index", &ConstSeedProxy2::index)
0254       .def_property_readonly("size", &ConstSeedProxy2::size)
0255       .def_property_readonly("empty", &ConstSeedProxy2::empty)
0256       .def_property_readonly("quality", &ConstSeedProxy2::quality)
0257       .def_property_readonly("vertexZ", &ConstSeedProxy2::vertexZ)
0258       .def_property_readonly(
0259           "spacePointIndices", [](const ConstSeedProxy2& self) {
0260             return spanToNumpy1d(self.spacePointIndices(), py::cast(self));
0261           });
0262
0263   // BoundTrackParameters (alias for
0264   // GenericBoundTrackParameters<ParticleHypothesis>)
0265   py::class_<BoundTrackParameters>(m, "BoundTrackParameters")
0266       .def_property_readonly("parameters",
0267                              [](const BoundTrackParameters& self) {
0268                                return BoundVector(self.parameters());
0269                              })
0270       .def_property_readonly(
0271           "covariance",
0272           [](const BoundTrackParameters& self) { return self.covariance(); })
0273       .def_property_readonly("referenceSurface",
0274                              [](const BoundTrackParameters& self)
0275                                  -> std::shared_ptr<const Surface> {
0276                                return self.referenceSurface().getSharedPtr();
0277                              })
0278       .def_property_readonly("particleHypothesis",
0279                              &BoundTrackParameters::particleHypothesis)
0280       .def_property_readonly("position", &BoundTrackParameters::position)
0281       .def_property_readonly("momentum", &BoundTrackParameters::momentum)
0282       .def_property_readonly("transverseMomentum",
0283                              &BoundTrackParameters::transverseMomentum);
0284
0285   // SourceLink - type-erased wrapper (minimal binding for type identity)
0286   // Note: SourceLink has no default constructor; concrete types (e.g.
0287   // IndexSourceLink from Examples) are typically used to construct it.
0288   py::class_<SourceLink>(m, "SourceLink");
0289 }
0290
0291 }  // namespace ActsPython