Merge pull request #211 from jcarpent/devel

jcarpent · web-flow · commit 5e6ac76c5eff · 2020-08-24T20:16:38.000+02:00
Add support for exporting Eigen::Ref to Python
diff --git a/include/eigenpy/details.hpp b/include/eigenpy/details.hpp
@@ -49,6 +49,20 @@ namespace boost { namespace python { namespace detail {
           std::size_t, value = sizeof(MatType));
   };
 
+  template<class MatType>
+  struct referent_size<Eigen::PlainObjectBase<MatType>&>
+  {
+      BOOST_STATIC_CONSTANT(
+          std::size_t, value = sizeof(MatType));
+  };
+
+  template<class MatType>
+  struct referent_size<Eigen::PlainObjectBase<MatType> >
+  {
+      BOOST_STATIC_CONSTANT(
+          std::size_t, value = sizeof(MatType));
+  };
+
 }}}
 
 namespace eigenpy
@@ -65,7 +79,13 @@ namespace eigenpy
   {
     if(check_registration<MatType>()) return;
     
+    // to-python
     EigenToPyConverter<MatType>::registration();
+#if EIGEN_VERSION_AT_LEAST(3,2,0)
+    EigenToPyConverter< Eigen::Ref<MatType> >::registration();
+#endif
+    
+    // from-python
     EigenFromPyConverter<MatType>::registration();
   }
 
diff --git a/include/eigenpy/eigen-from-python.hpp b/include/eigenpy/eigen-from-python.hpp
@@ -173,6 +173,14 @@ namespace boost { namespace python { namespace converter {
     EIGENPY_RVALUE_FROM_PYTHON_DATA_INIT(Derived const &)
   };
 
+  /// \brief Template specialization of rvalue_from_python_data
+  template<typename Derived>
+  struct rvalue_from_python_data<Eigen::PlainObjectBase<Derived> const &>
+  : rvalue_from_python_data_eigen<Derived const &>
+  {
+    EIGENPY_RVALUE_FROM_PYTHON_DATA_INIT(Derived const &)
+  };
+
   template<typename MatType, int Options, typename Stride>
   struct rvalue_from_python_data<Eigen::Ref<MatType,Options,Stride> &>
   : rvalue_from_python_storage<Eigen::Ref<MatType,Options,Stride> &>
@@ -425,6 +433,10 @@ namespace eigenpy
       typedef Eigen::EigenBase<MatType> EigenBase;
       EigenFromPy<EigenBase>::registration();
 
+      // Add conversion to Eigen::PlainObjectBase<MatType>
+      typedef Eigen::PlainObjectBase<MatType> PlainObjectBase;
+      EigenFromPy<PlainObjectBase>::registration();
+
 #if EIGEN_VERSION_AT_LEAST(3,2,0)
       // Add conversion to Eigen::Ref<MatType>
       typedef Eigen::Ref<MatType> RefType;
@@ -464,6 +476,20 @@ namespace eigenpy
        &EigenFromPy::construct,bp::type_id<Base>());
     }
   };
+    
+  template<typename MatType>
+  struct EigenFromPy< Eigen::PlainObjectBase<MatType> > : EigenFromPy<MatType>
+  {
+    typedef EigenFromPy<MatType> EigenFromPyDerived;
+    typedef Eigen::PlainObjectBase<MatType> Base;
+
+    static void registration()
+    {
+      bp::converter::registry::push_back
+      (reinterpret_cast<void *(*)(_object *)>(&EigenFromPy::convertible),
+       &EigenFromPy::construct,bp::type_id<Base>());
+    }
+  };
 
 #if EIGEN_VERSION_AT_LEAST(3,2,0)
 
diff --git a/include/eigenpy/eigen-to-python.hpp b/include/eigenpy/eigen-to-python.hpp
@@ -82,6 +82,36 @@ namespace eigenpy
     }
   };
 
