numpy.h 25.9 KB
Newer Older
Wenzel Jakob's avatar
Wenzel Jakob committed
1
/*
2
    pybind11/numpy.h: Basic NumPy support, vectorize() wrapper
Wenzel Jakob's avatar
Wenzel Jakob committed
3

4
    Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
Wenzel Jakob's avatar
Wenzel Jakob committed
5
6
7
8
9
10
11

    All rights reserved. Use of this source code is governed by a
    BSD-style license that can be found in the LICENSE file.
*/

#pragma once

12
13
#include "pybind11.h"
#include "complex.h"
14
15
#include <numeric>
#include <algorithm>
16
#include <array>
17
#include <cstdlib>
18
#include <cstring>
19
#include <sstream>
20
#include <initializer_list>
21

Wenzel Jakob's avatar
Wenzel Jakob committed
22
23
24
25
26
#if defined(_MSC_VER)
#pragma warning(push)
#pragma warning(disable: 4127) // warning C4127: Conditional expression is constant
#endif

27
NAMESPACE_BEGIN(pybind11)
28
29
namespace detail {
template <typename type, typename SFINAE = void> struct npy_format_descriptor { };
Ivan Smirnov's avatar
Ivan Smirnov committed
30
template <typename type> struct is_pod_struct;
31
}
Wenzel Jakob's avatar
Wenzel Jakob committed
32

Wenzel Jakob's avatar
Wenzel Jakob committed
33
class array : public buffer {
Wenzel Jakob's avatar
Wenzel Jakob committed
34
public:
Wenzel Jakob's avatar
Wenzel Jakob committed
35
36
37
38
39
40
41
    struct API {
        enum Entries {
            API_PyArray_Type = 2,
            API_PyArray_DescrFromType = 45,
            API_PyArray_FromAny = 69,
            API_PyArray_NewCopy = 85,
            API_PyArray_NewFromDescr = 94,
42
            API_PyArray_DescrNewFromType = 9,
43
            API_PyArray_DescrConverter = 174,
44
            API_PyArray_EquivTypes = 182,
45
            API_PyArray_GetArrayParamsFromObject = 278,
46
47
48
49
50
51
52
53
54
55
56
57

            NPY_C_CONTIGUOUS_ = 0x0001,
            NPY_F_CONTIGUOUS_ = 0x0002,
            NPY_ARRAY_FORCECAST_ = 0x0010,
            NPY_ENSURE_ARRAY_ = 0x0040,
            NPY_BOOL_ = 0,
            NPY_BYTE_, NPY_UBYTE_,
            NPY_SHORT_, NPY_USHORT_,
            NPY_INT_, NPY_UINT_,
            NPY_LONG_, NPY_ULONG_,
            NPY_LONGLONG_, NPY_ULONGLONG_,
            NPY_FLOAT_, NPY_DOUBLE_, NPY_LONGDOUBLE_,
58
59
60
            NPY_CFLOAT_, NPY_CDOUBLE_, NPY_CLONGDOUBLE_,
            NPY_OBJECT_ = 17,
            NPY_STRING_, NPY_UNICODE_, NPY_VOID_
Wenzel Jakob's avatar
Wenzel Jakob committed
61
62
63
        };

