Skip to content

Commit

Permalink
[eigen] Test overloads for dense vs. sparse
Browse files Browse the repository at this point in the history
  • Loading branch information
EricCousineau-TRI committed Nov 10, 2023
1 parent c1b0cdd commit 1ed181b
Show file tree
Hide file tree
Showing 3 changed files with 60 additions and 1 deletion.
3 changes: 2 additions & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -480,7 +480,8 @@ jobs:
- uses: actions/checkout@v3

- name: Add Python 3
run: apt-get update; apt-get install -y python3-dev python3-numpy python3-pytest python3-pip libeigen3-dev
# NOTE(drake): We explicitly test scipy here.
run: apt-get update; apt-get install -y python3-dev python3-numpy python3-pytest python3-pip libeigen3-dev python3-scipy

- name: Update pip
run: python3 -m pip install --upgrade pip
Expand Down
36 changes: 36 additions & 0 deletions tests/test_eigen_matrix.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,12 @@ typedef Eigen::Matrix<ADScalar, 5, 1> Vector5ADScalar;
typedef Eigen::Matrix<ADScalar, 1, 6> Vector6ADScalarR;
PYBIND11_NUMPY_OBJECT_DTYPE(ADScalar);

struct MyScalar {
int value{};
};
PYBIND11_NUMPY_OBJECT_DTYPE(MyScalar);
using VectorXMyScalar = Eigen::Matrix<MyScalar, Eigen::Dynamic, 1>;

// Sets/resets a testing reference matrix to have values of 10*r + c, where r and c are the
// (1-based) row/column number.
template <typename M>
Expand Down Expand Up @@ -516,4 +522,34 @@ TEST_SUBMODULE(eigen_matrix, m) {
py::module_::import("numpy").attr("ones")(10);
return v[0](5);
});

// test_eigen_dense_sparse_overload
// https://github.com/RobotLocomotion/drake/issues/20516
// Certain argument combinations seem to make overloads go sideways.
py::class_<MyScalar>(m, "MyScalar")
.def(py::init<int>(), py::arg("value"))
.def_readwrite("value", &MyScalar::value);

m.def(
"accept_matrix",
[](const Eigen::MatrixXd&) { return "dense"; });
m.def(
"accept_matrix",
[](const Eigen::SparseMatrix<double>&) { return "sparse"; });
m.def(
"accept_matrix_and_dtype_object",
[](const Eigen::MatrixXd&, const VectorXADScalar&) { return "dense"; });
m.def(
"accept_matrix_and_dtype_object",
[](const Eigen::SparseMatrix<double>&, const VectorXADScalar&) {
return "sparse";
});
m.def(
"accept_matrix_and_dtype_object",
[](const Eigen::MatrixXd&, const VectorXMyScalar&) { return "dense"; });
m.def(
"accept_matrix_and_dtype_object",
[](const Eigen::SparseMatrix<double>&, const VectorXMyScalar&) {
return "sparse";
});
}
22 changes: 22 additions & 0 deletions tests/test_eigen_matrix.py
Original file line number Diff line number Diff line change
Expand Up @@ -903,3 +903,25 @@ def test_custom_operator_new():
o = m.CustomOperatorNew()
np.testing.assert_allclose(o.a, 0.0)
np.testing.assert_allclose(o.b.diagonal(), 1.0)


def test_eigen_dense_sparse_overload():
"""Overloads that should prefer dense vs. sparse should work as
expected, even when combining with arguments like dtype=object."""

pytest.importorskip("scipy")
import scipy.sparse

A_dense = np.eye(2)
A_sparse = scipy.sparse.csc_matrix(np.eye(2))

assert m.accept_matrix(A_dense) == "dense"
assert m.accept_matrix(A_sparse) == "sparse"

x_ad = float_to_adscalar(np.ones(2), deriv=[1.0])
assert m.accept_matrix_and_dtype_object(A_dense, x_ad) == "dense"
assert m.accept_matrix_and_dtype_object(A_sparse, x_ad) == "sparse"

x_myscalar = np.array([m.MyScalar(1), m.MyScalar(2)])
assert m.accept_matrix_and_dtype_object(A_dense, x_myscalar) == "dense"
assert m.accept_matrix_and_dtype_object(A_sparse, x_myscalar) == "sparse"

0 comments on commit 1ed181b

Please sign in to comment.