+  template<typename MatType, int Options, typename Stride>
+  struct EigenToPy< Eigen::Ref<MatType,Options,Stride> >
+  {
+    static PyObject* convert(const Eigen::Ref<MatType,Options,Stride> & mat)
+    {
+      typedef Eigen::Ref<MatType,Options,Stride> EigenRef;
+      
+      assert( (mat.rows()<INT_MAX) && (mat.cols()<INT_MAX)
+             && "Matrix range larger than int ... should never happen." );
+      const npy_intp R = (npy_intp)mat.rows(), C = (npy_intp)mat.cols();
+      
+      PyArrayObject* pyArray;
+      // Allocate Python memory
+      if( ( ((!(C == 1) != !(R == 1)) && !MatType::IsVectorAtCompileTime) || MatType::IsVectorAtCompileTime)
+         && NumpyType::getType() == ARRAY_TYPE) // Handle array with a single dimension
+      {
+        npy_intp shape[1] = { C == 1 ? R : C };
+        pyArray = NumpyAllocator<EigenRef>::allocate(const_cast<EigenRef &>(mat),1,shape);
+      }
+      else
+      {
+        npy_intp shape[2] = { R,C };
+        pyArray = NumpyAllocator<EigenRef>::allocate(const_cast<EigenRef &>(mat),2,shape);
+      }
+      
+      // Create an instance (either np.array or np.matrix)
+      return NumpyType::make(pyArray).ptr();
+    }
+  };
+
   template<typename MatType>
   struct EigenToPyConverter
   {
diff --git a/include/eigenpy/numpy-allocator.hpp b/include/eigenpy/numpy-allocator.hpp
@@ -57,16 +57,41 @@ namespace eigenpy
       }
       else
       {
-        return NumpyAllocator<MatType>::allocate(mat.derived(),nd,shape);
+        return NumpyAllocator<MatType>::allocate(mat,nd,shape);
       }
     }
   };
 
 #if EIGEN_VERSION_AT_LEAST(3,2,0)
 
   template<typename MatType, int Options, typename Stride>
-  struct NumpyAllocator<Eigen::Ref<MatType,Options,Stride> > : NumpyAllocator<MatType &>
+  struct NumpyAllocator<Eigen::Ref<MatType,Options,Stride> >
   {
+    typedef Eigen::Ref<MatType,Options,Stride> RefType;
+    
+    static PyArrayObject * allocate(RefType & mat,
+                                    npy_intp nd, npy_intp * shape)
+    {
+      typedef typename RefType::Scalar Scalar;
+      enum { NPY_ARRAY_MEMORY_CONTIGUOUS = RefType::IsRowMajor ? NPY_ARRAY_CARRAY : NPY_ARRAY_FARRAY };
+      
+      if(NumpyType::sharedMemory())
+      {
+        const int Scalar_type_code = Register::getTypeCode<Scalar>();
+        PyArrayObject * pyArray = (PyArrayObject*) call_PyArray_New(getPyArrayType(),
+                                                                    static_cast<int>(nd),
+                                                                    shape,
+                                                                    Scalar_type_code,
+                                                                    mat.data(),
+                                                                    NPY_ARRAY_MEMORY_CONTIGUOUS | NPY_ARRAY_ALIGNED);
+        
+        return pyArray;
+      }
+      else
+      {
+        return NumpyAllocator<MatType>::allocate(mat,nd,shape);
+      }
+    }
   };
 
 #endif
@@ -88,23 +113,49 @@ namespace eigenpy
                                                                     static_cast<int>(nd),
                                                                     shape,
                                                                     Scalar_type_code,
-                                                                    const_cast<SimilarMatrixType &>(mat.derived()).data(),
+                                                                    const_cast<Scalar *>(mat.data()),
                                                                     NPY_ARRAY_MEMORY_CONTIGUOUS_RO | NPY_ARRAY_ALIGNED);
                                                                     
         return pyArray;
       }
       else
       {
-        return NumpyAllocator<MatType>::allocate(mat.derived(),nd,shape);
+        return NumpyAllocator<MatType>::allocate(mat,nd,shape);
       }
     }
   };
 
 #if EIGEN_VERSION_AT_LEAST(3,2,0)
 
   template<typename MatType, int Options, typename Stride>
-  struct NumpyAllocator<const Eigen::Ref<const MatType,Options,Stride> > : NumpyAllocator<const MatType &>
+  struct NumpyAllocator<const Eigen::Ref<const MatType,Options,Stride> >
   {
+    typedef const Eigen::Ref<const MatType,Options,Stride> RefType;
+    
+    template<typename SimilarMatrixType>
+    static PyArrayObject * allocate(RefType & mat,
+                                    npy_intp nd, npy_intp * shape)
+    {
+      typedef typename SimilarMatrixType::Scalar Scalar;
+      enum { NPY_ARRAY_MEMORY_CONTIGUOUS_RO = SimilarMatrixType::IsRowMajor ? NPY_ARRAY_CARRAY_RO : NPY_ARRAY_FARRAY_RO };
+      
+      if(NumpyType::sharedMemory())
+      {
+        const int Scalar_type_code = Register::getTypeCode<Scalar>();
+        PyArrayObject * pyArray = (PyArrayObject*) call_PyArray_New(getPyArrayType(),
+                                                                    static_cast<int>(nd),
+                                                                    shape,
+                                                                    Scalar_type_code,
+                                                                    const_cast<Scalar *>(mat.data()),
+                                                                    NPY_ARRAY_MEMORY_CONTIGUOUS_RO | NPY_ARRAY_ALIGNED);
+                                                                    
+        return pyArray;
+      }
+      else
+      {
+        return NumpyAllocator<MatType>::allocate(mat,nd,shape);
+      }
+    }
   };
 
 #endif