        static API lookup() {
64
65
            module m = module::import("numpy.core.multiarray");
            object c = (object) m.attr("_ARRAY_API");
66
#if PY_MAJOR_VERSION >= 3
67
            void **api_ptr = (void **) (c ? PyCapsule_GetPointer(c.ptr(), NULL) : nullptr);
68
#else
69
            void **api_ptr = (void **) (c ? PyCObject_AsVoidPtr(c.ptr()) : nullptr);
70
#endif
Wenzel Jakob's avatar
Wenzel Jakob committed
71
            API api;
72
73
74
75
76
77
#define DECL_NPY_API(Func) api.Func##_ = (decltype(api.Func##_)) api_ptr[API_##Func];
            DECL_NPY_API(PyArray_Type);
            DECL_NPY_API(PyArray_DescrFromType);
            DECL_NPY_API(PyArray_FromAny);
            DECL_NPY_API(PyArray_NewCopy);
            DECL_NPY_API(PyArray_NewFromDescr);
78
            DECL_NPY_API(PyArray_DescrNewFromType);
79
            DECL_NPY_API(PyArray_DescrConverter);
80
            DECL_NPY_API(PyArray_EquivTypes);
81
            DECL_NPY_API(PyArray_GetArrayParamsFromObject);
82
#undef DECL_NPY_API
Wenzel Jakob's avatar
Wenzel Jakob committed
83
84
85
            return api;
        }

86
        bool PyArray_Check_(PyObject *obj) const { return (bool) PyObject_TypeCheck(obj, PyArray_Type_); }
Wenzel Jakob's avatar
Wenzel Jakob committed
87

88
89
        PyObject *(*PyArray_DescrFromType_)(int);
        PyObject *(*PyArray_NewFromDescr_)
Wenzel Jakob's avatar
Wenzel Jakob committed
90
91
            (PyTypeObject *, PyObject *, int, Py_intptr_t *,
             Py_intptr_t *, void *, int, PyObject *);
92
        PyObject *(*PyArray_DescrNewFromType_)(int);
93
94
95
        PyObject *(*PyArray_NewCopy_)(PyObject *, int);
        PyTypeObject *PyArray_Type_;
        PyObject *(*PyArray_FromAny_) (PyObject *, PyObject *, int, int, int, PyObject *);
96
        int (*PyArray_DescrConverter_) (PyObject *, PyObject **);
97
        bool (*PyArray_EquivTypes_) (PyObject *, PyObject *);
98
99
        int (*PyArray_GetArrayParamsFromObject_)(PyObject *, PyObject *, char, PyObject **, int *,
                                                 Py_ssize_t *, PyObject **, PyObject *);
Wenzel Jakob's avatar
Wenzel Jakob committed
100
    };
Wenzel Jakob's avatar
Wenzel Jakob committed
101

102
    PYBIND11_OBJECT_DEFAULT(array, buffer, lookup_api().PyArray_Check_)
Wenzel Jakob's avatar
Wenzel Jakob committed
103

104
105
    enum {
        c_style = API::NPY_C_CONTIGUOUS_,
106
107
        f_style = API::NPY_F_CONTIGUOUS_,
        forcecast = API::NPY_ARRAY_FORCECAST_
108
109
    };

Wenzel Jakob's avatar
Wenzel Jakob committed
110
111
    template <typename Type> array(size_t size, const Type *ptr) {
        API& api = lookup_api();
112
        PyObject *descr = detail::npy_format_descriptor<Type>::dtype().release().ptr();
Wenzel Jakob's avatar
Wenzel Jakob committed
113
        Py_intptr_t shape = (Py_intptr_t) size;
114
115
116
        object tmp = object(api.PyArray_NewFromDescr_(
            api.PyArray_Type_, descr, 1, &shape, nullptr, (void *) ptr, 0, nullptr), false);
        if (!tmp)
Wenzel Jakob's avatar
Wenzel Jakob committed
117
            pybind11_fail("NumPy: unable to create array!");
118
119
        if (ptr)
            tmp = object(api.PyArray_NewCopy_(tmp.ptr(), -1 /* any order */), false);
Wenzel Jakob's avatar
Wenzel Jakob committed
120
        m_ptr = tmp.release().ptr();
Wenzel Jakob's avatar
Wenzel Jakob committed
121
122
123
    }

    array(const buffer_info &info) {
124
125
        auto& api = lookup_api();

126
        // _dtype_from_pep3118 returns dtypes with padding fields in, so we need to strip them
127
128
        auto numpy_internal = module::import("numpy.core._internal");
        auto dtype_from_fmt = (object) numpy_internal.attr("_dtype_from_pep3118");
Ivan Smirnov's avatar
Ivan Smirnov committed
129
        auto dtype = strip_padding_fields(dtype_from_fmt(pybind11::str(info.format)));
130
131

        object tmp(api.PyArray_NewFromDescr_(
Ivan Smirnov's avatar
Ivan Smirnov committed
132
            api.PyArray_Type_, dtype.release().ptr(), (int) info.ndim, (Py_intptr_t *) &info.shape[0],
133
134
135
136
137
138
            (Py_intptr_t *) &info.strides[0], info.ptr, 0, nullptr), false);
        if (!tmp)
            pybind11_fail("NumPy: unable to create array!");
        if (info.ptr)
            tmp = object(api.PyArray_NewCopy_(tmp.ptr(), -1 /* any order */), false);
        m_ptr = tmp.release().ptr();
Wenzel Jakob's avatar
Wenzel Jakob committed
139
140
    }

Ivan Smirnov's avatar
Ivan Smirnov committed
141
protected:
Wenzel Jakob's avatar
Wenzel Jakob committed
142
143
144
145
    static API &lookup_api() {
        static API api = API::lookup();
        return api;
    }
146
147

