Skip to content

Commit dfef38d

Browse files
jagermanjagerman
committed
Added eigen<->numpy referencing tests
1 parent cc3946e commit dfef38d

File tree

2 files changed

+612
-81
lines changed

2 files changed

+612
-81
lines changed

‎tests/test_eigen.cpp

Lines changed: 183 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -8,29 +8,96 @@
88
*/
99

1010
#include "pybind11_tests.h"
11+
#include "constructor_stats.h"
1112
#include <pybind11/eigen.h>
1213
#include <Eigen/Cholesky>
1314

14-
Eigen::VectorXf double_col(const Eigen::VectorXf& x)
15-
{ return 2.0f * x; }
15+
using MatrixXdR = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
1616

17-
Eigen::RowVectorXf double_row(const Eigen::RowVectorXf& x)
18-
{ return 2.0f * x; }
17+
// Different ways of passing via Eigen::Ref; the first and second are the Eigen-recommended
18+
Eigen::MatrixXd cholesky1(Eigen::Ref<MatrixXdR> x) { return x.llt().matrixL(); }
19+
Eigen::MatrixXd cholesky2(const Eigen::Ref<const MatrixXdR> &x) { return x.llt().matrixL(); }
20+
Eigen::MatrixXd cholesky3(const Eigen::Ref<MatrixXdR> &x) { return x.llt().matrixL(); }
21+
Eigen::MatrixXd cholesky4(Eigen::Ref<const MatrixXdR> x) { return x.llt().matrixL(); }
1922

20-
Eigen::MatrixXf double_mat_cm(const Eigen::MatrixXf& x)
21-
{ return 2.0f * x; }
23+
// Mutators: these add some value to the given element using Eigen, but Eigen should be mapping into
24+
// the numpy array data and so the result should show up there. There are three versions: one that
25+
// works on a contiguous-row matrix (numpy's default), one for a contiguous-column matrix, and one
26+
// for any matrix.
27+
void add_rm(Eigen::Ref<MatrixXdR> x, int r, int c, double v) { x(r,c) += v; }
28+
void add_cm(Eigen::Ref<Eigen::MatrixXd> x, int r, int c, double v) { x(r,c) += v; }
29+
void add_any(Eigen::Ref<Eigen::MatrixXd, 0, Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>> x, int r, int c, double v) { x(r,c) += v; }
2230

23-
// Different ways of passing via Eigen::Ref; the first and second are the Eigen-recommended
24-
Eigen::MatrixXd cholesky1(Eigen::Ref<Eigen::MatrixXd> &x) { return x.llt().matrixL(); }
25-
Eigen::MatrixXd cholesky2(const Eigen::Ref<const Eigen::MatrixXd> &x) { return x.llt().matrixL(); }
26-
Eigen::MatrixXd cholesky3(const Eigen::Ref<Eigen::MatrixXd> &x) { return x.llt().matrixL(); }
27-
Eigen::MatrixXd cholesky4(Eigen::Ref<const Eigen::MatrixXd> &x) { return x.llt().matrixL(); }
28-
Eigen::MatrixXd cholesky5(Eigen::Ref<Eigen::MatrixXd> x) { return x.llt().matrixL(); }
29-
Eigen::MatrixXd cholesky6(Eigen::Ref<const Eigen::MatrixXd> x) { return x.llt().matrixL(); }
31+
template <typename M> void reset_ref(M &x) {
32+
for (int i = 0; i < x.rows(); i++) for (int j = 0; j < x.cols(); j++)
33+
x(i, j) = 11 + 10*i + j;
34+
}
35+
36+
Eigen::MatrixXd &get_cm() {
37+
static Eigen::MatrixXd *x;
38+
if (!x) {
39+
x = new Eigen::MatrixXd(3, 3);
40+
reset_ref(*x);
41+
}
42+
return *x;
43+
}
44+
MatrixXdR &get_rm() {
45+
static MatrixXdR *x;
46+
if (!x) {
47+
x = new MatrixXdR(3, 3);
48+
reset_ref(*x);
49+
}
50+
return *x;
51+
}
52+
53+
// Returns a slice of get_cm matrix:
54+
py::EigenDMap<Eigen::Matrix2d> get_cm_corners() {
55+
auto &x = get_cm();
56+
return py::EigenDMap<Eigen::Matrix2d>(
57+
x.data(),
58+
py::EigenDStride(x.outerStride() * (x.rows() - 1), x.innerStride() * (x.cols() - 1)));
59+
}
60+
61+
py::EigenDMap<const Eigen::Matrix2d> get_cm_corners_const() {
62+
const auto &x = get_cm();
63+
return py::EigenDMap<const Eigen::Matrix2d>(
64+
x.data(),
65+
py::EigenDStride(x.outerStride() * (x.rows() - 1), x.innerStride() * (x.cols() - 1)));
66+
}
3067

