Merge pull request #26 from EricCousineau-TRI/issue/drake_9886

dtype object: Ensure fixed-size vectors are indexed correctly

Author: EricCousineau-TRI

include/pybind11/eigen.h

@@ -248,8 +248,9 @@ template <typename props> handle eigen_array_cast(typename props::Type const &sr
                 nullptr,
                 empty_base
             );
+            constexpr bool is_row = props::fixed_rows && props::rows == 1;
             for (ssize_t i = 0; i < src.size(); ++i) {
-                const Scalar src_val = props::fixed_rows ? src(0, i) : src(i, 0);
+                const Scalar src_val = is_row ? src(0, i) : src(i, 0);
                 auto value_ = reinterpret_steal<object>(make_caster<Scalar>::cast(src_val, policy, empty_base));
                 if (!value_)
                     return handle();

tests/test_eigen.cpp

@@ -28,6 +28,8 @@ using MatrixX = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>;
 
 typedef Eigen::Matrix<ADScalar, Eigen::Dynamic, 1> VectorXADScalar;
 typedef Eigen::Matrix<ADScalar, 1, Eigen::Dynamic> VectorXADScalarR;
+typedef Eigen::Matrix<ADScalar, 5, 1> Vector5ADScalar;
+typedef Eigen::Matrix<ADScalar, 1, 6> Vector6ADScalarR;
 PYBIND11_NUMPY_OBJECT_DTYPE(ADScalar);
 
 // Sets/resets a testing reference matrix to have values of 10*r + c, where r and c are the
@@ -112,14 +114,17 @@ TEST_SUBMODULE(eigen, m) {
     // various tests
     m.def("double_col", [](const Eigen::VectorXf &x) -> Eigen::VectorXf { return 2.0f * x; });
     m.def("double_adscalar_col", [](const VectorXADScalar &x) -> VectorXADScalar { return 2.0f * x; });
+    m.def("double_adscalar_col5", [](const Vector5ADScalar &x) -> Vector5ADScalar { return 2.0f * x; });
     m.def("double_row", [](const Eigen::RowVectorXf &x) -> Eigen::RowVectorXf { return 2.0f * x; });
     m.def("double_adscalar_row", [](const VectorXADScalarR &x) -> VectorXADScalarR { return 2.0f * x; });
+    m.def("double_adscalar_row6", [](const Vector6ADScalarR &x) -> Vector6ADScalarR { return 2.0f * x; });
     m.def("double_complex", [](const Eigen::VectorXcf &x) -> Eigen::VectorXcf { return 2.0f * x; });
     m.def("double_threec", [](py::EigenDRef<Eigen::Vector3f> x) { x *= 2; });
     m.def("double_threer", [](py::EigenDRef<Eigen::RowVector3f> x) { x *= 2; });
     m.def("double_mat_cm", [](Eigen::MatrixXf x) -> Eigen::MatrixXf { return 2.0f * x; });
     m.def("double_mat_rm", [](DenseMatrixR x) -> DenseMatrixR { return 2.0f * x; });
 
+
     // test_eigen_ref_to_python
     // Different ways of passing via Eigen::Ref; the first and second are the Eigen-recommended
     m.def("cholesky1", [](Eigen::Ref<MatrixXdR> x) -> Eigen::MatrixXd { return x.llt().matrixL(); });
@@ -270,8 +275,20 @@ TEST_SUBMODULE(eigen, m) {
         .def("value", [](const ADScalar & self) {
           return self.value();
         })
+        .def("__repr__", [](const ADScalar& self) {
+          return py::str("<ADScalar {} deriv={}>").format(self.value(), self.derivatives());
+        })
         ;
 
+    m.def("equal_to", [](double a, double b) { return a == b; });
+    // AutDiff's operator== only compares the value; we should compare the full scalar.
+    m.def("equal_to", [](const ADScalar& a, const ADScalar& b) {
+      auto& a_d = a.derivatives();
+      auto& b_d = b.derivatives();
+      return a.value() == b.value() && a_d.size() == b_d.size() &&
+          (a_d.array() == b_d.array()).all();
+    });
+
     // test_special_matrix_objects
     // Returns a DiagonalMatrix with diagonal (1,2,3,...)
     m.def("incr_diag", [](int k) {

tests/test_eigen.py

@@ -151,41 +151,52 @@ def test_nonunit_stride_from_python():
     np.testing.assert_array_equal(counting_mat, [[0., 2, 2], [6, 16, 10], [6, 14, 8]])
 
 
-def conv_double_to_adscalar(arr, vice_versa=False):
-    flat_arr = arr.flatten()
-    new_arr = np.zeros(flat_arr.shape, dtype=object)
+def float_to_adscalar(arr, deriv):
+    arr = np.asarray(arr)
+    assert arr.dtype == float
+    new_arr = [m.AutoDiffXd(x, deriv) for x in arr.flat]
+    return np.array(new_arr).reshape(arr.shape)
 
-    for i in range(0, flat_arr.shape[0]):
-        if vice_versa:
-            new_arr[i] = flat_arr[i].value()
-        else:
-            new_arr[i] = m.AutoDiffXd(flat_arr[i], np.ones(1))
 
-    return new_arr.reshape(arr.shape)
+def adscalar_to_float(arr):
+    arr = np.asarray(arr)
+    assert arr.dtype == object
+    new_arr = [x.value() for x in arr.flat]
+    return np.array(new_arr).reshape(arr.shape)
+
+
+def check_array(a, b):
+    a, b = (np.asarray(x) for x in (a, b))
+    assert a.shape == b.shape and a.dtype == b.dtype
+    for index, (ai, bi) in enumerate(zip(a.flat, b.flat)):
+        assert m.equal_to(ai, bi), index
 
 
 def test_eigen_passing_adscalar():
-    adscalar_mat = conv_double_to_adscalar(ref)
+    assert m.equal_to(1., 1.)
+    assert not m.equal_to(1., 1.1)
+    assert m.equal_to(m.AutoDiffXd(0, [1.]), m.AutoDiffXd(0, [1.]))
+    assert not m.equal_to(m.AutoDiffXd(0, [1.]), m.AutoDiffXd(0, [1.1]))
+
+    adscalar_mat = float_to_adscalar(ref, deriv=[1.])
     adscalar_vec_col = adscalar_mat[:, 0]
     adscalar_vec_row = adscalar_mat[0, :]
 
-    # Checking if a Python vector is getting doubled, when passed into a dynamic
+    # Checking if a Python vector is getting doubled, when passed into a dynamic or fixed
     # row or col vector in Eigen.
-    adscalar_double_col = m.double_adscalar_col(adscalar_vec_col)
-    adscalar_double_row = m.double_adscalar_row(adscalar_vec_row)
-    np.testing.assert_array_equal(conv_double_to_adscalar(adscalar_double_col, vice_versa=True),
-                                  2 * ref[:, 0])
-    np.testing.assert_array_equal(conv_double_to_adscalar(adscalar_double_row, vice_versa=True),
-                                  2 * ref[0, :])
+    double_adscalar_mat = float_to_adscalar(2 * ref, deriv=[2.])
+    check_array(m.double_adscalar_col(adscalar_vec_col), double_adscalar_mat[:, 0])
+    check_array(m.double_adscalar_col5(adscalar_vec_col), double_adscalar_mat[:, 0])
+    check_array(m.double_adscalar_row(adscalar_vec_row), double_adscalar_mat[0, :])
+    check_array(m.double_adscalar_row6(adscalar_vec_row), double_adscalar_mat[0, :])
 
     # Adding 7 to the a dynamic matrix using reference.
-    incremented_adscalar_mat = conv_double_to_adscalar(m.incr_adscalar_matrix(adscalar_mat, 7.),
-                                                       vice_versa=True)
-    np.testing.assert_array_equal(incremented_adscalar_mat, ref + 7)
+    incr_adscalar_mat = float_to_adscalar(ref + 7, deriv=[1.])
+    check_array(m.incr_adscalar_matrix(adscalar_mat, 7.), incr_adscalar_mat)
     # The original adscalar_mat remains unchanged in spite of passing by reference, since
     # `Eigen::Ref<const CType>` permits copying, and copying is the only valid operation for
     # `dtype=object`.
-    np.testing.assert_array_equal(conv_double_to_adscalar(adscalar_mat, vice_versa=True), ref)
+    check_array(adscalar_to_float(adscalar_mat), ref)
 
     # Changes in Python are not reflected in C++ when internal_reference is returned.
     # These conversions should be disabled at runtime.