记录numpy和c++的混合编程
准备记录numpy和c++的混合编程
#include <boost/python.hpp> #include <numpy/ndarrayobject.h> namespace bp = boost::python; void reference_contiguous_array(PyObject* in, PyArrayObject* in_con, double* ptr, int& count) { in_con = PyArray_GETCONTIGUOUS((PyArrayObject*)in); ptr = (double*)PyArray_DATA(in_con); int num_dim = PyArray_NDIM(in_con); npy_intp* pdim = PyArray_DIMS(in_con); count = 1; for (int i = 0; i < num_dim; i++) { count *= pdim[i]; } } void dereference(PyObject* o) { Py_DECREF(o); } PyObject* entry_square_matrix(PyObject* input_matrix) { // get the input array double* ptr; int count; PyArrayObject* input_contigous_array; reference_contiguous_array(input_matrix, input_contigous_array, ptr, count); // create the output array npy_intp dst_dim[1]; dst_dim[0] = count; PyObject* out_matrix = PyArray_SimpleNew(1, dst_dim, NPY_FLOAT64); double* ptr_out; PyArrayObject* output_contigous_array; reference_contiguous_array(out_matrix, output_contigous_array, ptr_out, count); for (int i = 0; i < count; i++) { ptr_out[i] = ptr[i] * ptr[i]; } dereference((PyObject*)input_contigous_array); dereference((PyObject*)output_contigous_array); return out_matrix; } BOOST_PYTHON_MODULE(_func) { import_array(); bp::def(square_matrix, entry_square_matrix); }