I am trying to wrap my C++ code using pybind11
. In C++ I have a class Matrix3D
which acts as a 3-D array (i.e. with shape [n,m,p]
). It has the following basic signature:
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);
};
To minimize the wrapper code I would like to cast this class directly to and from a NumPy-array (copies are no problem). For example, I would like to directly wrap a function of the following signature:
Matrix3D<double> func ( const Matrix3D<double>& );
using the wrapper code
#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();
}
Currently I have another function in-between that accepts and returns py::array_t<double>
. But I would like to avoid having to write a wrapper function for each function by replacing it by some template.
This has been done for the Eigen
-library (for arrays and (2-D) matrices). But the code is too involved for me to derive my own code from. Plus, I really only need to wrap only one, simple, class.
With the help of @kazemakase and @jagerman (the latter via the pybind11 forum) I have figured it out. The class itself should have a constructor that can copy from some input, here using an iterator:
#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());
}
To directly wrap a function of the following signature:
Matrix3D<double> func ( const Matrix3D<double>& );
the following wrapper code is needed
#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();
}
Note that function overloading is now also possible. For example if an overloaded function would exist with the following signature:
Matrix3D<int > func ( const Matrix3D<int >& );
Matrix3D<double> func ( const Matrix3D<double>& );
The following wrapper function definition would be needed:
m.def("func", py::overload_cast<Matrix3D<int >&>(&func), "Function description" );
m.def("func", py::overload_cast<Matrix3D<double>&>(&func), "Function description" );
I am not familiar with pybind11 but became interested after reading this question. From the documentanion it looks like you will have to write your own type caster. This apparently is a rather advanced topic but seems doable with some effort.
Stripped from the documentation, this is the shell of such a converter for convertiong the C++ type 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
It seems all you have to do is to replace inty
with Matrix3D<double>
and implement load()
and cast()
.
Let's see how they did it for Eigen (eigen.h, line 236 onward):
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;
}
This does not look too difficult. First they make sure the input is of type array_t<Scalar>
(probably array_t<double>
in your case). Then they check the dimensions and some conformatibility (you can probably skip the latter). And finally create the Eigen matrix. Since copying is not a problem, at this point simply create a new Martix3D<double>
instance and fill it with the data from the numpy array.
There are different implementations of the cast()
function for different cases of l-value and const-ness. I guess it is sufficient to do only one implementation that creates a copy of the data in a new numpy array, if that is fine. See function eigen_array_cast()
how to return an array as handle
return type.
I have not tested any of this and there may be more to the process than it seems. Hopefully this will serve as a starting point.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With