Passing a NumPy Array into C++
Imagine we want to create a NumPy array in python and pass it into C++ to operate on.
In C++, we refer to the incoming array using the pybind11::array_t<T>
class.
We then create a view (the pybind11 docs call it a proxy) that allows us to index
the array. To do this, we need to specify the dimensions of the array via the template
parameter.
Note that this array is immutable because of the method we used to access its contents.
#include <iostream>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
double sum(pybind11::array_t<double>& a) {
auto a_arr = a.unchecked<2>();
double suma{0.0};
for (pybind11::ssize_t i = 0; i < a_arr.shape(0); ++i) {
for (pybind11::ssize_t j = 0; j < a_arr.shape(1); ++j) {
suma += a_arr(i, j);
}
}
return suma;
}
PYBIND11_MODULE(sum, m) {
m.doc() = "sum the elements of a 2-dimensional array";
m.def("sum", &sum);
}
To compile it, we can essentially the same GNUmakefile
:
SOURCES := $(wildcard *.cpp)
OBJECTS := $(SOURCES:.cpp=.o)
BASE := sum
HEADERS := $(wildcard *.H)
PYBIND_INCLUDES := $(shell python -m pybind11 --includes)
PYBIND_SUFFIX := $(shell python3-config --extension-suffix)
PYTHON_LIBRARY := ${BASE}${PYBIND_SUFFIX}
ALL: ${PYTHON_LIBRARY}
CXXFLAGS := -O3 -Wall -Wextra -shared -std=c++17 -fPIC ${PYBIND_INCLUDES}
%.o : %.cpp
g++ ${CXXFLAGS} -c $<
${PYTHON_LIBRARY}: ${OBJECTS} ${HEADERS}
g++ ${CXXFLAGS} -o $@ $<
print-%: ; @echo $* is $($*)
We can use this in python as:
import numpy as np
import sum
a = np.ones((5, 4), dtype=np.float64)
print(sum.sum(a))
b = np.ones(5, dtype=np.float64)
print(sum.sum(b))
Notice that the creation of the view / proxy into the array fails if we pass in an array with the incorrect number of dimensions.
Mutable array
Here’s a similar C++ function, but now we use
mutable_unchecked<DIM>()
to get our view into the passed array.
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
void update(pybind11::array_t<double>& a) {
auto a_arr = a.mutable_unchecked<2>();
for (pybind11::ssize_t i = 0; i < a_arr.shape(0); ++i) {
for (pybind11::ssize_t j = 0; j < a_arr.shape(1); ++j) {
a_arr(i, j) *= 100;
}
}
}
PYBIND11_MODULE(mutate, m) {
m.doc() = "change the of a 2-dimensional array";
m.def("update", &update);
}
And again, the same GNUmakefile
(just changing the module name): GNUmakefile
.
Here’s how we can use it from python:
import numpy as np
from mutate import update
a = np.arange(12, dtype=np.float64).reshape(4, 3)
update(a)
print(a)
Note
It doesn’t matter if we declare our function as:
void update(pybind11::array_t<double>& a)
or
void update(pybind11::array_t<double> a)
The making the array_t
a reference doesn’t matter, since array_t
is
just a wrapper that has a pointer to a buffer where the data is stored.