    template <typename T, typename SFINAE> friend struct detail::npy_format_descriptor;
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

    static object strip_padding_fields(object dtype) {
        // Recursively strip all void fields with empty names that are generated for
        // padding fields (as of NumPy v1.11).
        auto fields = dtype.attr("fields").cast<object>();
        if (fields.ptr() == Py_None)
            return dtype;

        struct field_descr { pybind11::str name; object format; int_ offset; };
        std::vector<field_descr> field_descriptors;

        auto items = fields.attr("items").cast<object>();
        for (auto field : items()) {
            auto spec = object(field, true).cast<tuple>();
            auto name = spec[0].cast<pybind11::str>();
            auto format = spec[1].cast<tuple>()[0].cast<object>();
            auto offset = spec[1].cast<tuple>()[1].cast<int_>();
            if (!len(name) && (std::string) dtype.attr("kind").cast<pybind11::str>() == "V")
166
                continue;
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
            field_descriptors.push_back({name, strip_padding_fields(format), offset});
        }

        std::sort(field_descriptors.begin(), field_descriptors.end(),
                  [](const field_descr& a, const field_descr& b) {
                      return (int) a.offset < (int) b.offset;
                  });

        list names, formats, offsets;
        for (auto& descr : field_descriptors) {
            names.append(descr.name);
            formats.append(descr.format);
            offsets.append(descr.offset);
        }
        auto args = dict();
        args["names"] = names; args["formats"] = formats; args["offsets"] = offsets;
        args["itemsize"] = dtype.attr("itemsize").cast<int_>();

        PyObject *descr = nullptr;
        if (!lookup_api().PyArray_DescrConverter_(args.release().ptr(), &descr) || !descr)
            pybind11_fail("NumPy: failed to create structured dtype");
        return object(descr, false);
    }
Wenzel Jakob's avatar
Wenzel Jakob committed
190
191
};

192
template <typename T, int ExtraFlags = array::forcecast> class array_t : public array {
Wenzel Jakob's avatar
Wenzel Jakob committed
193
public:
194
    PYBIND11_OBJECT_CVT(array_t, array, is_non_null, m_ptr = ensure(m_ptr));
195
    array_t() : array() { }
Johan Mabille's avatar
Johan Mabille committed
196
    array_t(const buffer_info& info) : array(info) {}
Wenzel Jakob's avatar
Wenzel Jakob committed
197
    static bool is_non_null(PyObject *ptr) { return ptr != nullptr; }
198
    static PyObject *ensure(PyObject *ptr) {
199
200
        if (ptr == nullptr)
            return nullptr;
Wenzel Jakob's avatar
Wenzel Jakob committed
201
        API &api = lookup_api();
202
        PyObject *descr = detail::npy_format_descriptor<T>::dtype().release().ptr();
203
204
        PyObject *result = api.PyArray_FromAny_(ptr, descr, 0, 0,
                                                API::NPY_ENSURE_ARRAY_ | ExtraFlags, nullptr);
205
206
        if (!result)
            PyErr_Clear();
207
208
        Py_DECREF(ptr);
        return result;
Wenzel Jakob's avatar
Wenzel Jakob committed
209
210
211
    }
};

212
213
template <typename T>
struct format_descriptor<T, typename std::enable_if<detail::is_pod_struct<T>::value>::type> {
214
215
216
217
218
219
220
221
    static const char *format() { return detail::npy_format_descriptor<T>::format(); }
};

template <size_t N> struct format_descriptor<char[N]> {
    static const char *format() { PYBIND11_DESCR s = detail::_<N>() + detail::_("s"); return s.text(); }
};
template <size_t N> struct format_descriptor<std::array<char, N>> {
    static const char *format() { PYBIND11_DESCR s = detail::_<N>() + detail::_("s"); return s.text(); }
222
223
};

224
225
template <typename T>
object dtype_of() {
226
    return detail::npy_format_descriptor<T>::dtype();
227
228
}

229
NAMESPACE_BEGIN(detail)
Ivan Smirnov's avatar
Ivan Smirnov committed
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
template <typename T> struct is_std_array : std::false_type { };
template <typename T, size_t N> struct is_std_array<std::array<T, N>> : std::true_type { };

