使用pybind11将NumPy数组强制转换为自定义C ++矩阵类

汤姆·德·格斯

我正在尝试使用包装我的C ++代码pybind11。在C ++中，我有一个Matrix3D充当3-D数组的类（即具有shape [n,m,p]）。它具有以下基本签名：

template <class T> class Matrix3D
{

  public:

    std::vector<T> data;
    std::vector<size_t> shape;
    std::vector<size_t> strides;

    Matrix3D<T>();
    Matrix3D<T>(std::vector<size_t>);
    Matrix3D<T>(const Matrix3D<T>&);

    T& operator() (int,int,int);

};

为了最大程度地减少包装器代码，我想将此类直接与NumPy数组进行转换（复制没有问题）。例如，我想直接包装以下签名的功能：

Matrix3D<double> func ( const Matrix3D<double>& );

使用包装器代码

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>

namespace py = pybind11;

PYBIND11_PLUGIN(example) {
  py::module m("example", "Module description");
  m.def("func", &func, "Function description" );
  return m.ptr();
}

目前，我之间还有另一个函数接受和返回py::array_t<double>。但是我想避免必须通过用一些模板替换每个包装函数来编写包装函数。

对于Eigen-library（对于数组和（2-D）矩阵）已完成此操作。但是对于我来说，代码太复杂了，无法从中获取自己的代码。另外，我真的只需要包装一个简单的类。

汤姆·德·格斯

在@kazemakase和@jagerman（后者通过pybind11论坛）的帮助下，我已经弄清楚了。类本身应该具有可以从某些输入复制的构造函数，这里使用迭代器：

#include <vector>
#include <assert.h>
#include <iterator>


template <class T> class Matrix3D
{
public:

  std::vector<T>      data;
  std::vector<size_t> shape;
  std::vector<size_t> strides;

  Matrix3D<T>() = default;

  template<class Iterator>
  Matrix3D<T>(const std::vector<size_t> &shape, Iterator first, Iterator last);
};


template <class T>
template<class Iterator>
Matrix3D<T>::Matrix3D(const std::vector<size_t> &shape_, Iterator first, Iterator last)
{
  shape = shape_;

  assert( shape.size() == 3 );

  strides.resize(3);

  strides[0] = shape[2]*shape[1];
  strides[1] = shape[2];
  strides[2] = 1;

  int size = shape[0] * shape[1] * shape[2];

  assert( last-first == size );

  data.resize(size);

  std::copy(first, last, data.begin());
}

直接包装以下签名的功能：

Matrix3D<double> func ( const Matrix3D<double>& );

需要以下包装器代码

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>

namespace py = pybind11;

namespace pybind11 { namespace detail {
  template <typename T> struct type_caster<Matrix3D<T>>
  {
    public:

      PYBIND11_TYPE_CASTER(Matrix3D<T>, _("Matrix3D<T>"));

      // Conversion part 1 (Python -> C++)
      bool load(py::handle src, bool convert) 
      {
        if ( !convert and !py::array_t<T>::check_(src) )
          return false;

        auto buf = py::array_t<T, py::array::c_style | py::array::forcecast>::ensure(src);
        if ( !buf )
          return false;

        auto dims = buf.ndim();
        if ( dims != 3  )
          return false;

        std::vector<size_t> shape(3);

        for ( int i = 0 ; i < 3 ; ++i )
          shape[i] = buf.shape()[i];

        value = Matrix3D<T>(shape, buf.data(), buf.data()+buf.size());

        return true;
      }

      //Conversion part 2 (C++ -> Python)
      static py::handle cast(const Matrix3D<T>& src, py::return_value_policy policy, py::handle parent) 
      {

        std::vector<size_t> shape  (3);
        std::vector<size_t> strides(3);

        for ( int i = 0 ; i < 3 ; ++i ) {
          shape  [i] = src.shape  [i];
          strides[i] = src.strides[i]*sizeof(T);
        }

        py::array a(std::move(shape), std::move(strides), src.data.data() );

        return a.release();

      }
  };
}} // namespace pybind11::detail

PYBIND11_PLUGIN(example) {
    py::module m("example", "Module description");
    m.def("func", &func, "Function description" );
    return m.ptr();
}

注意，现在函数重载也是可能的。例如，如果存在带有以下签名的重载函数：

Matrix3D<int   > func ( const Matrix3D<int   >& );
Matrix3D<double> func ( const Matrix3D<double>& );

需要以下包装器函数定义：

m.def("func", py::overload_cast<Matrix3D<int   >&>(&func), "Function description" );
m.def("func", py::overload_cast<Matrix3D<double>&>(&func), "Function description" );

本文收集自互联网，转载请注明来源。

如有侵权，请联系 [email protected] 删除。