diff --git a/unittest/eigen_ref.cpp b/unittest/eigen_ref.cpp
@@ -37,6 +37,14 @@ void fill(Eigen::Ref<MatType> mat, const typename MatType::Scalar & value)
   mat.fill(value);
 }
 
+template<typename MatType>
+Eigen::Ref<MatType> asRef(const int rows, const int cols)
+{
+  static MatType mat(rows,cols);
+  std::cout << "mat:\n" << mat << std::endl;
+  return mat;
+}
+
 BOOST_PYTHON_MODULE(eigen_ref)
 {
   namespace bp = boost::python;
@@ -56,4 +64,6 @@ BOOST_PYTHON_MODULE(eigen_ref)
   bp::def("fillVec3", fill<Vector3d>);
   bp::def("fillVec", fill<VectorXd>);
   bp::def("fill", fill<MatrixXd>);
+  
+  bp::def("asRef", asRef<MatrixXd>);
 }
diff --git a/unittest/matrix.cpp b/unittest/matrix.cpp
@@ -81,6 +81,12 @@ MatrixDerived base(const Eigen::MatrixBase<MatrixDerived> & m)
   return m.derived();
 }
 
+template<typename MatrixDerived>
+MatrixDerived plain(const Eigen::PlainObjectBase<MatrixDerived> & m)
+{
+  return m.derived();
+}
+
 template<typename Scalar>
 Eigen::Matrix<Scalar,6,6> matrix6(const Scalar & value)
 {
@@ -123,6 +129,9 @@ BOOST_PYTHON_MODULE(matrix)
   
   bp::def("base", base<VectorXd>);
   bp::def("base", base<MatrixXd>);
+  
+  bp::def("plain", plain<VectorXd>);
+  bp::def("plain", plain<MatrixXd>);
 
   bp::def("matrix6", matrix6<double>);
 }
diff --git a/unittest/python/test_eigen_ref.py b/unittest/python/test_eigen_ref.py
@@ -8,6 +8,16 @@ def test(mat):
   printMatrix(mat)
   assert np.array_equal(mat,np.full(mat.shape,1.))
 
+  A_ref = asRef(mat.shape[0],mat.shape[1])
+  A_ref.fill(1.)
+  A_ref2 = asRef(mat.shape[0],mat.shape[1])
+
+  assert np.array_equal(A_ref,A_ref2)
+
+  A_ref2.fill(0)
+  assert np.array_equal(A_ref,A_ref2)
+  
+
 rows = 10
 cols = 30
 
diff --git a/unittest/python/test_matrix.py b/unittest/python/test_matrix.py
@@ -36,6 +36,10 @@
 Mref_from_base = eigenpy.base(Mref)
 assert( np.array_equal(Mref,Mref_from_base) );
 
+# Test plain function
+Mref_from_plain = eigenpy.plain(Mref)
+assert( np.array_equal(Mref,Mref_from_plain) );
+
 if verbose: print("===> Matrix 8x8")
 M = Mref
 assert( np.array_equal(M,eigenpy.reflex(M,verbose)) );

Original file line number	Diff line number	Diff line change
`@@ -81,6 +81,12 @@ MatrixDerived base(const Eigen::MatrixBase<MatrixDerived> & m)`
`81`	`81`	`return m.derived();`
`82`	`82`	`}`
`83`	`83`
	`84`	`+template<typename MatrixDerived>`
	`85`	`+MatrixDerived plain(const Eigen::PlainObjectBase<MatrixDerived> & m)`
	`86`	`+{`
	`87`	`+ return m.derived();`
	`88`	`+}`
	`89`	`+`
`84`	`90`	`template<typename Scalar>`
`85`	`91`	`Eigen::Matrix<Scalar,6,6> matrix6(const Scalar & value)`
`86`	`92`	`{`
`@@ -123,6 +129,9 @@ BOOST_PYTHON_MODULE(matrix)`
`123`	`129`
`124`	`130`	`bp::def("base", base<VectorXd>);`
`125`	`131`	`bp::def("base", base<MatrixXd>);`
	`132`	`+`
	`133`	`+ bp::def("plain", plain<VectorXd>);`
	`134`	`+ bp::def("plain", plain<MatrixXd>);`
`126`	`135`
`127`	`136`	`bp::def("matrix6", matrix6<double>);`
`128`	`137`	`}`