template <typename T>
struct is_pod_struct {
    enum { value = std::is_pod<T>::value && // offsetof only works correctly for POD types
           !std::is_array<T>::value &&
           !is_std_array<T>::value &&
           !std::is_integral<T>::value &&
           !std::is_same<T, float>::value &&
           !std::is_same<T, double>::value &&
           !std::is_same<T, bool>::value &&
           !std::is_same<T, std::complex<float>>::value &&
           !std::is_same<T, std::complex<double>>::value };
};
245

246
247
template <typename T> struct npy_format_descriptor<T, typename std::enable_if<std::is_integral<T>::value>::type> {
private:
Johan Mabille's avatar
Johan Mabille committed
248
    constexpr static const int values[8] = {
249
250
251
        array::API::NPY_BYTE_, array::API::NPY_UBYTE_, array::API::NPY_SHORT_,    array::API::NPY_USHORT_,
        array::API::NPY_INT_,  array::API::NPY_UINT_,  array::API::NPY_LONGLONG_, array::API::NPY_ULONGLONG_ };
public:
252
253
254
255
256
    enum { value = values[detail::log2(sizeof(T)) * 2 + (std::is_unsigned<T>::value ? 1 : 0)] };
    static object dtype() {
        if (auto ptr = array::lookup_api().PyArray_DescrFromType_(value))
            return object(ptr, true);
        pybind11_fail("Unsupported buffer format!");
257
    }
258
259
260
261
    template <typename T2 = T, typename std::enable_if<std::is_signed<T2>::value, int>::type = 0>
    static PYBIND11_DESCR name() { return _("int") + _<sizeof(T)*8>(); }
    template <typename T2 = T, typename std::enable_if<!std::is_signed<T2>::value, int>::type = 0>
    static PYBIND11_DESCR name() { return _("uint") + _<sizeof(T)*8>(); }
262
263
264
265
};
template <typename T> constexpr const int npy_format_descriptor<
    T, typename std::enable_if<std::is_integral<T>::value>::type>::values[8];

266
#define DECL_FMT(Type, NumPyName, Name) template<> struct npy_format_descriptor<Type> { \
267
268
269
270
271
    enum { value = array::API::NumPyName }; \
    static object dtype() { \
        if (auto ptr = array::lookup_api().PyArray_DescrFromType_(value)) \
            return object(ptr, true); \
        pybind11_fail("Unsupported buffer format!"); \
272
    } \
273
    static PYBIND11_DESCR name() { return _(Name); } }
274
275
276
277
278
DECL_FMT(float, NPY_FLOAT_, "float32");
DECL_FMT(double, NPY_DOUBLE_, "float64");
DECL_FMT(bool, NPY_BOOL_, "bool");
DECL_FMT(std::complex<float>, NPY_CFLOAT_, "complex64");
DECL_FMT(std::complex<double>, NPY_CDOUBLE_, "complex128");
279
280
#undef DECL_FMT

281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
#define DECL_CHAR_FMT \
    static PYBIND11_DESCR name() { return _("S") + _<N>(); } \
    static object dtype() { \
        auto& api = array::lookup_api(); \
        PyObject *descr = nullptr; \
        PYBIND11_DESCR fmt = _("S") + _<N>(); \
        pybind11::str py_fmt(fmt.text()); \
        if (!api.PyArray_DescrConverter_(py_fmt.release().ptr(), &descr) || !descr) \
            pybind11_fail("NumPy: failed to create string dtype"); \
        return object(descr, false); \
    } \
    static const char *format() { PYBIND11_DESCR s = _<N>() + _("s"); return s.text(); }
template <size_t N> struct npy_format_descriptor<char[N]> { DECL_CHAR_FMT };
template <size_t N> struct npy_format_descriptor<std::array<char, N>> { DECL_CHAR_FMT };
#undef DECL_CHAR_FMT

297
298
struct field_descriptor {
    const char *name;
299
    size_t offset;
300
301
    size_t size;
    const char *format;
302
    object descr;
303
304
};

