blob: 32b50e3d0e387151f06f8f7122b0245840681aa1 [file] [log] [blame]
Ivan Smirnovbb4015d2016-06-19 15:50:31 +01001/*
2 example/example20.cpp -- Usage of structured numpy dtypes
3
4 Copyright (c) 2016 Ivan Smirnov
5
6 All rights reserved. Use of this source code is governed by a
7 BSD-style license that can be found in the LICENSE file.
8*/
9
10#include "example.h"
11
12#include <pybind11/numpy.h>
13#include <cstdint>
14#include <iostream>
15
16namespace py = pybind11;
17
Ivan Smirnov4f164212016-06-22 01:07:20 +010018struct SimpleStruct {
Ivan Smirnovbb4015d2016-06-19 15:50:31 +010019 bool x;
20 uint32_t y;
21 float z;
22};
23
Ivan Smirnov4f164212016-06-22 01:07:20 +010024std::ostream& operator<<(std::ostream& os, const SimpleStruct& v) {
25 return os << "s:" << v.x << "," << v.y << "," << v.z;
Ivan Smirnov669e1422016-06-21 21:05:29 +010026}
27
Ivan Smirnovbb4015d2016-06-19 15:50:31 +010028struct PackedStruct {
29 bool x;
30 uint32_t y;
31 float z;
32} __attribute__((packed));
33
Ivan Smirnov669e1422016-06-21 21:05:29 +010034std::ostream& operator<<(std::ostream& os, const PackedStruct& v) {
Ivan Smirnov4f164212016-06-22 01:07:20 +010035 return os << "p:" << v.x << "," << v.y << "," << v.z;
Ivan Smirnov669e1422016-06-21 21:05:29 +010036}
37
Ivan Smirnovbb4015d2016-06-19 15:50:31 +010038struct NestedStruct {
Ivan Smirnov4f164212016-06-22 01:07:20 +010039 SimpleStruct a;
Ivan Smirnovbb4015d2016-06-19 15:50:31 +010040 PackedStruct b;
Ivan Smirnov7f913ae2016-06-19 16:41:15 +010041} __attribute__((packed));
42
Ivan Smirnov669e1422016-06-21 21:05:29 +010043std::ostream& operator<<(std::ostream& os, const NestedStruct& v) {
Ivan Smirnov4f164212016-06-22 01:07:20 +010044 return os << "n:a=" << v.a << ";b=" << v.b;
Ivan Smirnov669e1422016-06-21 21:05:29 +010045}
46
Ivan Smirnov8fa09cb2016-07-06 00:28:12 +010047struct PartialStruct {
48 bool x;
49 uint32_t y;
50 float z;
51 long dummy2;
52};
53
54struct PartialNestedStruct {
55 long dummy1;
56 PartialStruct a;
57 long dummy2;
58};
59
Ivan Smirnovd0bafd92016-06-26 16:35:28 +010060struct UnboundStruct { };
61
Ivan Smirnov7f913ae2016-06-19 16:41:15 +010062template <typename T>
63py::array mkarray_via_buffer(size_t n) {
64 return py::array(py::buffer_info(nullptr, sizeof(T),
Ivan Smirnov5e71e172016-06-26 12:42:34 +010065 py::format_descriptor<T>::format(),
Ivan Smirnov7f913ae2016-06-19 16:41:15 +010066 1, { n }, { sizeof(T) }));
67}
Ivan Smirnovbb4015d2016-06-19 15:50:31 +010068
69template <typename S>
Ivan Smirnov8fa09cb2016-07-06 00:28:12 +010070py::array_t<S, 0> create_recarray(size_t n) {
Ivan Smirnov7f913ae2016-06-19 16:41:15 +010071 auto arr = mkarray_via_buffer<S>(n);
72 auto ptr = static_cast<S*>(arr.request().ptr);
Ivan Smirnovbb4015d2016-06-19 15:50:31 +010073 for (size_t i = 0; i < n; i++) {
Ivan Smirnov7f913ae2016-06-19 16:41:15 +010074 ptr[i].x = i % 2; ptr[i].y = (uint32_t) i; ptr[i].z = (float) i * 1.5f;
Ivan Smirnovbb4015d2016-06-19 15:50:31 +010075 }
76 return arr;
77}
78
Ivan Smirnovd0bafd92016-06-26 16:35:28 +010079std::string get_format_unbound() {
80 return py::format_descriptor<UnboundStruct>::format();
81}
82
Ivan Smirnov8fa09cb2016-07-06 00:28:12 +010083py::array_t<NestedStruct, 0> create_nested(size_t n) {
Ivan Smirnov7f913ae2016-06-19 16:41:15 +010084 auto arr = mkarray_via_buffer<NestedStruct>(n);
85 auto ptr = static_cast<NestedStruct*>(arr.request().ptr);
86 for (size_t i = 0; i < n; i++) {
87 ptr[i].a.x = i % 2; ptr[i].a.y = (uint32_t) i; ptr[i].a.z = (float) i * 1.5f;
88 ptr[i].b.x = (i + 1) % 2; ptr[i].b.y = (uint32_t) (i + 1); ptr[i].b.z = (float) (i + 1) * 1.5f;
89 }
90 return arr;
Ivan Smirnovbdc99022016-06-19 16:54:07 +010091}
Ivan Smirnov7f913ae2016-06-19 16:41:15 +010092
Ivan Smirnov8fa09cb2016-07-06 00:28:12 +010093py::array_t<PartialNestedStruct, 0> create_partial_nested(size_t n) {
94 auto arr = mkarray_via_buffer<PartialNestedStruct>(n);
95 auto ptr = static_cast<PartialNestedStruct*>(arr.request().ptr);
96 for (size_t i = 0; i < n; i++) {
97 ptr[i].a.x = i % 2; ptr[i].a.y = (uint32_t) i; ptr[i].a.z = (float) i * 1.5f;
98 }
99 return arr;
100}
101
Ivan Smirnov669e1422016-06-21 21:05:29 +0100102template <typename S>
Ivan Smirnov8fa09cb2016-07-06 00:28:12 +0100103void print_recarray(py::array_t<S, 0> arr) {
Ivan Smirnov669e1422016-06-21 21:05:29 +0100104 auto buf = arr.request();
105 auto ptr = static_cast<S*>(buf.ptr);
106 for (size_t i = 0; i < buf.size; i++)
107 std::cout << ptr[i] << std::endl;
108}
Ivan Smirnovbdc99022016-06-19 16:54:07 +0100109
110void print_format_descriptors() {
Ivan Smirnov5e71e172016-06-26 12:42:34 +0100111 std::cout << py::format_descriptor<SimpleStruct>::format() << std::endl;
112 std::cout << py::format_descriptor<PackedStruct>::format() << std::endl;
113 std::cout << py::format_descriptor<NestedStruct>::format() << std::endl;
Ivan Smirnov8fa09cb2016-07-06 00:28:12 +0100114 std::cout << py::format_descriptor<PartialStruct>::format() << std::endl;
115 std::cout << py::format_descriptor<PartialNestedStruct>::format() << std::endl;
Ivan Smirnov7f913ae2016-06-19 16:41:15 +0100116}
117
Ivan Smirnov4f164212016-06-22 01:07:20 +0100118void print_dtypes() {
119 auto to_str = [](py::object obj) {
120 return (std::string) (py::str) ((py::object) obj.attr("__str__"))();
121 };
122 std::cout << to_str(py::dtype_of<SimpleStruct>()) << std::endl;
123 std::cout << to_str(py::dtype_of<PackedStruct>()) << std::endl;
124 std::cout << to_str(py::dtype_of<NestedStruct>()) << std::endl;
Ivan Smirnov8fa09cb2016-07-06 00:28:12 +0100125 std::cout << to_str(py::dtype_of<PartialStruct>()) << std::endl;
126 std::cout << to_str(py::dtype_of<PartialNestedStruct>()) << std::endl;
Ivan Smirnov4f164212016-06-22 01:07:20 +0100127}
128
Ivan Smirnovbb4015d2016-06-19 15:50:31 +0100129void init_ex20(py::module &m) {
Ivan Smirnov5412a052016-07-02 16:18:42 +0100130 PYBIND11_NUMPY_DTYPE(SimpleStruct, x, y, z);
131 PYBIND11_NUMPY_DTYPE(PackedStruct, x, y, z);
132 PYBIND11_NUMPY_DTYPE(NestedStruct, a, b);
Ivan Smirnov8fa09cb2016-07-06 00:28:12 +0100133 PYBIND11_NUMPY_DTYPE(PartialStruct, x, y, z);
134 PYBIND11_NUMPY_DTYPE(PartialNestedStruct, a);
Ivan Smirnovbb4015d2016-06-19 15:50:31 +0100135
Ivan Smirnov4f164212016-06-22 01:07:20 +0100136 m.def("create_rec_simple", &create_recarray<SimpleStruct>);
Ivan Smirnovbb4015d2016-06-19 15:50:31 +0100137 m.def("create_rec_packed", &create_recarray<PackedStruct>);
Ivan Smirnov7f913ae2016-06-19 16:41:15 +0100138 m.def("create_rec_nested", &create_nested);
Ivan Smirnov8fa09cb2016-07-06 00:28:12 +0100139 m.def("create_rec_partial", &create_recarray<PartialStruct>);
140 m.def("create_rec_partial_nested", &create_partial_nested);
Ivan Smirnovbdc99022016-06-19 16:54:07 +0100141 m.def("print_format_descriptors", &print_format_descriptors);
Ivan Smirnov4f164212016-06-22 01:07:20 +0100142 m.def("print_rec_simple", &print_recarray<SimpleStruct>);
Ivan Smirnov669e1422016-06-21 21:05:29 +0100143 m.def("print_rec_packed", &print_recarray<PackedStruct>);
144 m.def("print_rec_nested", &print_recarray<NestedStruct>);
Ivan Smirnov4f164212016-06-22 01:07:20 +0100145 m.def("print_dtypes", &print_dtypes);
Ivan Smirnovd0bafd92016-06-26 16:35:28 +0100146 m.def("get_format_unbound", &get_format_unbound);
Ivan Smirnovbb4015d2016-06-19 15:50:31 +0100147}