31-
typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> MatrixXfRowMajor;
32-
MatrixXfRowMajor double_mat_rm(const MatrixXfRowMajor& x)
33-
{ return 2.0f * x; }
68+
void reset_refs() {
69+
reset_ref(get_cm());
70+
reset_ref(get_rm());
71+
}
72+
73+
// These return a Eigen::Ref<> that numpy is allowed to mutate to change Eigen's values
74+
Eigen::Ref<Eigen::MatrixXd> get_cm_ref() { return get_cm(); }
75+
Eigen::Ref<MatrixXdR> get_rm_ref() { return get_rm(); }
76+
// These return a Eigen::Ref<const> that should become read-only on the numpy side
77+
Eigen::Ref<const Eigen::MatrixXd> get_cm_const_ref() { return get_cm(); }
78+
Eigen::Ref<const MatrixXdR> get_rm_const_ref() { return get_rm(); }
79+
80+
// Increments all elements in a matrix by some value then returns the same reference
81+
Eigen::Ref<Eigen::MatrixXd> incr_matrix(Eigen::Ref<Eigen::MatrixXd> m, double v) {
82+
m += Eigen::MatrixXd::Constant(m.rows(), m.cols(), v);
83+
return m;
84+
}
85+
// Same, but takes any strides
86+
Eigen::Ref<Eigen::MatrixXd, 0, Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>> incr_matrix_any(Eigen::Ref<Eigen::MatrixXd, 0, Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>> m, double v) {
87+
m += Eigen::MatrixXd::Constant(m.rows(), m.cols(), v);
88+
return m;
89+
}
90+
91+
// Returns an eigen slice of even rows
92+
Eigen::Ref<Eigen::MatrixXd, 0, Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>> even_rows(Eigen::Ref<Eigen::MatrixXd, 0, Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>> m) {
93+
return Eigen::Map<Eigen::MatrixXd, 0, Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>>(m.data(), (m.rows() + 1) / 2, m.cols(),
94+
Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>(m.outerStride(), 2 * m.innerStride()));
95+
}
96+
// Returns an eigen slice of even columns
97+
Eigen::Ref<Eigen::MatrixXd, 0, Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>> even_cols(Eigen::Ref<Eigen::MatrixXd, 0, Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>> m) {
98+
return Eigen::Map<Eigen::MatrixXd, 0, Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>>(m.data(), m.rows(), (m.cols() + 1) / 2,
99+
Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic>(2 * m.outerStride(), m.innerStride()));
100+
}
34101