305
template <typename T>
Ivan Smirnov's avatar
Ivan Smirnov committed
306
struct npy_format_descriptor<T, typename std::enable_if<is_pod_struct<T>::value>::type> {
307
    static PYBIND11_DESCR name() { return _("struct"); }
308

309
310
    static object dtype() {
        if (!dtype_())
311
            pybind11_fail("NumPy: unsupported buffer format!");
312
        return object(dtype_(), true);
313
314
    }

315
    static const char* format() {
316
317
        if (!dtype_())
            pybind11_fail("NumPy: unsupported buffer format!");
318
        return format_().c_str();
319
320
321
    }

    static void register_dtype(std::initializer_list<field_descriptor> fields) {
322
        auto& api = array::lookup_api();
323
324
        auto args = dict();
        list names { }, offsets { }, formats { };
325
326
327
        for (auto field : fields) {
            if (!field.descr)
                pybind11_fail("NumPy: unsupported field dtype");
328
329
            names.append(str(field.name));
            offsets.append(int_(field.offset));
330
            formats.append(field.descr);
331
        }
332
        args["names"] = names; args["offsets"] = offsets; args["formats"] = formats;
333
        args["itemsize"] = int_(sizeof(T));
334
335
        // This is essentially the same as calling np.dtype() constructor in Python and passing
        // it a dict of the form {'names': ..., 'formats': ..., 'offsets': ...}.
336
        if (!api.PyArray_DescrConverter_(args.release().ptr(), &dtype_()) || !dtype_())
337
            pybind11_fail("NumPy: failed to create structured dtype");
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359

        // There is an existing bug in NumPy (as of v1.11): trailing bytes are
        // not encoded explicitly into the format string. This will supposedly
        // get fixed in v1.12; for further details, see these:
        // - https://github.com/numpy/numpy/issues/7797
        // - https://github.com/numpy/numpy/pull/7798
        // Because of this, we won't use numpy's logic to generate buffer format
        // strings and will just do it ourselves.
        std::vector<field_descriptor> ordered_fields(fields);
        std::sort(ordered_fields.begin(), ordered_fields.end(),
                  [](const field_descriptor& a, const field_descriptor &b) {
                      return a.offset < b.offset;
                  });
        size_t offset = 0;
        std::ostringstream oss;
        oss << "T{";
        for (auto& field : ordered_fields) {
            if (field.offset > offset)
                oss << (field.offset - offset) << 'x';
            // note that '=' is required to cover the case of unaligned fields
            oss << '=' << field.format << ':' << field.name << ':';
            offset = field.offset + field.size;
Ivan Smirnov's avatar
Ivan Smirnov committed
360
        }
361
362
363
        if (sizeof(T) > offset)
            oss << (sizeof(T) - offset) << 'x';
        oss << '}';
364
        format_() = oss.str();
365
366
367

        // Sanity check: verify that NumPy properly parses our buffer format string
        auto arr =  array(buffer_info(nullptr, sizeof(T), format(), 1, { 0 }, { sizeof(T) }));
368
369
370
        auto fixed_dtype = array::strip_padding_fields(object(dtype_(), true));
        if (!api.PyArray_EquivTypes_(dtype_(), fixed_dtype.ptr()))
            pybind11_fail("NumPy: invalid buffer descriptor!");
371
372
373
    }

private:
374
    static inline PyObject*& dtype_() { static PyObject *ptr = nullptr; return ptr; }
375
    static inline std::string& format_() { static std::string s; return s; }
376
377
};

378
// Extract name, offset and format descriptor for a struct field
379
#define PYBIND11_FIELD_DESCRIPTOR(Type, Field) \
380
    ::pybind11::detail::field_descriptor { \
381
382
        #Field, offsetof(Type, Field), sizeof(decltype(static_cast<Type*>(0)->Field)), \
        ::pybind11::format_descriptor<decltype(static_cast<Type*>(0)->Field)>::format(), \
383
        ::pybind11::detail::npy_format_descriptor<decltype(static_cast<Type*>(0)->Field)>::dtype() \
384
    }
385
386
387

