The source code and dockerfile for the GSW2024 AI Lab.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
This repo is archived. You can view files and clone it, but cannot push or open issues/pull-requests.

315 lines
16 KiB

4 weeks ago
  1. /*
  2. tests/eigen.cpp -- automatic conversion of Eigen types
  3. Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
  4. All rights reserved. Use of this source code is governed by a
  5. BSD-style license that can be found in the LICENSE file.
  6. */
  7. #include "pybind11_tests.h"
  8. #include "constructor_stats.h"
  9. #include <pybind11/eigen.h>
  10. #include <pybind11/stl.h>
  11. #include <Eigen/Cholesky>
  12. using MatrixXdR = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
  13. // Sets/resets a testing reference matrix to have values of 10*r + c, where r and c are the
  14. // (1-based) row/column number.
  15. template <typename M> void reset_ref(M &x) {
  16. for (int i = 0; i < x.rows(); i++) for (int j = 0; j < x.cols(); j++)
  17. x(i, j) = 11 + 10*i + j;
  18. }
  19. // Returns a static, column-major matrix
  20. Eigen::MatrixXd &get_cm() {
  21. static Eigen::MatrixXd *x;
  22. if (!x) {
  23. x = new Eigen::MatrixXd(3, 3);
  24. reset_ref(*x);
  25. }
  26. return *x;
  27. }
  28. // Likewise, but row-major
  29. MatrixXdR &get_rm() {
  30. static MatrixXdR *x;
  31. if (!x) {
  32. x = new MatrixXdR(3, 3);
  33. reset_ref(*x);
  34. }
  35. return *x;
  36. }
  37. // Resets the values of the static matrices returned by get_cm()/get_rm()
  38. void reset_refs() {
  39. reset_ref(get_cm());
  40. reset_ref(get_rm());
  41. }
  42. // Returns element 2,1 from a matrix (used to test copy/nocopy)
  43. double get_elem(Eigen::Ref<const Eigen::MatrixXd> m) { return m(2, 1); };
  44. // Returns a matrix with 10*r + 100*c added to each matrix element (to help test that the matrix
  45. // reference is referencing rows/columns correctly).
  46. template <typename MatrixArgType> Eigen::MatrixXd adjust_matrix(MatrixArgType m) {
  47. Eigen::MatrixXd ret(m);
  48. for (int c = 0; c < m.cols(); c++) for (int r = 0; r < m.rows(); r++)
  49. ret(r, c) += 10*r + 100*c;
  50. return ret;
  51. }
  52. struct CustomOperatorNew {
  53. CustomOperatorNew() = default;
  54. Eigen::Matrix4d a = Eigen::Matrix4d::Zero();
  55. Eigen::Matrix4d b = Eigen::Matrix4d::Identity();
  56. EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
  57. };
  58. test_initializer eigen([](py::module &m) {
  59. typedef Eigen::Matrix<float, 5, 6, Eigen::RowMajor> FixedMatrixR;
  60. typedef Eigen::Matrix<float, 5, 6> FixedMatrixC;
  61. typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> DenseMatrixR;
  62. typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic> DenseMatrixC;
  63. typedef Eigen::Matrix<float, 4, Eigen::Dynamic> FourRowMatrixC;
  64. typedef Eigen::Matrix<float, Eigen::Dynamic, 4> FourColMatrixC;
  65. typedef Eigen::Matrix<float, 4, Eigen::Dynamic> FourRowMatrixR;
  66. typedef Eigen::Matrix<float, Eigen::Dynamic, 4> FourColMatrixR;
  67. typedef Eigen::SparseMatrix<float, Eigen::RowMajor> SparseMatrixR;
  68. typedef Eigen::SparseMatrix<float> SparseMatrixC;
  69. m.attr("have_eigen") = true;
  70. m.def("double_col", [](const Eigen::VectorXf &x) -> Eigen::VectorXf { return 2.0f * x; });
  71. m.def("double_row", [](const Eigen::RowVectorXf &x) -> Eigen::RowVectorXf { return 2.0f * x; });
  72. m.def("double_complex", [](const Eigen::VectorXcf &x) -> Eigen::VectorXcf { return 2.0f * x; });
  73. m.def("double_threec", [](py::EigenDRef<Eigen::Vector3f> x) { x *= 2; });
  74. m.def("double_threer", [](py::EigenDRef<Eigen::RowVector3f> x) { x *= 2; });
  75. m.def("double_mat_cm", [](Eigen::MatrixXf x) -> Eigen::MatrixXf { return 2.0f * x; });
  76. m.def("double_mat_rm", [](DenseMatrixR x) -> DenseMatrixR { return 2.0f * x; });
  77. // Different ways of passing via Eigen::Ref; the first and second are the Eigen-recommended
  78. m.def("cholesky1", [](Eigen::Ref<MatrixXdR> x) -> Eigen::MatrixXd { return x.llt().matrixL(); });
  79. m.def("cholesky2", [](const Eigen::Ref<const MatrixXdR> &x) -> Eigen::MatrixXd { return x.llt().matrixL(); });
  80. m.def("cholesky3", [](const Eigen::Ref<MatrixXdR> &x) -> Eigen::MatrixXd { return x.llt().matrixL(); });
  81. m.def("cholesky4", [](Eigen::Ref<const MatrixXdR> x) -> Eigen::MatrixXd { return x.llt().matrixL(); });
  82. // Mutators: these add some value to the given element using Eigen, but Eigen should be mapping into
  83. // the numpy array data and so the result should show up there. There are three versions: one that
  84. // works on a contiguous-row matrix (numpy's default), one for a contiguous-column matrix, and one
  85. // for any matrix.
  86. auto add_rm = [](Eigen::Ref<MatrixXdR> x, int r, int c, double v) { x(r,c) += v; };
  87. auto add_cm = [](Eigen::Ref<Eigen::MatrixXd> x, int r, int c, double v) { x(r,c) += v; };
  88. // Mutators (Eigen maps into numpy variables):
  89. m.def("add_rm", add_rm); // Only takes row-contiguous
  90. m.def("add_cm", add_cm); // Only takes column-contiguous
  91. // Overloaded versions that will accept either row or column contiguous:
  92. m.def("add1", add_rm);
  93. m.def("add1", add_cm);
  94. m.def("add2", add_cm);
  95. m.def("add2", add_rm);
  96. // This one accepts a matrix of any stride:
  97. m.def("add_any", [](py::EigenDRef<Eigen::MatrixXd> x, int r, int c, double v) { x(r,c) += v; });
  98. // Return mutable references (numpy maps into eigen varibles)
  99. m.def("get_cm_ref", []() { return Eigen::Ref<Eigen::MatrixXd>(get_cm()); });
  100. m.def("get_rm_ref", []() { return Eigen::Ref<MatrixXdR>(get_rm()); });
  101. // The same references, but non-mutable (numpy maps into eigen variables, but is !writeable)
  102. m.def("get_cm_const_ref", []() { return Eigen::Ref<const Eigen::MatrixXd>(get_cm()); });
  103. m.def("get_rm_const_ref", []() { return Eigen::Ref<const MatrixXdR>(get_rm()); });
  104. // Just the corners (via a Map instead of a Ref):
  105. m.def("get_cm_corners", []() {
  106. auto &x = get_cm();
  107. return py::EigenDMap<Eigen::Matrix2d>(
  108. x.data(),
  109. py::EigenDStride(x.outerStride() * (x.rows() - 1), x.innerStride() * (x.cols() - 1)));
  110. });
  111. m.def("get_cm_corners_const", []() {
  112. const auto &x = get_cm();
  113. return py::EigenDMap<const Eigen::Matrix2d>(
  114. x.data(),
  115. py::EigenDStride(x.outerStride() * (x.rows() - 1), x.innerStride() * (x.cols() - 1)));
  116. });
  117. m.def("reset_refs", reset_refs); // Restores get_{cm,rm}_ref to original values
  118. // Increments and returns ref to (same) matrix
  119. m.def("incr_matrix", [](Eigen::Ref<Eigen::MatrixXd> m, double v) {
  120. m += Eigen::MatrixXd::Constant(m.rows(), m.cols(), v);
  121. return m;
  122. }, py::return_value_policy::reference);
  123. // Same, but accepts a matrix of any strides
  124. m.def("incr_matrix_any", [](py::EigenDRef<Eigen::MatrixXd> m, double v) {
  125. m += Eigen::MatrixXd::Constant(m.rows(), m.cols(), v);
  126. return m;
  127. }, py::return_value_policy::reference);
  128. // Returns an eigen slice of even rows
  129. m.def("even_rows", [](py::EigenDRef<Eigen::MatrixXd> m) {
  130. return py::EigenDMap<Eigen::MatrixXd>(
  131. m.data(), (m.rows() + 1) / 2, m.cols(),
  132. py::EigenDStride(m.outerStride(), 2 * m.innerStride()));
  133. }, py::return_value_policy::reference);
  134. // Returns an eigen slice of even columns
  135. m.def("even_cols", [](py::EigenDRef<Eigen::MatrixXd> m) {
  136. return py::EigenDMap<Eigen::MatrixXd>(
  137. m.data(), m.rows(), (m.cols() + 1) / 2,
  138. py::EigenDStride(2 * m.outerStride(), m.innerStride()));
  139. }, py::return_value_policy::reference);
  140. // Returns diagonals: a vector-like object with an inner stride != 1
  141. m.def("diagonal", [](const Eigen::Ref<const Eigen::MatrixXd> &x) { return x.diagonal(); });
  142. m.def("diagonal_1", [](const Eigen::Ref<const Eigen::MatrixXd> &x) { return x.diagonal<1>(); });
  143. m.def("diagonal_n", [](const Eigen::Ref<const Eigen::MatrixXd> &x, int index) { return x.diagonal(index); });
  144. // Return a block of a matrix (gives non-standard strides)
  145. m.def("block", [](const Eigen::Ref<const Eigen::MatrixXd> &x, int start_row, int start_col, int block_rows, int block_cols) {
  146. return x.block(start_row, start_col, block_rows, block_cols);
  147. });
  148. // return value referencing/copying tests:
  149. class ReturnTester {
  150. Eigen::MatrixXd mat = create();
  151. public:
  152. ReturnTester() { print_created(this); }
  153. ~ReturnTester() { print_destroyed(this); }
  154. static Eigen::MatrixXd create() { return Eigen::MatrixXd::Ones(10, 10); }
  155. static const Eigen::MatrixXd createConst() { return Eigen::MatrixXd::Ones(10, 10); }
  156. Eigen::MatrixXd &get() { return mat; }
  157. Eigen::MatrixXd *getPtr() { return &mat; }
  158. const Eigen::MatrixXd &view() { return mat; }
  159. const Eigen::MatrixXd *viewPtr() { return &mat; }
  160. Eigen::Ref<Eigen::MatrixXd> ref() { return mat; }
  161. Eigen::Ref<const Eigen::MatrixXd> refConst() { return mat; }
  162. Eigen::Block<Eigen::MatrixXd> block(int r, int c, int nrow, int ncol) { return mat.block(r, c, nrow, ncol); }
  163. Eigen::Block<const Eigen::MatrixXd> blockConst(int r, int c, int nrow, int ncol) const { return mat.block(r, c, nrow, ncol); }
  164. py::EigenDMap<Eigen::Matrix2d> corners() { return py::EigenDMap<Eigen::Matrix2d>(mat.data(),
  165. py::EigenDStride(mat.outerStride() * (mat.outerSize()-1), mat.innerStride() * (mat.innerSize()-1))); }
  166. py::EigenDMap<const Eigen::Matrix2d> cornersConst() const { return py::EigenDMap<const Eigen::Matrix2d>(mat.data(),
  167. py::EigenDStride(mat.outerStride() * (mat.outerSize()-1), mat.innerStride() * (mat.innerSize()-1))); }
  168. };
  169. using rvp = py::return_value_policy;
  170. py::class_<ReturnTester>(m, "ReturnTester")
  171. .def(py::init<>())
  172. .def_static("create", &ReturnTester::create)
  173. .def_static("create_const", &ReturnTester::createConst)
  174. .def("get", &ReturnTester::get, rvp::reference_internal)
  175. .def("get_ptr", &ReturnTester::getPtr, rvp::reference_internal)
  176. .def("view", &ReturnTester::view, rvp::reference_internal)
  177. .def("view_ptr", &ReturnTester::view, rvp::reference_internal)
  178. .def("copy_get", &ReturnTester::get) // Default rvp: copy
  179. .def("copy_view", &ReturnTester::view) // "
  180. .def("ref", &ReturnTester::ref) // Default for Ref is to reference
  181. .def("ref_const", &ReturnTester::refConst) // Likewise, but const
  182. .def("ref_safe", &ReturnTester::ref, rvp::reference_internal)
  183. .def("ref_const_safe", &ReturnTester::refConst, rvp::reference_internal)
  184. .def("copy_ref", &ReturnTester::ref, rvp::copy)
  185. .def("copy_ref_const", &ReturnTester::refConst, rvp::copy)
  186. .def("block", &ReturnTester::block)
  187. .def("block_safe", &ReturnTester::block, rvp::reference_internal)
  188. .def("block_const", &ReturnTester::blockConst, rvp::reference_internal)
  189. .def("copy_block", &ReturnTester::block, rvp::copy)
  190. .def("corners", &ReturnTester::corners, rvp::reference_internal)
  191. .def("corners_const", &ReturnTester::cornersConst, rvp::reference_internal)
  192. ;
  193. // Returns a DiagonalMatrix with diagonal (1,2,3,...)
  194. m.def("incr_diag", [](int k) {
  195. Eigen::DiagonalMatrix<int, Eigen::Dynamic> m(k);
  196. for (int i = 0; i < k; i++) m.diagonal()[i] = i+1;
  197. return m;
  198. });
  199. // Returns a SelfAdjointView referencing the lower triangle of m
  200. m.def("symmetric_lower", [](const Eigen::MatrixXi &m) {
  201. return m.selfadjointView<Eigen::Lower>();
  202. });
  203. // Returns a SelfAdjointView referencing the lower triangle of m
  204. m.def("symmetric_upper", [](const Eigen::MatrixXi &m) {
  205. return m.selfadjointView<Eigen::Upper>();
  206. });
  207. // Test matrix for various functions below.
  208. Eigen::MatrixXf mat(5, 6);
  209. mat << 0, 3, 0, 0, 0, 11,
  210. 22, 0, 0, 0, 17, 11,
  211. 7, 5, 0, 1, 0, 11,
  212. 0, 0, 0, 0, 0, 11,
  213. 0, 0, 14, 0, 8, 11;
  214. m.def("fixed_r", [mat]() -> FixedMatrixR { return FixedMatrixR(mat); });
  215. m.def("fixed_r_const", [mat]() -> const FixedMatrixR { return FixedMatrixR(mat); });
  216. m.def("fixed_c", [mat]() -> FixedMatrixC { return FixedMatrixC(mat); });
  217. m.def("fixed_copy_r", [](const FixedMatrixR &m) -> FixedMatrixR { return m; });
  218. m.def("fixed_copy_c", [](const FixedMatrixC &m) -> FixedMatrixC { return m; });
  219. m.def("fixed_mutator_r", [](Eigen::Ref<FixedMatrixR>) {});
  220. m.def("fixed_mutator_c", [](Eigen::Ref<FixedMatrixC>) {});
  221. m.def("fixed_mutator_a", [](py::EigenDRef<FixedMatrixC>) {});
  222. m.def("dense_r", [mat]() -> DenseMatrixR { return DenseMatrixR(mat); });
  223. m.def("dense_c", [mat]() -> DenseMatrixC { return DenseMatrixC(mat); });
  224. m.def("dense_copy_r", [](const DenseMatrixR &m) -> DenseMatrixR { return m; });
  225. m.def("dense_copy_c", [](const DenseMatrixC &m) -> DenseMatrixC { return m; });
  226. m.def("sparse_r", [mat]() -> SparseMatrixR { return Eigen::SparseView<Eigen::MatrixXf>(mat); });
  227. m.def("sparse_c", [mat]() -> SparseMatrixC { return Eigen::SparseView<Eigen::MatrixXf>(mat); });
  228. m.def("sparse_copy_r", [](const SparseMatrixR &m) -> SparseMatrixR { return m; });
  229. m.def("sparse_copy_c", [](const SparseMatrixC &m) -> SparseMatrixC { return m; });
  230. m.def("partial_copy_four_rm_r", [](const FourRowMatrixR &m) -> FourRowMatrixR { return m; });
  231. m.def("partial_copy_four_rm_c", [](const FourColMatrixR &m) -> FourColMatrixR { return m; });
  232. m.def("partial_copy_four_cm_r", [](const FourRowMatrixC &m) -> FourRowMatrixC { return m; });
  233. m.def("partial_copy_four_cm_c", [](const FourColMatrixC &m) -> FourColMatrixC { return m; });
  234. // Test that we can cast a numpy object to a Eigen::MatrixXd explicitly
  235. m.def("cpp_copy", [](py::handle m) { return m.cast<Eigen::MatrixXd>()(1, 0); });
  236. m.def("cpp_ref_c", [](py::handle m) { return m.cast<Eigen::Ref<Eigen::MatrixXd>>()(1, 0); });
  237. m.def("cpp_ref_r", [](py::handle m) { return m.cast<Eigen::Ref<MatrixXdR>>()(1, 0); });
  238. m.def("cpp_ref_any", [](py::handle m) { return m.cast<py::EigenDRef<Eigen::MatrixXd>>()(1, 0); });
  239. // Test that we can prevent copying into an argument that would normally copy: First a version
  240. // that would allow copying (if types or strides don't match) for comparison:
  241. m.def("get_elem", &get_elem);
  242. // Now this alternative that calls the tells pybind to fail rather than copy:
  243. m.def("get_elem_nocopy", [](Eigen::Ref<const Eigen::MatrixXd> m) -> double { return get_elem(m); },
  244. py::arg().noconvert());
  245. // Also test a row-major-only no-copy const ref:
  246. m.def("get_elem_rm_nocopy", [](Eigen::Ref<const Eigen::Matrix<long, -1, -1, Eigen::RowMajor>> &m) -> long { return m(2, 1); },
  247. py::arg().noconvert());
  248. // Issue #738: 1xN or Nx1 2D matrices were neither accepted nor properly copied with an
  249. // incompatible stride value on the length-1 dimension--but that should be allowed (without
  250. // requiring a copy!) because the stride value can be safely ignored on a size-1 dimension.
  251. m.def("iss738_f1", &adjust_matrix<const Eigen::Ref<const Eigen::MatrixXd> &>, py::arg().noconvert());
  252. m.def("iss738_f2", &adjust_matrix<const Eigen::Ref<const Eigen::Matrix<double, -1, -1, Eigen::RowMajor>> &>, py::arg().noconvert());
  253. // Make sure named arguments are working properly:
  254. m.def("matrix_multiply", [](const py::EigenDRef<const Eigen::MatrixXd> A, const py::EigenDRef<const Eigen::MatrixXd> B)
  255. -> Eigen::MatrixXd {
  256. if (A.cols() != B.rows()) throw std::domain_error("Nonconformable matrices!");
  257. return A * B;
  258. }, py::arg("A"), py::arg("B"));
  259. py::class_<CustomOperatorNew>(m, "CustomOperatorNew")
  260. .def(py::init<>())
  261. .def_readonly("a", &CustomOperatorNew::a)
  262. .def_readonly("b", &CustomOperatorNew::b);
  263. // test_eigen_ref_life_support
  264. // In case of a failure (the caster's temp array does not live long enough), creating
  265. // a new array (np.ones(10)) increases the chances that the temp array will be garbage
  266. // collected and/or that its memory will be overridden with different values.
  267. m.def("get_elem_direct", [](Eigen::Ref<const Eigen::VectorXd> v) {
  268. py::module::import("numpy").attr("ones")(10);
  269. return v(5);
  270. });
  271. m.def("get_elem_indirect", [](std::vector<Eigen::Ref<const Eigen::VectorXd>> v) {
  272. py::module::import("numpy").attr("ones")(10);
  273. return v[0](5);
  274. });
  275. });