使用pybind11将numpy阵列铸造到自定义C 矩阵级

Cast NumPy array to/from custom C++ Matrix-class using pybind11

本文关键字：自定义 pybind11 numpy 阵列 使用      更新时间：2023-10-16

我正在尝试使用pybind11包装C 代码。在C 中，我有一个Matrix3D类，该类充当3-D数组（即形状[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-array中施放此类（副本没问题）。例如，我想直接包装以下签名的函数：

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-纤维（对于阵列和（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" );

我对pybind11不熟悉，但在阅读了这个问题后变得很感兴趣。从文档中看，您似乎必须编写自己的类型施法者。这显然是一个相当高级的话题，但似乎可以通过一些努力可行。

从文档中剥离，这是Convertiong C 类型inty的转换器的外壳：

namespace pybind11 { namespace detail {
    template <> struct type_caster<inty> {
    public:
        PYBIND11_TYPE_CASTER(inty, _("inty"));    
        // Conversion part 1 (Python->C++)
        bool load(handle src, bool);
        //Conversion part 2 (C++ -> Python)
        static handle cast(inty src, return_value_policy, handle);
    };
}} // namespace pybind11::detail

看来，您要做的就是用Matrix3D<double>替换inty并实现load()和cast()。

让我们看看他们是如何为eigen做的（eigen.h，第236行）：

bool load(handle src, bool) {
    auto buf = array_t<Scalar>::ensure(src);
    if (!buf)
        return false;
    auto dims = buf.ndim();
    if (dims < 1 || dims > 2)
        return false;
    auto fits = props::conformable(buf);
    if (!fits)
        return false; // Non-comformable vector/matrix types
    value = Eigen::Map<const Type, 0, EigenDStride>(buf.data(), fits.rows, fits.cols, fits.stride);
    return true;
}

这看起来并不困难。首先，他们确保输入是array_t<Scalar>类型（在您的情况下可能是array_t<double>）。然后他们检查尺寸和一些符合性（您可能会跳过后者）。最后创建特征矩阵。由于复制不是问题，因此此时只需创建一个新的Martix3D<double>实例，然后将其填充Numpy数组中的数据。

对于不同的l值和const-Ness的cast()功能有不同的实现。我想只能执行一个实现，该实现可以在新的Numpy数组中创建数据副本，如果可以的话。请参阅函数eigen_array_cast()如何返回数组为handle返回类型。

我尚未对此进行任何测试，并且该过程可能比看起来更多。希望这将成为起点。

相关文章：