The source code and dockerfile for the GSW2024 AI Lab.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
This repo is archived. You can view files and clone it, but cannot push or open issues/pull-requests.

366 lines
14 KiB

2 months ago
  1. .. _numpy:
  2. NumPy
  3. #####
  4. Buffer protocol
  5. ===============
  6. Python supports an extremely general and convenient approach for exchanging
  7. data between plugin libraries. Types can expose a buffer view [#f2]_, which
  8. provides fast direct access to the raw internal data representation. Suppose we
  9. want to bind the following simplistic Matrix class:
  10. .. code-block:: cpp
  11. class Matrix {
  12. public:
  13. Matrix(size_t rows, size_t cols) : m_rows(rows), m_cols(cols) {
  14. m_data = new float[rows*cols];
  15. }
  16. float *data() { return m_data; }
  17. size_t rows() const { return m_rows; }
  18. size_t cols() const { return m_cols; }
  19. private:
  20. size_t m_rows, m_cols;
  21. float *m_data;
  22. };
  23. The following binding code exposes the ``Matrix`` contents as a buffer object,
  24. making it possible to cast Matrices into NumPy arrays. It is even possible to
  25. completely avoid copy operations with Python expressions like
  26. ``np.array(matrix_instance, copy = False)``.
  27. .. code-block:: cpp
  28. py::class_<Matrix>(m, "Matrix", py::buffer_protocol())
  29. .def_buffer([](Matrix &m) -> py::buffer_info {
  30. return py::buffer_info(
  31. m.data(), /* Pointer to buffer */
  32. sizeof(float), /* Size of one scalar */
  33. py::format_descriptor<float>::format(), /* Python struct-style format descriptor */
  34. 2, /* Number of dimensions */
  35. { m.rows(), m.cols() }, /* Buffer dimensions */
  36. { sizeof(float) * m.rows(), /* Strides (in bytes) for each index */
  37. sizeof(float) }
  38. );
  39. });
  40. Supporting the buffer protocol in a new type involves specifying the special
  41. ``py::buffer_protocol()`` tag in the ``py::class_`` constructor and calling the
  42. ``def_buffer()`` method with a lambda function that creates a
  43. ``py::buffer_info`` description record on demand describing a given matrix
  44. instance. The contents of ``py::buffer_info`` mirror the Python buffer protocol
  45. specification.
  46. .. code-block:: cpp
  47. struct buffer_info {
  48. void *ptr;
  49. ssize_t itemsize;
  50. std::string format;
  51. ssize_t ndim;
  52. std::vector<ssize_t> shape;
  53. std::vector<ssize_t> strides;
  54. };
  55. To create a C++ function that can take a Python buffer object as an argument,
  56. simply use the type ``py::buffer`` as one of its arguments. Buffers can exist
  57. in a great variety of configurations, hence some safety checks are usually
  58. necessary in the function body. Below, you can see an basic example on how to
  59. define a custom constructor for the Eigen double precision matrix
  60. (``Eigen::MatrixXd``) type, which supports initialization from compatible
  61. buffer objects (e.g. a NumPy matrix).
  62. .. code-block:: cpp
  63. /* Bind MatrixXd (or some other Eigen type) to Python */
  64. typedef Eigen::MatrixXd Matrix;
  65. typedef Matrix::Scalar Scalar;
  66. constexpr bool rowMajor = Matrix::Flags & Eigen::RowMajorBit;
  67. py::class_<Matrix>(m, "Matrix", py::buffer_protocol())
  68. .def("__init__", [](Matrix &m, py::buffer b) {
  69. typedef Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic> Strides;
  70. /* Request a buffer descriptor from Python */
  71. py::buffer_info info = b.request();
  72. /* Some sanity checks ... */
  73. if (info.format != py::format_descriptor<Scalar>::format())
  74. throw std::runtime_error("Incompatible format: expected a double array!");
  75. if (info.ndim != 2)
  76. throw std::runtime_error("Incompatible buffer dimension!");
  77. auto strides = Strides(
  78. info.strides[rowMajor ? 0 : 1] / (py::ssize_t)sizeof(Scalar),
  79. info.strides[rowMajor ? 1 : 0] / (py::ssize_t)sizeof(Scalar));
  80. auto map = Eigen::Map<Matrix, 0, Strides>(
  81. static_cast<Scalar *>(info.ptr), info.shape[0], info.shape[1], strides);
  82. new (&m) Matrix(map);
  83. });
  84. For reference, the ``def_buffer()`` call for this Eigen data type should look
  85. as follows:
  86. .. code-block:: cpp
  87. .def_buffer([](Matrix &m) -> py::buffer_info {
  88. return py::buffer_info(
  89. m.data(), /* Pointer to buffer */
  90. sizeof(Scalar), /* Size of one scalar */
  91. py::format_descriptor<Scalar>::format(), /* Python struct-style format descriptor */
  92. 2, /* Number of dimensions */
  93. { m.rows(), m.cols() }, /* Buffer dimensions */
  94. { sizeof(Scalar) * (rowMajor ? m.cols() : 1),
  95. sizeof(Scalar) * (rowMajor ? 1 : m.rows()) }
  96. /* Strides (in bytes) for each index */
  97. );
  98. })
  99. For a much easier approach of binding Eigen types (although with some
  100. limitations), refer to the section on :doc:`/advanced/cast/eigen`.
  101. .. seealso::
  102. The file :file:`tests/test_buffers.cpp` contains a complete example
  103. that demonstrates using the buffer protocol with pybind11 in more detail.
  104. .. [#f2] http://docs.python.org/3/c-api/buffer.html
  105. Arrays
  106. ======
  107. By exchanging ``py::buffer`` with ``py::array`` in the above snippet, we can
  108. restrict the function so that it only accepts NumPy arrays (rather than any
  109. type of Python object satisfying the buffer protocol).
  110. In many situations, we want to define a function which only accepts a NumPy
  111. array of a certain data type. This is possible via the ``py::array_t<T>``
  112. template. For instance, the following function requires the argument to be a
  113. NumPy array containing double precision values.
  114. .. code-block:: cpp
  115. void f(py::array_t<double> array);
  116. When it is invoked with a different type (e.g. an integer or a list of
  117. integers), the binding code will attempt to cast the input into a NumPy array
  118. of the requested type. Note that this feature requires the
  119. :file:`pybind11/numpy.h` header to be included.
  120. Data in NumPy arrays is not guaranteed to packed in a dense manner;
  121. furthermore, entries can be separated by arbitrary column and row strides.
  122. Sometimes, it can be useful to require a function to only accept dense arrays
  123. using either the C (row-major) or Fortran (column-major) ordering. This can be
  124. accomplished via a second template argument with values ``py::array::c_style``
  125. or ``py::array::f_style``.
  126. .. code-block:: cpp
  127. void f(py::array_t<double, py::array::c_style | py::array::forcecast> array);
  128. The ``py::array::forcecast`` argument is the default value of the second
  129. template parameter, and it ensures that non-conforming arguments are converted
  130. into an array satisfying the specified requirements instead of trying the next
  131. function overload.
  132. Structured types
  133. ================
  134. In order for ``py::array_t`` to work with structured (record) types, we first
  135. need to register the memory layout of the type. This can be done via
  136. ``PYBIND11_NUMPY_DTYPE`` macro, called in the plugin definition code, which
  137. expects the type followed by field names:
  138. .. code-block:: cpp
  139. struct A {
  140. int x;
  141. double y;
  142. };
  143. struct B {
  144. int z;
  145. A a;
  146. };
  147. // ...
  148. PYBIND11_MODULE(test, m) {
  149. // ...
  150. PYBIND11_NUMPY_DTYPE(A, x, y);
  151. PYBIND11_NUMPY_DTYPE(B, z, a);
  152. /* now both A and B can be used as template arguments to py::array_t */
  153. }
  154. The structure should consist of fundamental arithmetic types, ``std::complex``,
  155. previously registered substructures, and arrays of any of the above. Both C++
  156. arrays and ``std::array`` are supported. While there is a static assertion to
  157. prevent many types of unsupported structures, it is still the user's
  158. responsibility to use only "plain" structures that can be safely manipulated as
  159. raw memory without violating invariants.
  160. Vectorizing functions
  161. =====================
  162. Suppose we want to bind a function with the following signature to Python so
  163. that it can process arbitrary NumPy array arguments (vectors, matrices, general
  164. N-D arrays) in addition to its normal arguments:
  165. .. code-block:: cpp
  166. double my_func(int x, float y, double z);
  167. After including the ``pybind11/numpy.h`` header, this is extremely simple:
  168. .. code-block:: cpp
  169. m.def("vectorized_func", py::vectorize(my_func));
  170. Invoking the function like below causes 4 calls to be made to ``my_func`` with
  171. each of the array elements. The significant advantage of this compared to
  172. solutions like ``numpy.vectorize()`` is that the loop over the elements runs
  173. entirely on the C++ side and can be crunched down into a tight, optimized loop
  174. by the compiler. The result is returned as a NumPy array of type
  175. ``numpy.dtype.float64``.
  176. .. code-block:: pycon
  177. >>> x = np.array([[1, 3],[5, 7]])
  178. >>> y = np.array([[2, 4],[6, 8]])
  179. >>> z = 3
  180. >>> result = vectorized_func(x, y, z)
  181. The scalar argument ``z`` is transparently replicated 4 times. The input
  182. arrays ``x`` and ``y`` are automatically converted into the right types (they
  183. are of type ``numpy.dtype.int64`` but need to be ``numpy.dtype.int32`` and
  184. ``numpy.dtype.float32``, respectively).
  185. .. note::
  186. Only arithmetic, complex, and POD types passed by value or by ``const &``
  187. reference are vectorized; all other arguments are passed through as-is.
  188. Functions taking rvalue reference arguments cannot be vectorized.
  189. In cases where the computation is too complicated to be reduced to
  190. ``vectorize``, it will be necessary to create and access the buffer contents
  191. manually. The following snippet contains a complete example that shows how this
  192. works (the code is somewhat contrived, since it could have been done more
  193. simply using ``vectorize``).
  194. .. code-block:: cpp
  195. #include <pybind11/pybind11.h>
  196. #include <pybind11/numpy.h>
  197. namespace py = pybind11;
  198. py::array_t<double> add_arrays(py::array_t<double> input1, py::array_t<double> input2) {
  199. auto buf1 = input1.request(), buf2 = input2.request();
  200. if (buf1.ndim != 1 || buf2.ndim != 1)
  201. throw std::runtime_error("Number of dimensions must be one");
  202. if (buf1.size != buf2.size)
  203. throw std::runtime_error("Input shapes must match");
  204. /* No pointer is passed, so NumPy will allocate the buffer */
  205. auto result = py::array_t<double>(buf1.size);
  206. auto buf3 = result.request();
  207. double *ptr1 = (double *) buf1.ptr,
  208. *ptr2 = (double *) buf2.ptr,
  209. *ptr3 = (double *) buf3.ptr;
  210. for (size_t idx = 0; idx < buf1.shape[0]; idx++)
  211. ptr3[idx] = ptr1[idx] + ptr2[idx];
  212. return result;
  213. }
  214. PYBIND11_MODULE(test, m) {
  215. m.def("add_arrays", &add_arrays, "Add two NumPy arrays");
  216. }
  217. .. seealso::
  218. The file :file:`tests/test_numpy_vectorize.cpp` contains a complete
  219. example that demonstrates using :func:`vectorize` in more detail.
  220. Direct access
  221. =============
  222. For performance reasons, particularly when dealing with very large arrays, it
  223. is often desirable to directly access array elements without internal checking
  224. of dimensions and bounds on every access when indices are known to be already
  225. valid. To avoid such checks, the ``array`` class and ``array_t<T>`` template
  226. class offer an unchecked proxy object that can be used for this unchecked
  227. access through the ``unchecked<N>`` and ``mutable_unchecked<N>`` methods,
  228. where ``N`` gives the required dimensionality of the array:
  229. .. code-block:: cpp
  230. m.def("sum_3d", [](py::array_t<double> x) {
  231. auto r = x.unchecked<3>(); // x must have ndim = 3; can be non-writeable
  232. double sum = 0;
  233. for (ssize_t i = 0; i < r.shape(0); i++)
  234. for (ssize_t j = 0; j < r.shape(1); j++)
  235. for (ssize_t k = 0; k < r.shape(2); k++)
  236. sum += r(i, j, k);
  237. return sum;
  238. });
  239. m.def("increment_3d", [](py::array_t<double> x) {
  240. auto r = x.mutable_unchecked<3>(); // Will throw if ndim != 3 or flags.writeable is false
  241. for (ssize_t i = 0; i < r.shape(0); i++)
  242. for (ssize_t j = 0; j < r.shape(1); j++)
  243. for (ssize_t k = 0; k < r.shape(2); k++)
  244. r(i, j, k) += 1.0;
  245. }, py::arg().noconvert());
  246. To obtain the proxy from an ``array`` object, you must specify both the data
  247. type and number of dimensions as template arguments, such as ``auto r =
  248. myarray.mutable_unchecked<float, 2>()``.
  249. If the number of dimensions is not known at compile time, you can omit the
  250. dimensions template parameter (i.e. calling ``arr_t.unchecked()`` or
  251. ``arr.unchecked<T>()``. This will give you a proxy object that works in the
  252. same way, but results in less optimizable code and thus a small efficiency
  253. loss in tight loops.
  254. Note that the returned proxy object directly references the array's data, and
  255. only reads its shape, strides, and writeable flag when constructed. You must
  256. take care to ensure that the referenced array is not destroyed or reshaped for
  257. the duration of the returned object, typically by limiting the scope of the
  258. returned instance.
  259. The returned proxy object supports some of the same methods as ``py::array`` so
  260. that it can be used as a drop-in replacement for some existing, index-checked
  261. uses of ``py::array``:
  262. - ``r.ndim()`` returns the number of dimensions
  263. - ``r.data(1, 2, ...)`` and ``r.mutable_data(1, 2, ...)``` returns a pointer to
  264. the ``const T`` or ``T`` data, respectively, at the given indices. The
  265. latter is only available to proxies obtained via ``a.mutable_unchecked()``.
  266. - ``itemsize()`` returns the size of an item in bytes, i.e. ``sizeof(T)``.
  267. - ``ndim()`` returns the number of dimensions.
  268. - ``shape(n)`` returns the size of dimension ``n``
  269. - ``size()`` returns the total number of elements (i.e. the product of the shapes).
  270. - ``nbytes()`` returns the number of bytes used by the referenced elements
  271. (i.e. ``itemsize()`` times ``size()``).
  272. .. seealso::
  273. The file :file:`tests/test_numpy_array.cpp` contains additional examples
  274. demonstrating the use of this feature.