35102
test_initializer eigen([](py::module &m) {
36103
typedef Eigen::Matrix<float, 5, 6, Eigen::RowMajor> FixedMatrixR;
@@ -51,16 +118,46 @@ test_initializer eigen([](py::module &m) {
51118
mat << 0, 3, 0, 0, 0, 11, 22, 0, 0, 0, 17, 11, 7, 5, 0, 1, 0, 11, 0,
52119
0, 0, 0, 0, 11, 0, 0, 14, 0, 8, 11;
53120

54-
m.def("double_col", &double_col);
55-
m.def("double_row", &double_row);
56-
m.def("double_mat_cm", &double_mat_cm);
57-
m.def("double_mat_rm", &double_mat_rm);
121+
m.def("double_col", [](const Eigen::VectorXf &x) -> Eigen::VectorXf { return 2.0f * x; });
122+
m.def("double_row", [](const Eigen::RowVectorXf &x) -> Eigen::RowVectorXf { return 2.0f * x; });
123+
m.def("double_threec", [](py::EigenDRef<Eigen::Vector3f> x) { x *= 2; });
124+
m.def("double_threer", [](py::EigenDRef<Eigen::RowVector3f> x) { x *= 2; });
125+
m.def("double_mat_cm", [](Eigen::MatrixXf x) -> Eigen::MatrixXf { return 2.0f * x; });
126+
m.def("double_mat_rm", [](DenseMatrixR x) -> DenseMatrixR { return 2.0f * x; });
58127
m.def("cholesky1", &cholesky1);
59128
m.def("cholesky2", &cholesky2);
60129
m.def("cholesky3", &cholesky3);
61130
m.def("cholesky4", &cholesky4);
62-
m.def("cholesky5", &cholesky5);
63-
m.def("cholesky6", &cholesky6);
131+
132+
// Mutators (Eigen maps into numpy variables):
133+
m.def("add_rm", add_rm); // Only takes row-contiguous
134+
m.def("add_cm", add_cm); // Only takes column-contiguous
135+
m.def("add_any", add_any); // Accepts anything
136+
// Overloaded versions that will accept either row or column contiguous:
137+
m.def("add1", add_rm);
138+
m.def("add1", add_cm);
139+
m.def("add2", add_cm);
140+
m.def("add2", add_rm);
141+
142+
// Return mutable references (numpy maps into eigen varibles)
143+
m.def("get_cm_ref", []() { return Eigen::Ref<Eigen::MatrixXd>(get_cm()); });
144+
m.def("get_rm_ref", []() { return Eigen::Ref<MatrixXdR>(get_rm()); });
145+
// The same references, but non-mutable (numpy maps into eigen variables, but is !writeable)
146+
m.def("get_cm_const_ref", []() { return Eigen::Ref<const Eigen::MatrixXd>(get_cm()); });
147+
m.def("get_rm_const_ref", []() { return Eigen::Ref<const MatrixXdR>(get_rm()); });
148+
// Just the corners (via a Map instead of a Ref):
149+
m.def("get_cm_corners", get_cm_corners);
150+
m.def("get_cm_corners_const", get_cm_corners_const);
151+
152+
m.def("reset_refs", reset_refs); // Restores get_{cm,rm}_ref to original values
153+
154+
// Increments and returns ref to matrix
155+
m.def("incr_matrix", incr_matrix, py::return_value_policy::reference);
156+
m.def("incr_matrix_any", incr_matrix_any, py::return_value_policy::reference);
157+
158+
// Even row/column slicing
159+
m.def("even_rows", even_rows, py::return_value_policy::reference);
160+
m.def("even_cols", even_cols, py::return_value_policy::reference);
64161

65162
// Returns diagonals: a vector-like object with an inner stride != 1
66163
m.def("diagonal", [](const Eigen::Ref<const Eigen::MatrixXd> &x) { return x.diagonal(); });
@@ -72,6 +169,52 @@ test_initializer eigen([](py::module &m) {
72169
return x.block(start_row, start_col, block_rows, block_cols);
73170
});
74171

172+
// return value referencing/copying tests:
173+
class ReturnTester {
174+
Eigen::MatrixXd mat = create();
175+
public:
176+
ReturnTester() { print_created(this); }
177+
~ReturnTester() { print_destroyed(this); }
178+
static Eigen::MatrixXd create() { return Eigen::MatrixXd::Ones(10, 10); }
179+
static const Eigen::MatrixXd createConst() { return Eigen::MatrixXd::Ones(10, 10); }
180+
Eigen::MatrixXd &get() { return mat; }
181+
Eigen::MatrixXd *getPtr() { return &mat; }
182+
const Eigen::MatrixXd &view() { return mat; }
183+
const Eigen::MatrixXd *viewPtr() { return &mat; }
184+
Eigen::Ref<Eigen::MatrixXd> ref() { return mat; }
185+
Eigen::Ref<const Eigen::MatrixXd> refConst() { return mat; }
186+
Eigen::Block<Eigen::MatrixXd> block(int r, int c, int nrow, int ncol) { return mat.block(r, c, nrow, ncol); }
187+
Eigen::Block<const Eigen::MatrixXd> blockConst(int r, int c, int nrow, int ncol) const { return mat.block(r, c, nrow, ncol); }
188+
py::EigenDMap<Eigen::Matrix2d> corners() { return py::EigenDMap<Eigen::Matrix2d>(mat.data(),
189+
py::EigenDStride(mat.outerStride() * (mat.outerSize()-1), mat.innerStride() * (mat.innerSize()-1))); }
190+
py::EigenDMap<const Eigen::Matrix2d> cornersConst() const { return py::EigenDMap<const Eigen::Matrix2d>(mat.data(),
191+
py::EigenDStride(mat.outerStride() * (mat.outerSize()-1), mat.innerStride() * (mat.innerSize()-1))); }
192+
};
193+
using rvp = py::return_value_policy;
194+
py::class_<ReturnTester>(m, "ReturnTester")
195+
.def(py::init<>())
196+
.def_static("create", &ReturnTester::create)
197+
.def_static("create_const", &ReturnTester::createConst)
198+
.def("get", &ReturnTester::get, rvp::reference_internal)
199+
.def("get_ptr", &ReturnTester::getPtr, rvp::reference_internal)
200+
.def("view", &ReturnTester::view, rvp::reference_internal)
201+
.def("view_ptr", &ReturnTester::view, rvp::reference_internal)
202+
.def("copy_get", &ReturnTester::get) // Default rvp: copy
203+
.def("copy_view", &ReturnTester::view) // "
204+
.def("ref", &ReturnTester::ref) // Default for Ref is to reference
205+
.def("ref_const", &ReturnTester::refConst) // Likewise, but const
206+
.def("ref_safe", &ReturnTester::ref, rvp::reference_internal)
207+
.def("ref_const_safe", &ReturnTester::refConst, rvp::reference_internal)
208+
.def("copy_ref", &ReturnTester::ref, rvp::copy)
209+
.def("copy_ref_const", &ReturnTester::refConst, rvp::copy)
210+
.def("block", &ReturnTester::block)
211+
.def("block_safe", &ReturnTester::block, rvp::reference_internal)
212+
.def("block_const", &ReturnTester::blockConst, rvp::reference_internal)
213+
.def("copy_block", &ReturnTester::block, rvp::copy)
214+
.def("corners", &ReturnTester::corners, rvp::reference_internal)
215+
.def("corners_const", &ReturnTester::cornersConst, rvp::reference_internal)
216+
;
217+
75218
// Returns a DiagonalMatrix with diagonal (1,2,3,...)
76219
m.def("incr_diag", [](int k) {
77220
Eigen::DiagonalMatrix<int, Eigen::Dynamic> m(k);
@@ -92,18 +235,26 @@ test_initializer eigen([](py::module &m) {
92235
return FixedMatrixR(mat);
93236
});
94237

238+
m.def("fixed_r_const", [mat]() -> const FixedMatrixR {
239+
return FixedMatrixR(mat);
240+
});
241+
95242
m.def("fixed_c", [mat]() -> FixedMatrixC {
96243
return FixedMatrixC(mat);
97244
});
98245

99-
m.def("fixed_passthrough_r", [](const FixedMatrixR &m) -> FixedMatrixR {
246+
m.def("fixed_copy_r", [](const FixedMatrixR &m) -> FixedMatrixR {
100247
return m;
101248
});
102249

103-
m.def("fixed_passthrough_c", [](const FixedMatrixC &m) -> FixedMatrixC {
250+
m.def("fixed_copy_c", [](const FixedMatrixC &m) -> FixedMatrixC {
104251
return m;
105252
});
106253

254+
m.def("fixed_mutator_r", [](Eigen::Ref<FixedMatrixR>) {});
255+
m.def("fixed_mutator_c", [](Eigen::Ref<FixedMatrixC>) {});
256+
m.def("fixed_mutator_a", [](py::EigenDRef<FixedMatrixC>) {});
257+
107258
m.def("dense_r", [mat]() -> DenseMatrixR {
108259
return DenseMatrixR(mat);
109260
});
@@ -112,11 +263,11 @@ test_initializer eigen([](py::module &m) {
112263
return DenseMatrixC(mat);
113264
});
114265

115-
m.def("dense_passthrough_r", [](const DenseMatrixR &m) -> DenseMatrixR {
266+
m.def("dense_copy_r", [](const DenseMatrixR &m) -> DenseMatrixR {
116267
return m;
117268
});
118269

119-
m.def("dense_passthrough_c", [](const DenseMatrixC &m) -> DenseMatrixC {
270+
m.def("dense_copy_c", [](const DenseMatrixC &m) -> DenseMatrixC {
120271
return m;
121272
});
122273

@@ -128,24 +279,24 @@ test_initializer eigen([](py::module &m) {
128279
return Eigen::SparseView<Eigen::MatrixXf>(mat);
129280
});
130281

131-
m.def("sparse_passthrough_r", [](const SparseMatrixR &m) -> SparseMatrixR {
282+
m.def("sparse_copy_r", [](const SparseMatrixR &m) -> SparseMatrixR {
132283
return m;
133284
});
134285

135-
m.def("sparse_passthrough_c", [](const SparseMatrixC &m) -> SparseMatrixC {
286+
m.def("sparse_copy_c", [](const SparseMatrixC &m) -> SparseMatrixC {
136287
return m;
137288
});
138289

139-
m.def("partial_passthrough_four_rm_r", [](const FourRowMatrixR &m) -> FourRowMatrixR {
290+
m.def("partial_copy_four_rm_r", [](const FourRowMatrixR &m) -> FourRowMatrixR {
140291
return m;
141292
});
142-
m.def("partial_passthrough_four_rm_c", [](const FourColMatrixR &m) -> FourColMatrixR {
293+
m.def("partial_copy_four_rm_c", [](const FourColMatrixR &m) -> FourColMatrixR {
143294
return m;
144295
});
145-
m.def("partial_passthrough_four_cm_r", [](const FourRowMatrixC &m) -> FourRowMatrixC {
296+
m.def("partial_copy_four_cm_r", [](const FourRowMatrixC &m) -> FourRowMatrixC {
146297
return m;
147298
});
148-
m.def("partial_passthrough_four_cm_c", [](const FourColMatrixC &m) -> FourColMatrixC {
299+
m.def("partial_copy_four_cm_c", [](const FourColMatrixC &m) -> FourColMatrixC {
149300
return m;
150301
});
151302
});

0 commit comments

Comments
 (0)