// The main idea of this macro is borrowed from https://github.com/swansontec/map-macro
// (C) William Swanson, Paul Fultz
388
389
390
391
392
393
394
395
396
397
398
399
400
#define PYBIND11_EVAL0(...) __VA_ARGS__
#define PYBIND11_EVAL1(...) PYBIND11_EVAL0 (PYBIND11_EVAL0 (PYBIND11_EVAL0 (__VA_ARGS__)))
#define PYBIND11_EVAL2(...) PYBIND11_EVAL1 (PYBIND11_EVAL1 (PYBIND11_EVAL1 (__VA_ARGS__)))
#define PYBIND11_EVAL3(...) PYBIND11_EVAL2 (PYBIND11_EVAL2 (PYBIND11_EVAL2 (__VA_ARGS__)))
#define PYBIND11_EVAL4(...) PYBIND11_EVAL3 (PYBIND11_EVAL3 (PYBIND11_EVAL3 (__VA_ARGS__)))
#define PYBIND11_EVAL(...)  PYBIND11_EVAL4 (PYBIND11_EVAL4 (PYBIND11_EVAL4 (__VA_ARGS__)))
#define PYBIND11_MAP_END(...)
#define PYBIND11_MAP_OUT
#define PYBIND11_MAP_COMMA ,
#define PYBIND11_MAP_GET_END() 0, PYBIND11_MAP_END
#define PYBIND11_MAP_NEXT0(test, next, ...) next PYBIND11_MAP_OUT
#define PYBIND11_MAP_NEXT1(test, next) PYBIND11_MAP_NEXT0 (test, next, 0)
#define PYBIND11_MAP_NEXT(test, next)  PYBIND11_MAP_NEXT1 (PYBIND11_MAP_GET_END test, next)
401
#ifdef _MSC_VER // MSVC is not as eager to expand macros, hence this workaround
402
403
#define PYBIND11_MAP_LIST_NEXT1(test, next) \
    PYBIND11_EVAL0 (PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0))
404
#else
405
406
#define PYBIND11_MAP_LIST_NEXT1(test, next) \
    PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0)
407
#endif
408
409
410
411
412
413
#define PYBIND11_MAP_LIST_NEXT(test, next) \
    PYBIND11_MAP_LIST_NEXT1 (PYBIND11_MAP_GET_END test, next)
#define PYBIND11_MAP_LIST0(f, t, x, peek, ...) \
    f(t, x) PYBIND11_MAP_LIST_NEXT (peek, PYBIND11_MAP_LIST1) (f, t, peek, __VA_ARGS__)
#define PYBIND11_MAP_LIST1(f, t, x, peek, ...) \
    f(t, x) PYBIND11_MAP_LIST_NEXT (peek, PYBIND11_MAP_LIST0) (f, t, peek, __VA_ARGS__)
414
// PYBIND11_MAP_LIST(f, t, a1, a2, ...) expands to f(t, a1), f(t, a2), ...
415
416
#define PYBIND11_MAP_LIST(f, t, ...) \
    PYBIND11_EVAL (PYBIND11_MAP_LIST1 (f, t, __VA_ARGS__, (), 0))
417

418
#define PYBIND11_NUMPY_DTYPE(Type, ...) \
419
    ::pybind11::detail::npy_format_descriptor<Type>::register_dtype \
420
        ({PYBIND11_MAP_LIST (PYBIND11_FIELD_DESCRIPTOR, Type, __VA_ARGS__)})
421

422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
template  <class T>
using array_iterator = typename std::add_pointer<T>::type;

template <class T>
array_iterator<T> array_begin(const buffer_info& buffer) {
    return array_iterator<T>(reinterpret_cast<T*>(buffer.ptr));
}

template <class T>
array_iterator<T> array_end(const buffer_info& buffer) {
    return array_iterator<T>(reinterpret_cast<T*>(buffer.ptr) + buffer.size);
}

class common_iterator {
public:
    using container_type = std::vector<size_t>;
    using value_type = container_type::value_type;
    using size_type = container_type::size_type;

    common_iterator() : p_ptr(0), m_strides() {}
442

443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
    common_iterator(void* ptr, const container_type& strides, const std::vector<size_t>& shape)
        : p_ptr(reinterpret_cast<char*>(ptr)), m_strides(strides.size()) {
        m_strides.back() = static_cast<value_type>(strides.back());
        for (size_type i = m_strides.size() - 1; i != 0; --i) {
            size_type j = i - 1;
            value_type s = static_cast<value_type>(shape[i]);
            m_strides[j] = strides[j] + m_strides[i] - strides[i] * s;
        }
    }

    void increment(size_type dim) {
        p_ptr += m_strides[dim];
    }

    void* data() const {
        return p_ptr;
    }

private:
    char* p_ptr;
    container_type m_strides;
};

