Merge pull request #279 from ManifoldFR/modify-block-test

Test: modify a matrix block through Python subclass

include/eigenpy/numpy-allocator.hpp

@@ -70,9 +70,15 @@ struct NumpyAllocator<Eigen::Ref<MatType, Options, Stride> > {
 
     if (NumpyType::sharedMemory()) {
       const int Scalar_type_code = Register::getTypeCode<Scalar>();
+      Eigen::DenseIndex inner_stride = mat.innerStride(),
+                        outer_stride = mat.outerStride();
+
+      const int elsize = call_PyArray_DescrFromType(Scalar_type_code)->elsize;
+      npy_intp strides[2] = {elsize * inner_stride, elsize * outer_stride};
+
       PyArrayObject *pyArray = (PyArrayObject *)call_PyArray_New(
           getPyArrayType(), static_cast<int>(nd), shape, Scalar_type_code,
-          mat.data(), NPY_ARRAY_MEMORY_CONTIGUOUS | NPY_ARRAY_ALIGNED);
+          strides, mat.data(), NPY_ARRAY_MEMORY_CONTIGUOUS | NPY_ARRAY_ALIGNED);
 
       return pyArray;
     } else {
@@ -125,9 +131,15 @@ struct NumpyAllocator<const Eigen::Ref<const MatType, Options, Stride> > {
 
     if (NumpyType::sharedMemory()) {
       const int Scalar_type_code = Register::getTypeCode<Scalar>();
+      Eigen::DenseIndex inner_stride = mat.innerStride(),
+                        outer_stride = mat.outerStride();
+
+      const int elsize = call_PyArray_DescrFromType(Scalar_type_code)->elsize;
+      npy_intp strides[2] = {elsize * inner_stride, elsize * outer_stride};
+
       PyArrayObject *pyArray = (PyArrayObject *)call_PyArray_New(
           getPyArrayType(), static_cast<int>(nd), shape, Scalar_type_code,
-          const_cast<Scalar *>(mat.data()),
+          strides, const_cast<Scalar *>(mat.data()),
           NPY_ARRAY_MEMORY_CONTIGUOUS_RO | NPY_ARRAY_ALIGNED);
 
       return pyArray;

include/eigenpy/numpy.hpp

 /*
- * Copyright 2020-2021 INRIA
+ * Copyright 2020-2022 INRIA
  */
 
 #ifndef __eigenpy_numpy_hpp__
@@ -109,6 +109,11 @@ EIGENPY_DLLAPI PyObject* call_PyArray_New(PyTypeObject* py_type_ptr, int nd,
                                           npy_intp* shape, int np_type,
                                           void* data_ptr, int options);
 
+EIGENPY_DLLAPI PyObject* call_PyArray_New(PyTypeObject* py_type_ptr, int nd,
+                                          npy_intp* shape, int np_type,
+                                          npy_intp* strides, void* data_ptr,
+                                          int options);
+
 EIGENPY_DLLAPI int call_PyArray_ObjectType(PyObject*, int);
 
 EIGENPY_DLLAPI PyTypeObject* getPyArrayType();
@@ -143,6 +148,14 @@ inline PyObject* call_PyArray_New(PyTypeObject* py_type_ptr, int nd,
                      options, NULL);
 }
 
+inline PyObject* call_PyArray_New(PyTypeObject* py_type_ptr, int nd,
+                                  npy_intp* shape, int np_type,
+                                  npy_intp* strides, void* data_ptr,
+                                  int options) {
+  return PyArray_New(py_type_ptr, nd, shape, np_type, strides, data_ptr, 0,
+                     options, NULL);
+}
+
 inline int call_PyArray_ObjectType(PyObject* obj, int val) {
   return PyArray_ObjectType(obj, val);
 }

src/numpy.cpp

 /*
- * Copyright 2020-2021 INRIA
+ * Copyright 2020-2022 INRIA
  */
 
 #include "eigenpy/numpy.hpp"
@@ -31,6 +31,13 @@ PyObject* call_PyArray_New(PyTypeObject* py_type_ptr, int nd, npy_intp* shape,
                      options, NULL);
 }
 
+PyObject* call_PyArray_New(PyTypeObject* py_type_ptr, int nd, npy_intp* shape,
+                           int np_type, npy_intp* strides, void* data_ptr,
+                           int options) {
+  return PyArray_New(py_type_ptr, nd, shape, np_type, strides, data_ptr, 0,
+                     options, NULL);
+}
+
 int call_PyArray_ObjectType(PyObject* obj, int val) {
   return PyArray_ObjectType(obj, val);
 }

unittest/eigen_ref.cpp

@@ -30,6 +30,28 @@ void setOnes(Eigen::Ref<MatType> mat) {
   mat.setOnes();
 }
 
+template <typename MatType>
+Eigen::Ref<MatType> getBlock(Eigen::Ref<MatType> mat, Eigen::DenseIndex i,
+                             Eigen::DenseIndex j, Eigen::DenseIndex n,
+                             Eigen::DenseIndex m) {
+  return mat.block(i, j, n, m);
+}
+
+template <typename MatType>
+Eigen::Ref<MatType> editBlock(Eigen::Ref<MatType> mat, Eigen::DenseIndex i,
+                              Eigen::DenseIndex j, Eigen::DenseIndex n,
+                              Eigen::DenseIndex m) {
+  typename Eigen::Ref<MatType>::BlockXpr B = mat.block(i, j, n, m);
+  int k = 0;
+  for (int i = 0; i < B.rows(); ++i) {
+    for (int j = 0; j < B.cols(); ++j) {
+      B(i, j) = k++;
+    }
+  }
+  std::cout << "B:\n" << B << std::endl;
+  return mat;
+}
+
 template <typename MatType>
 void fill(Eigen::Ref<MatType> mat, const typename MatType::Scalar& value) {
   mat.fill(value);
@@ -52,6 +74,18 @@ const Eigen::Ref<const MatType> asConstRef(Eigen::Ref<MatType> mat) {
   return Eigen::Ref<const MatType>(mat);
 }
 
+struct modify_block {
+  MatrixXd J;
+  modify_block() : J(10, 10) { J.setZero(); }
+  void modify(int n, int m) { call(J.topLeftCorner(n, m)); }
+  virtual void call(Eigen::Ref<MatrixXd> mat) = 0;
+};
+
+struct modify_wrap : modify_block, bp::wrapper<modify_block> {
+  modify_wrap() : modify_block() {}
+  void call(Eigen::Ref<MatrixXd> mat) { this->get_override("call")(mat); }
+};
+
 BOOST_PYTHON_MODULE(eigen_ref) {
   namespace bp = boost::python;
   eigenpy::enableEigenPy();
@@ -77,4 +111,12 @@ BOOST_PYTHON_MODULE(eigen_ref) {
           (Eigen::Ref<MatrixXd>(*)(Eigen::Ref<MatrixXd>))asRef<MatrixXd>);
   bp::def("asConstRef", (const Eigen::Ref<const MatrixXd> (*)(
                             Eigen::Ref<MatrixXd>))asConstRef<MatrixXd>);
+
+  bp::def("getBlock", &getBlock<MatrixXd>);
+  bp::def("editBlock", &editBlock<MatrixXd>);
+
+  bp::class_<modify_wrap, boost::noncopyable>("modify_block", bp::init<>())
+      .def_readonly("J", &modify_block::J)
+      .def("modify", &modify_block::modify)
+      .def("call", bp::pure_virtual(&modify_wrap::call));
 }

unittest/python/test_eigen_ref.py

@@ -24,6 +24,44 @@ def test(mat):
     const_ref = asConstRef(mat)
     assert np.all(const_ref == mat)
 
+    mat.fill(0.0)
+    fill(mat[:3, :2], 1.0)
+
+    assert np.all(mat[:3, :2] == np.ones((3, 2)))
+
+    mat.fill(0.0)
+    fill(mat[:2, :3], 1.0)
+
+    assert np.all(mat[:2, :3] == np.ones((2, 3)))
+
+    mat.fill(0.0)
+    mat_as_C_order = np.array(mat, order="F")
+    getBlock(mat_as_C_order, 0, 0, 3, 2)[:, :] = 1.0
+
+    assert np.all(mat_as_C_order[:3, :2] == np.ones((3, 2)))
+
+    mat_as_C_order[:3, :2] = 0.0
+    mat_copy = mat_as_C_order.copy()
+    editBlock(mat_as_C_order, 0, 0, 3, 2)
+    mat_copy[:3, :2] = np.arange(6).reshape(3, 2)
+
+    assert np.all(mat_as_C_order == mat_copy)
+
+    class ModifyBlockImpl(modify_block):
+        def __init__(self):
+            super().__init__()
+
+        def call(self, mat):
+            n, m = mat.shape
+            mat[:, :] = np.arange(n * m).reshape(n, m)
+
+    modify = ModifyBlockImpl()
+    modify.modify(2, 3)
+    Jref = np.zeros((10, 10))
+    Jref[:2, :3] = np.arange(6).reshape(2, 3)
+
+    assert np.array_equal(Jref, modify.J)
+
 
 rows = 10
 cols = 30