Examples/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/SpacePointContainer2.hpp"
0010 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0011 #include "ActsExamples/EventData/ProtoTrack.hpp"
0012 #include "ActsExamples/EventData/Track.hpp"
0013 #include "ActsPython/Utilities/WhiteBoardRegistry.hpp"
0014
0015 #include <pybind11/numpy.h>
0016 #include <pybind11/pybind11.h>
0017 #include <pybind11/stl_bind.h>
0018
0019 // Prevent stl.h's list_caster-based type_caster<std::vector<T>> from matching
0020 // ProtoTrackContainer, which would break py::cast<std::unique_ptr<T>> needed
0021 // by WhiteBoardRegistry. The full specialization takes priority over stl.h's
0022 // partial specialization regardless of include order.
0023 PYBIND11_MAKE_OPAQUE(ActsExamples::ProtoTrackContainer)
0024
0025 namespace py = pybind11;
0026
0027 using namespace ActsExamples;
0028
0029 namespace ActsPython {
0030
0031 void addEventData(py::module& mex) {
0032   py::class_<Acts::TrackStateType>(mex, "TrackStateTypeFlags")
0033       .def_property_readonly("isMeasurement",
0034                              &Acts::TrackStateType::isMeasurement)
0035       .def_property_readonly("isOutlier", &Acts::TrackStateType::isOutlier)
0036       .def_property_readonly("isHole", &Acts::TrackStateType::isHole)
0037       .def_property_readonly("hasMaterial", &Acts::TrackStateType::hasMaterial)
0038       .def_property_readonly("isSharedHit", &Acts::TrackStateType::isSharedHit)
0039       .def_property_readonly("hasParameters",
0040                              &Acts::TrackStateType::hasParameters);
0041
0042   py::class_<ConstTrackStateProxy>(mex, "ConstTrackStateProxy")
0043       .def_property_readonly(
0044           "typeFlags",
0045           [](const ConstTrackStateProxy& self) {
0046             return Acts::TrackStateType{self.typeFlags().raw()};
0047           })
0048       .def_property_readonly("hasPredicted",
0049                              &ConstTrackStateProxy::hasPredicted)
0050       .def_property_readonly("hasFiltered", &ConstTrackStateProxy::hasFiltered)
0051       .def_property_readonly("hasSmoothed", &ConstTrackStateProxy::hasSmoothed)
0052       .def_property_readonly("predicted",
0053                              [](const ConstTrackStateProxy& self) {
0054                                return Acts::BoundVector{self.predicted()};
0055                              })
0056       .def_property_readonly("filtered",
0057                              [](const ConstTrackStateProxy& self) {
0058                                return Acts::BoundVector{self.filtered()};
0059                              })
0060       .def_property_readonly("smoothed",
0061                              [](const ConstTrackStateProxy& self) {
0062                                return Acts::BoundVector{self.smoothed()};
0063                              })
0064       .def_property_readonly(
0065           "predictedCovariance",
0066           [](const ConstTrackStateProxy& self) {
0067             return Acts::BoundMatrix{self.predictedCovariance()};
0068           })
0069       .def_property_readonly(
0070           "filteredCovariance",
0071           [](const ConstTrackStateProxy& self) {
0072             return Acts::BoundMatrix{self.filteredCovariance()};
0073           })
0074       .def_property_readonly(
0075           "smoothedCovariance",
0076           [](const ConstTrackStateProxy& self) {
0077             return Acts::BoundMatrix{self.smoothedCovariance()};
0078           })
0079       .def_property_readonly("pathLength", &ConstTrackStateProxy::pathLength);
0080
0081   auto constTrackProxy =
0082       py::class_<ConstTrackProxy>(mex, "ConstTrackProxy")
0083           .def_property_readonly("index", &ConstTrackProxy::index)
0084           .def_property_readonly(
0085               "tipIndex",
0086               [](const ConstTrackProxy& self) { return self.tipIndex(); })
0087           .def_property_readonly(
0088               "stemIndex",
0089               [](const ConstTrackProxy& self) { return self.stemIndex(); })
0090           .def_property_readonly("referenceSurface",
0091                                  &ConstTrackProxy::referenceSurface)
0092           .def_property_readonly("hasReferenceSurface",
0093                                  &ConstTrackProxy::hasReferenceSurface)
0094           // Convert the parameters to a BoundVector so they're accessible by
0095           // value
0096           .def_property_readonly("parameters",
0097                                  [](const ConstTrackProxy& self) {
0098                                    return Acts::BoundVector{self.parameters()};
0099                                  })
0100           // Convert the covariance to a BoundMatrix so they're accessible by
0101           // value
0102           .def_property_readonly("covariance",
0103                                  [](const ConstTrackProxy& self) {
0104                                    return Acts::BoundMatrix{self.covariance()};
0105                                  })
0106           .def_property_readonly("particleHypothesis",
0107                                  &ConstTrackProxy::particleHypothesis)
0108           .def_property_readonly(
0109               "nMeasurements",
0110               [](const ConstTrackProxy& self) { return self.nMeasurements(); })
0111           .def_property_readonly(
0112               "nHoles",
0113               [](const ConstTrackProxy& self) { return self.nHoles(); })
0114           .def_property_readonly("isForwardLinked",
0115                                  [](const ConstTrackProxy& self) {
0116                                    return self.isForwardLinked();
0117                                  })
0118           .def_property_readonly("trackStatesReversed",
0119                                  py::cpp_function(
0120                                      [](const ConstTrackProxy& self) {
0121                                        auto range = self.trackStatesReversed();
0122                                        return py::make_iterator(range.begin(),
0123                                                                 range.end());
0124                                      },
0125                                      py::keep_alive<0, 1>()))
0126           .def_property_readonly("trackStates",
0127                                  py::cpp_function(
0128                                      [](const ConstTrackProxy& self) {
0129                                        auto range = self.trackStates();
0130                                        return py::make_iterator(range.begin(),
0131                                                                 range.end());
0132                                      },
0133                                      py::keep_alive<0, 1>()));
0134
0135   // Mark a numpy array as non-writeable and return it.
0136   const auto readOnly = [](auto arr) {
0137     arr.attr("flags").attr("writeable") = py::bool_(false);
0138     return arr;
0139   };
0140
0141   // Factory for zero-copy 1-D views over a plain SoA column.
0142   // `accessor` is called with the backend and must return a const-ref to the
0143   // desired std::vector member.  dtype and stride are derived automatically.
0144   const auto col1D = [&readOnly](auto accessor) {
0145     return [accessor, &readOnly](const py::object& self_py) {
0146       const auto& backend =
0147           self_py.cast<const ConstTrackContainer&>().container();
0148       const auto N = static_cast<py::ssize_t>(backend.size_impl());
0149       const auto& vec = accessor(backend);
0150       using T = typename std::decay_t<decltype(vec)>::value_type;
0151       if (N == 0) {
0152         return readOnly(py::array_t<T>(py::ssize_t{0}));
0153       }
0154       return readOnly(py::array_t<T>({N}, {static_cast<py::ssize_t>(sizeof(T))},
0155                                      vec.data(), self_py));
0156     };
0157   };
0158
0159   auto constTrackContainer =
0160       py::classh<ConstTrackContainer>(mex, "ConstTrackContainer")
0161           .def("__len__", &ConstTrackContainer::size)
0162           .def("__iter__",
0163                [](const ConstTrackContainer& self) {
0164                  return py::make_iterator(self.begin(), self.end());
0165                })
0166           .def("__getitem__", py::overload_cast<ConstTrackContainer::IndexType>(
0167                                   &ConstTrackContainer::getTrack, py::const_))
0168
0169           // Zero-copy numpy array views of the underlying SoA columns.
0170           // The returned arrays are read-only and keep the container alive via
0171           // the numpy base-object mechanism as long as any array is alive.
0172
0173           // shape (N, 6), float64 — bound track parameters [loc0, loc1, phi,
0174           // theta, q/p, t]. Strides account for potential Eigen alignment
0175           // padding between entries.
0176           .def_property_readonly(
0177               "parameters",
0178               [&readOnly](const py::object& self_py) -> py::array_t<double> {
0179                 using CoeffsType = Acts::detail_tsp::FixedSizeTypes<
0180                     Acts::eBoundSize>::Coefficients;
0181                 const auto& backend =
0182                     self_py.cast<const ConstTrackContainer&>().container();
0183                 const auto N = static_cast<py::ssize_t>(backend.size_impl());
0184                 if (N == 0) {
0185                   return readOnly(py::array_t<double>({N, py::ssize_t{6}}));
0186                 }
0187                 return readOnly(py::array_t<double>(
0188                     {N, py::ssize_t{6}},
0189                     {static_cast<py::ssize_t>(sizeof(CoeffsType)),
0190                      static_cast<py::ssize_t>(sizeof(double))},
0191                     backend.m_params[0].data(), self_py));
0192               })
0193
0194           // shape (N, 6, 6), float64, column-major sub-matrices (Eigen
0195           // default). arr[k, i, j] is row i, column j of track k's covariance.
0196           .def_property_readonly(
0197               "covariance",
0198               [&readOnly](const py::object& self_py) -> py::array_t<double> {
0199                 using CovType = Acts::detail_tsp::FixedSizeTypes<
0200                     Acts::eBoundSize>::Covariance;
0201                 const auto& backend =
0202                     self_py.cast<const ConstTrackContainer&>().container();
0203                 const auto N = static_cast<py::ssize_t>(backend.size_impl());
0204                 if (N == 0) {
0205                   return readOnly(
0206                       py::array_t<double>({N, py::ssize_t{6}, py::ssize_t{6}}));
0207                 }
0208                 // Eigen column-major: row stride = 1 double, col stride = 6
0209                 // doubles.
0210                 constexpr py::ssize_t dbl = sizeof(double);
0211                 return readOnly(py::array_t<double>(
0212                     {N, py::ssize_t{6}, py::ssize_t{6}},
0213                     {static_cast<py::ssize_t>(sizeof(CovType)), dbl, 6 * dbl},
0214                     backend.m_cov[0].data(), self_py));
0215               })
0216
0217           .def_property_readonly(
0218               "tipIndex",
0219               col1D([](const auto& b) -> const auto& { return b.m_tipIndex; }))
0220           .def_property_readonly(
0221               "stemIndex",
0222               col1D([](const auto& b) -> const auto& { return b.m_stemIndex; }))
0223           .def_property_readonly("nMeasurements",
0224                                  col1D([](const auto& b) -> const auto& {
0225                                    return b.m_nMeasurements;
0226                                  }))
0227           .def_property_readonly(
0228               "nHoles",
0229               col1D([](const auto& b) -> const auto& { return b.m_nHoles; }))
0230           .def_property_readonly(
0231               "chi2",
0232               col1D([](const auto& b) -> const auto& { return b.m_chi2; }))
0233           .def_property_readonly("ndf", col1D([](const auto& b) -> const auto& {
0234                                    return b.m_ndf;
0235                                  }));
0236
0237   WhiteBoardRegistry::registerClass(constTrackContainer);
0238
0239   py::bind_vector<ProtoTrack>(mex, "ProtoTrack");
0240
0241   auto protoTrackContainer =
0242       py::bind_vector<ProtoTrackContainer, py::smart_holder>(
0243           mex, "ProtoTrackContainer");
0244   WhiteBoardRegistry::registerClass(protoTrackContainer);
0245
0246   mex.attr("kTrackIndexInvalid") = Acts::kTrackIndexInvalid;
0247
0248   py::class_<IndexSourceLink>(mex, "IndexSourceLink")
0249       .def("FromSourceLink",
0250            [](Acts::SourceLink const& sl) { return sl.get<IndexSourceLink>(); })
0251       .def("index", &IndexSourceLink::index)
0252       .def("geometryId", &IndexSourceLink::geometryId);
0253
0254   py::class_<TrackProxy>(mex, "TrackProxy")
0255       .def_property(
0256           "referenceSurface",
0257           [](const TrackProxy& self) -> const Acts::Surface& {
0258             return self.referenceSurface();
0259           },
0260           [](TrackProxy& self, std::shared_ptr<const Acts::Surface> srf) {
0261             self.setReferenceSurface(std::move(srf));
0262           })
0263       .def_property(
0264           "parameters",
0265           [](const TrackProxy& self) {
0266             return Acts::BoundVector{self.parameters()};
0267           },
0268           [](TrackProxy& self, const Acts::BoundVector& v) {
0269             self.parameters() = v;
0270           })
0271       .def_property(
0272           "covariance",
0273           [](const TrackProxy& self) {
0274             return Acts::BoundMatrix{self.covariance()};
0275           },
0276           [](TrackProxy& self, const Acts::BoundMatrix& m) {
0277             self.covariance() = m;
0278           })
0279       .def_property(
0280           "particleHypothesis",
0281           [](const TrackProxy& self) { return self.particleHypothesis(); },
0282           [](TrackProxy& self, const Acts::ParticleHypothesis& hyp) {
0283             self.setParticleHypothesis(hyp);
0284           })
0285       .def_property(
0286           "nMeasurements",
0287           [](const TrackProxy& self) -> std::uint32_t {
0288             return self.nMeasurements();
0289           },
0290           [](TrackProxy& self, std::uint32_t n) { self.nMeasurements() = n; })
0291       .def_property(
0292           "nHoles",
0293           [](const TrackProxy& self) -> std::uint32_t { return self.nHoles(); },
0294           [](TrackProxy& self, std::uint32_t n) { self.nHoles() = n; })
0295       .def_property(
0296           "chi2", [](const TrackProxy& self) -> float { return self.chi2(); },
0297           [](TrackProxy& self, float v) { self.chi2() = v; });
0298
0299   py::class_<TrackContainer>(mex, "TrackContainer")
0300       .def(py::init([]() {
0301         return TrackContainer{std::make_shared<Acts::VectorTrackContainer>(),
0302                               std::make_shared<Acts::VectorMultiTrajectory>()};
0303       }))
0304       .def("__len__", &TrackContainer::size)
0305       .def("makeTrack", &TrackContainer::makeTrack, py::keep_alive<0, 1>())
0306       .def("makeConst", [](TrackContainer& self) {
0307         return ConstTrackContainer{
0308             std::make_shared<Acts::ConstVectorTrackContainer>(
0309                 std::move(self.container())),
0310             std::make_shared<Acts::ConstVectorMultiTrajectory>(
0311                 std::move(self.trackStateContainer()))};
0312       });
0313 }
0314
0315 }  // namespace ActsPython