466
template <size_t N> class multi_array_iterator {
467
468
469
public:
    using container_type = std::vector<size_t>;

470
471
472
473
474
    multi_array_iterator(const std::array<buffer_info, N> &buffers,
                         const std::vector<size_t> &shape)
        : m_shape(shape.size()), m_index(shape.size(), 0),
          m_common_iterator() {

475
        // Manual copy to avoid conversion warning if using std::copy
476
        for (size_t i = 0; i < shape.size(); ++i)
477
478
479
            m_shape[i] = static_cast<container_type::value_type>(shape[i]);

        container_type strides(shape.size());
480
        for (size_t i = 0; i < N; ++i)
481
482
483
484
485
486
487
488
489
            init_common_iterator(buffers[i], shape, m_common_iterator[i], strides);
    }

    multi_array_iterator& operator++() {
        for (size_t j = m_index.size(); j != 0; --j) {
            size_t i = j - 1;
            if (++m_index[i] != m_shape[i]) {
                increment_common_iterator(i);
                break;
490
            } else {
491
492
493
494
495
496
                m_index[i] = 0;
            }
        }
        return *this;
    }

497
    template <size_t K, class T> const T& data() const {
498
499
500
501
502
503
504
        return *reinterpret_cast<T*>(m_common_iterator[K].data());
    }

private:

    using common_iter = common_iterator;

505
506
507
    void init_common_iterator(const buffer_info &buffer,
                              const std::vector<size_t> &shape,
                              common_iter &iterator, container_type &strides) {
508
509
510
511
512
513
514
        auto buffer_shape_iter = buffer.shape.rbegin();
        auto buffer_strides_iter = buffer.strides.rbegin();
        auto shape_iter = shape.rbegin();
        auto strides_iter = strides.rbegin();

        while (buffer_shape_iter != buffer.shape.rend()) {
            if (*shape_iter == *buffer_shape_iter)
515
                *strides_iter = static_cast<size_t>(*buffer_strides_iter);
516
517
518
519
520
521
522
523
524
525
526
527
528
529
            else
                *strides_iter = 0;

            ++buffer_shape_iter;
            ++buffer_strides_iter;
            ++shape_iter;
            ++strides_iter;
        }

        std::fill(strides_iter, strides.rend(), 0);
        iterator = common_iter(buffer.ptr, strides, shape);
    }

    void increment_common_iterator(size_t dim) {
530
        for (auto &iter : m_common_iterator)
531
532
533
534
535
536
537
538
539
            iter.increment(dim);
    }

    container_type m_shape;
    container_type m_index;
    std::array<common_iter, N> m_common_iterator;
};

template <size_t N>
540
541
bool broadcast(const std::array<buffer_info, N>& buffers, size_t& ndim, std::vector<size_t>& shape) {
    ndim = std::accumulate(buffers.begin(), buffers.end(), size_t(0), [](size_t res, const buffer_info& buf) {
542
543
544
        return std::max(res, buf.ndim);
    });

545
    shape = std::vector<size_t>(ndim, 1);
546
547
548
549
    bool trivial_broadcast = true;
    for (size_t i = 0; i < N; ++i) {
        auto res_iter = shape.rbegin();
        bool i_trivial_broadcast = (buffers[i].size == 1) || (buffers[i].ndim == ndim);
550
551
552
553
        for (auto shape_iter = buffers[i].shape.rbegin();
             shape_iter != buffers[i].shape.rend(); ++shape_iter, ++res_iter) {

            if (*res_iter == 1)
554
                *res_iter = *shape_iter;
555
            else if ((*shape_iter != 1) && (*res_iter != *shape_iter))
556
                pybind11_fail("pybind11::vectorize: incompatible size/dimension of inputs!");
557

558
559
560
561
562
563
564
            i_trivial_broadcast = i_trivial_broadcast && (*res_iter == *shape_iter);
        }
        trivial_broadcast = trivial_broadcast && i_trivial_broadcast;
    }
    return trivial_broadcast;
}

