pybind11和numpy进行交互

使用一个遵循buffer protocol的对象就可以和numpy交互了.
 
这个buffer_protocol要有哪些东西呢? 要有如下接口:

  
struct buffer_info {
    void *ptr;
    ssize_t itemsize;
    std::string format;
    ssize_t ndim;
    std::vector<ssize_t> shape;
    std::vector<ssize_t> strides;
};

 

 
其实就是一个指向数组的指针+各个维度的信息就可以了. 然后我们就可以用指针+偏移来访问数字中的任意位置上的数字了.
 
下面是一个可以跑的例子:
 
 1 #include <pybind11/pybind11.h>
 2 #include <pybind11/numpy.h>
 3 
 4 namespace py = pybind11;
 5 
 6 py::array_t<double> add_arrays(py::array_t<double> input1, py::array_t<double> input2) {
 7     py::buffer_info buf1 = input1.request(), buf2 = input2.request();
 8 
 9     if (buf1.ndim != 1 || buf2.ndim != 1)
10         throw std::runtime_error("Number of dimensions must be one");
11 
12     if (buf1.size != buf2.size)
13         throw std::runtime_error("Input shapes must match");
14 
15     /* No pointer is passed, so NumPy will allocate the buffer */
16     auto result = py::array_t<double>(buf1.size);
17 
18     py::buffer_info buf3 = result.request();
19 
20     double *ptr1 = (double *) buf1.ptr,
21            *ptr2 = (double *) buf2.ptr,
22            *ptr3 = (double *) buf3.ptr;
23 
24     for (size_t idx = 0; idx < buf1.shape[0]; idx++)
25         ptr3[idx] = ptr1[idx] + ptr2[idx];
26 
27     return result;
28 }
29 
30 PYBIND11_MODULE(test, m) {
31     m.def("add_arrays", &add_arrays, "Add two NumPy arrays");
32 }

 

 
array_t里的buf就是一个兼容的接口.
buf中可以得到指针和对应数字的维度信息.
 
 
为了方便我们甚至可以使用Eigen当作我们兼容numpy的接口:
 
 
 1 #include <pybind11/pybind11.h>
 2 #include <pybind11/eigen.h>
 3 
 4 #include <Eigen/LU>
 5 
 6 // N.B. this would equally work with Eigen-types that are not predefined. For example replacing
 7 // all occurrences of "Eigen::MatrixXd" with "MatD", with the following definition:
 8 //
 9 //  typedef Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> MatD;
10 
11 Eigen::MatrixXd inv(const Eigen::MatrixXd &xs)
12 {
13   return xs.inverse();
14 }
15 
16 double det(const Eigen::MatrixXd &xs)
17 {
18   return xs.determinant();
19 }
20 
21 namespace py = pybind11;
22 
23 PYBIND11_MODULE(example,m)
24 {
25   m.doc() = "pybind11 example plugin";
26 
27   m.def("inv", &inv);
28 
29   m.def("det", &det);
30 }

 

 
 
更多参考:
 
 
posted @ 2019-07-04 10:08  酱_油  阅读(4434)  评论(0)  编辑  收藏  举报