565
566
567
568
template <typename Func, typename Return, typename... Args>
struct vectorize_helper {
    typename std::remove_reference<Func>::type f;

569
570
    template <typename T>
    vectorize_helper(T&&f) : f(std::forward<T>(f)) { }
Wenzel Jakob's avatar
Wenzel Jakob committed
571

572
    object operator()(array_t<Args, array::c_style | array::forcecast>... args) {
573
574
        return run(args..., typename make_index_sequence<sizeof...(Args)>::type());
    }
Wenzel Jakob's avatar
Wenzel Jakob committed
575

576
    template <size_t ... Index> object run(array_t<Args, array::c_style | array::forcecast>&... args, index_sequence<Index...> index) {
Wenzel Jakob's avatar
Wenzel Jakob committed
577
        /* Request buffers from all parameters */
578
        const size_t N = sizeof...(Args);
579

Wenzel Jakob's avatar
Wenzel Jakob committed
580
581
582
        std::array<buffer_info, N> buffers {{ args.request()... }};

        /* Determine dimensions parameters of output array */
583
        size_t ndim = 0;
584
585
        std::vector<size_t> shape(0);
        bool trivial_broadcast = broadcast(buffers, ndim, shape);
586

587
        size_t size = 1;
Wenzel Jakob's avatar
Wenzel Jakob committed
588
589
        std::vector<size_t> strides(ndim);
        if (ndim > 0) {
590
            strides[ndim-1] = sizeof(Return);
591
            for (size_t i = ndim - 1; i > 0; --i) {
592
593
594
595
                strides[i - 1] = strides[i] * shape[i];
                size *= shape[i];
            }
            size *= shape[0];
Wenzel Jakob's avatar
Wenzel Jakob committed
596
597
        }

598
        if (size == 1)
599
            return cast(f(*((Args *) buffers[Index].ptr)...));
Wenzel Jakob's avatar
Wenzel Jakob committed
600

601
        array result(buffer_info(nullptr, sizeof(Return),
602
            format_descriptor<Return>::format(),
Wenzel Jakob's avatar
Wenzel Jakob committed
603
            ndim, shape, strides));
604
605
606
607

        buffer_info buf = result.request();
        Return *output = (Return *) buf.ptr;

608
        if (trivial_broadcast) {
609
610
611
            /* Call the function */
            for (size_t i=0; i<size; ++i) {
                output[i] = f((buffers[Index].size == 1
612
613
                               ? *((Args *) buffers[Index].ptr)
                               : ((Args *) buffers[Index].ptr)[i])...);
614
            }
615
        } else {
616
617
            apply_broadcast<N, Index...>(buffers, buf, index);
        }
618
619

        return result;
620
    }
621
622

    template <size_t N, size_t... Index>
623
624
    void apply_broadcast(const std::array<buffer_info, N> &buffers,
                         buffer_info &output, index_sequence<Index...>) {
625
626
627
628
629
630
        using input_iterator = multi_array_iterator<N>;
        using output_iterator = array_iterator<Return>;

        input_iterator input_iter(buffers, output.shape);
        output_iterator output_end = array_end<Return>(output);

631
632
        for (output_iterator iter = array_begin<Return>(output);
             iter != output_end; ++iter, ++input_iter) {
633
634
635
            *iter = f((input_iter.template data<Index, Args>())...);
        }
    }
636
637
};

638
template <typename T, int Flags> struct handle_type_name<array_t<T, Flags>> {
639
    static PYBIND11_DESCR name() { return _("numpy.ndarray[") + type_caster<T>::name() + _("]"); }
640
641
};

642
NAMESPACE_END(detail)
Wenzel Jakob's avatar
Wenzel Jakob committed
643

644
645
646
template <typename Func, typename Return, typename... Args>
detail::vectorize_helper<Func, Return, Args...> vectorize(const Func &f, Return (*) (Args ...)) {
    return detail::vectorize_helper<Func, Return, Args...>(f);
Wenzel Jakob's avatar
Wenzel Jakob committed
647
648
}

649
650
651
template <typename Return, typename... Args>
detail::vectorize_helper<Return (*) (Args ...), Return, Args...> vectorize(Return (*f) (Args ...)) {
    return vectorize<Return (*) (Args ...), Return, Args...>(f, f);
Wenzel Jakob's avatar
Wenzel Jakob committed
652
653
654
655
656
657
658
659
}

template <typename func> auto vectorize(func &&f) -> decltype(
        vectorize(std::forward<func>(f), (typename detail::remove_class<decltype(&std::remove_reference<func>::type::operator())>::type *) nullptr)) {
    return vectorize(std::forward<func>(f), (typename detail::remove_class<decltype(
                   &std::remove_reference<func>::type::operator())>::type *) nullptr);
}

660
NAMESPACE_END(pybind11)
Wenzel Jakob's avatar
Wenzel Jakob committed
661
662
663
664

#if defined(_MSC_VER)
#pragma warning(pop)
#endif