numpy.h 54.5 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 <string>
21
#include <initializer_list>
22
#include <functional>
23
#include <utility>
24
#include <typeindex>
25

Wenzel Jakob's avatar
Wenzel Jakob committed
26
#if defined(_MSC_VER)
27
28
#  pragma warning(push)
#  pragma warning(disable: 4127) // warning C4127: Conditional expression is constant
Wenzel Jakob's avatar
Wenzel Jakob committed
29
30
#endif

31
32
33
34
35
36
/* This will be true on all flat address space platforms and allows us to reduce the
   whole npy_intp / size_t / Py_intptr_t business down to just size_t for all size
   and dimension types (e.g. shape, strides, indexing), instead of inflicting this
   upon the library user. */
static_assert(sizeof(size_t) == sizeof(Py_intptr_t), "size_t != Py_intptr_t");

37
NAMESPACE_BEGIN(pybind11)
38
39
40

class array; // Forward declaration

41
NAMESPACE_BEGIN(detail)
42
template <typename type, typename SFINAE = void> struct npy_format_descriptor;
Wenzel Jakob's avatar
Wenzel Jakob committed
43

44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
struct PyArrayDescr_Proxy {
    PyObject_HEAD
    PyObject *typeobj;
    char kind;
    char type;
    char byteorder;
    char flags;
    int type_num;
    int elsize;
    int alignment;
    char *subarray;
    PyObject *fields;
    PyObject *names;
};

struct PyArray_Proxy {
    PyObject_HEAD
    char *data;
    int nd;
    ssize_t *dimensions;
    ssize_t *strides;
    PyObject *base;
    PyObject *descr;
    int flags;
};

70
71
72
73
74
75
76
77
struct PyVoidScalarObject_Proxy {
    PyObject_VAR_HEAD
    char *obval;
    PyArrayDescr_Proxy *descr;
    int flags;
    PyObject *base;
};

78
79
80
81
82
83
84
85
struct numpy_type_info {
    PyObject* dtype_ptr;
    std::string format_str;
};

struct numpy_internals {
    std::unordered_map<std::type_index, numpy_type_info> registered_dtypes;

86
87
    numpy_type_info *get_type_info(const std::type_info& tinfo, bool throw_if_missing = true) {
        auto it = registered_dtypes.find(std::type_index(tinfo));
88
89
90
        if (it != registered_dtypes.end())
            return &(it->second);
        if (throw_if_missing)
91
            pybind11_fail(std::string("NumPy type info missing for ") + tinfo.name());
92
93
        return nullptr;
    }
94
95
96
97

    template<typename T> numpy_type_info *get_type_info(bool throw_if_missing = true) {
        return get_type_info(typeid(typename std::remove_cv<T>::type), throw_if_missing);
    }
98
99
};

Ivan Smirnov's avatar
Ivan Smirnov committed
100
101
inline PYBIND11_NOINLINE void load_numpy_internals(numpy_internals* &ptr) {
    ptr = &get_or_create_shared_data<numpy_internals>("_numpy_internals");
102
103
104
}

inline numpy_internals& get_numpy_internals() {
Ivan Smirnov's avatar
Ivan Smirnov committed
105
106
107
    static numpy_internals* ptr = nullptr;
    if (!ptr)
        load_numpy_internals(ptr);
108
109
110
    return *ptr;
}

111
112
struct npy_api {
    enum constants {
113
114
        NPY_ARRAY_C_CONTIGUOUS_ = 0x0001,
        NPY_ARRAY_F_CONTIGUOUS_ = 0x0002,
115
        NPY_ARRAY_OWNDATA_ = 0x0004,
116
        NPY_ARRAY_FORCECAST_ = 0x0010,
117
        NPY_ARRAY_ENSUREARRAY_ = 0x0040,
118
119
        NPY_ARRAY_ALIGNED_ = 0x0100,
        NPY_ARRAY_WRITEABLE_ = 0x0400,
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
        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_,
        NPY_CFLOAT_, NPY_CDOUBLE_, NPY_CLONGDOUBLE_,
        NPY_OBJECT_ = 17,
        NPY_STRING_, NPY_UNICODE_, NPY_VOID_
    };

    static npy_api& get() {
        static npy_api api = lookup();
        return api;
    }

137
138
139
140
141
142
    bool PyArray_Check_(PyObject *obj) const {
        return (bool) PyObject_TypeCheck(obj, PyArray_Type_);
    }
    bool PyArrayDescr_Check_(PyObject *obj) const {
        return (bool) PyObject_TypeCheck(obj, PyArrayDescr_Type_);
    }
143

144
    unsigned int (*PyArray_GetNDArrayCFeatureVersion_)();
145
146
147
148
149
150
151
    PyObject *(*PyArray_DescrFromType_)(int);
    PyObject *(*PyArray_NewFromDescr_)
        (PyTypeObject *, PyObject *, int, Py_intptr_t *,
         Py_intptr_t *, void *, int, PyObject *);
    PyObject *(*PyArray_DescrNewFromType_)(int);
    PyObject *(*PyArray_NewCopy_)(PyObject *, int);
    PyTypeObject *PyArray_Type_;
152
    PyTypeObject *PyVoidArrType_Type_;
153
    PyTypeObject *PyArrayDescr_Type_;
154
    PyObject *(*PyArray_DescrFromScalar_)(PyObject *);
155
156
157
158
159
    PyObject *(*PyArray_FromAny_) (PyObject *, PyObject *, int, int, int, PyObject *);
    int (*PyArray_DescrConverter_) (PyObject *, PyObject **);
    bool (*PyArray_EquivTypes_) (PyObject *, PyObject *);
    int (*PyArray_GetArrayParamsFromObject_)(PyObject *, PyObject *, char, PyObject **, int *,
                                             Py_ssize_t *, PyObject **, PyObject *);
160
    PyObject *(*PyArray_Squeeze_)(PyObject *);
Jason Rhinelander's avatar
Jason Rhinelander committed
161
    int (*PyArray_SetBaseObject_)(PyObject *, PyObject *);
162
163
private:
    enum functions {
164
        API_PyArray_GetNDArrayCFeatureVersion = 211,
165
        API_PyArray_Type = 2,
166
        API_PyArrayDescr_Type = 3,
167
        API_PyVoidArrType_Type = 39,
168
        API_PyArray_DescrFromType = 45,
169
        API_PyArray_DescrFromScalar = 57,
170
171
172
173
174
175
176
        API_PyArray_FromAny = 69,
        API_PyArray_NewCopy = 85,
        API_PyArray_NewFromDescr = 94,
        API_PyArray_DescrNewFromType = 9,
        API_PyArray_DescrConverter = 174,
        API_PyArray_EquivTypes = 182,
        API_PyArray_GetArrayParamsFromObject = 278,
Jason Rhinelander's avatar
Jason Rhinelander committed
177
178
        API_PyArray_Squeeze = 136,
        API_PyArray_SetBaseObject = 282
179
180
181
182
    };

    static npy_api lookup() {
        module m = module::import("numpy.core.multiarray");
183
        auto c = m.attr("_ARRAY_API");
184
#if PY_MAJOR_VERSION >= 3
185
        void **api_ptr = (void **) PyCapsule_GetPointer(c.ptr(), NULL);
186
#else
187
        void **api_ptr = (void **) PyCObject_AsVoidPtr(c.ptr());
188
#endif
189
        npy_api api;
190
#define DECL_NPY_API(Func) api.Func##_ = (decltype(api.Func##_)) api_ptr[API_##Func];
191
192
193
        DECL_NPY_API(PyArray_GetNDArrayCFeatureVersion);
        if (api.PyArray_GetNDArrayCFeatureVersion_() < 0x7)
            pybind11_fail("pybind11 numpy support requires numpy >= 1.7.0");
194
        DECL_NPY_API(PyArray_Type);
195
        DECL_NPY_API(PyVoidArrType_Type);
196
        DECL_NPY_API(PyArrayDescr_Type);
197
        DECL_NPY_API(PyArray_DescrFromType);
198
        DECL_NPY_API(PyArray_DescrFromScalar);
199
200
201
202
203
204
205
        DECL_NPY_API(PyArray_FromAny);
        DECL_NPY_API(PyArray_NewCopy);
        DECL_NPY_API(PyArray_NewFromDescr);
        DECL_NPY_API(PyArray_DescrNewFromType);
        DECL_NPY_API(PyArray_DescrConverter);
        DECL_NPY_API(PyArray_EquivTypes);
        DECL_NPY_API(PyArray_GetArrayParamsFromObject);
206
        DECL_NPY_API(PyArray_Squeeze);
Jason Rhinelander's avatar
Jason Rhinelander committed
207
        DECL_NPY_API(PyArray_SetBaseObject);
208
#undef DECL_NPY_API
209
210
211
        return api;
    }
};
Wenzel Jakob's avatar
Wenzel Jakob committed
212

213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
inline PyArray_Proxy* array_proxy(void* ptr) {
    return reinterpret_cast<PyArray_Proxy*>(ptr);
}

inline const PyArray_Proxy* array_proxy(const void* ptr) {
    return reinterpret_cast<const PyArray_Proxy*>(ptr);
}

inline PyArrayDescr_Proxy* array_descriptor_proxy(PyObject* ptr) {
   return reinterpret_cast<PyArrayDescr_Proxy*>(ptr);
}

inline const PyArrayDescr_Proxy* array_descriptor_proxy(const PyObject* ptr) {
   return reinterpret_cast<const PyArrayDescr_Proxy*>(ptr);
}

inline bool check_flags(const void* ptr, int flag) {
    return (flag == (array_proxy(ptr)->flags & flag));
}

233
234
235
236
237
238
239
240
241
242
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_complex : std::false_type { };
template <typename T> struct is_complex<std::complex<T>> : std::true_type { };

template <typename T> using is_pod_struct = all_of<
    std::is_pod<T>, // since we're accessing directly in memory we need a POD type
    satisfies_none_of<T, std::is_reference, std::is_array, is_std_array, std::is_arithmetic, is_complex, std::is_enum>
>;

243
244
245
246
247
248
249
template <size_t Dim = 0, typename Strides> size_t byte_offset_unsafe(const Strides &) { return 0; }
template <size_t Dim = 0, typename Strides, typename... Ix>
size_t byte_offset_unsafe(const Strides &strides, size_t i, Ix... index) {
    return i * strides[Dim] + byte_offset_unsafe<Dim + 1>(strides, index...);
}

/** Proxy class providing unsafe, unchecked const access to array data.  This is constructed through
250
251
 * the `unchecked<T, N>()` method of `array` or the `unchecked<N>()` method of `array_t<T>`.  `Dims`
 * will be -1 for dimensions determined at runtime.
252
 */
253
template <typename T, ssize_t Dims>
254
255
class unchecked_reference {
protected:
256
    static constexpr bool Dynamic = Dims < 0;
257
258
    const unsigned char *data_;
    // Storing the shape & strides in local variables (i.e. these arrays) allows the compiler to
259
260
261
262
    // make large performance gains on big, nested loops, but requires compile-time dimensions
    conditional_t<Dynamic, const size_t *, std::array<size_t, (size_t) Dims>>
        shape_, strides_;
    const size_t dims_;
263
264

    friend class pybind11::array;
265
266
267
268
269
    // Constructor for compile-time dimensions:
    template <bool Dyn = Dynamic>
    unchecked_reference(const void *data, const size_t *shape, const size_t *strides, enable_if_t<!Dyn, size_t>)
    : data_{reinterpret_cast<const unsigned char *>(data)}, dims_{Dims} {
        for (size_t i = 0; i < dims_; i++) {
270
271
272
273
            shape_[i] = shape[i];
            strides_[i] = strides[i];
        }
    }
274
275
276
277
    // Constructor for runtime dimensions:
    template <bool Dyn = Dynamic>
    unchecked_reference(const void *data, const size_t *shape, const size_t *strides, enable_if_t<Dyn, size_t> dims)
    : data_{reinterpret_cast<const unsigned char *>(data)}, shape_{shape}, strides_{strides}, dims_{dims} {}
278
279

public:
280
281
282
    /** Unchecked const reference access to data at the given indices.  For a compile-time known
     * number of dimensions, this requires the correct number of arguments; for run-time
     * dimensionality, this is not checked (and so is up to the caller to use safely).
283
     */
284
285
286
287
    template <typename... Ix> const T &operator()(Ix... index) const {
        static_assert(sizeof...(Ix) == Dims || Dynamic,
                "Invalid number of indices for unchecked array reference");
        return *reinterpret_cast<const T *>(data_ + byte_offset_unsafe(strides_, size_t(index)...));
288
289
290
291
    }
    /** Unchecked const reference access to data; this operator only participates if the reference
     * is to a 1-dimensional array.  When present, this is exactly equivalent to `obj(index)`.
     */
292
    template <size_t D = Dims, typename = enable_if_t<D == 1 || Dynamic>>
293
294
    const T &operator[](size_t index) const { return operator()(index); }

295
296
297
298
299
300
    /// Pointer access to the data at the given indices.
    template <typename... Ix> const T *data(Ix... ix) const { return &operator()(size_t(ix)...); }

    /// Returns the item size, i.e. sizeof(T)
    constexpr static size_t itemsize() { return sizeof(T); }

301
302
303
304
    /// Returns the shape (i.e. size) of dimension `dim`
    size_t shape(size_t dim) const { return shape_[dim]; }

    /// Returns the number of dimensions of the array
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
    size_t ndim() const { return dims_; }

    /// Returns the total number of elements in the referenced array, i.e. the product of the shapes
    template <bool Dyn = Dynamic>
    enable_if_t<!Dyn, size_t> size() const {
        return std::accumulate(shape_.begin(), shape_.end(), (size_t) 1, std::multiplies<size_t>());
    }
    template <bool Dyn = Dynamic>
    enable_if_t<Dyn, size_t> size() const {
        return std::accumulate(shape_, shape_ + ndim(), (size_t) 1, std::multiplies<size_t>());
    }

    /// Returns the total number of bytes used by the referenced data.  Note that the actual span in
    /// memory may be larger if the referenced array has non-contiguous strides (e.g. for a slice).
    size_t nbytes() const {
        return size() * itemsize();
    }
322
323
};

324
template <typename T, ssize_t Dims>
325
326
327
328
class unchecked_mutable_reference : public unchecked_reference<T, Dims> {
    friend class pybind11::array;
    using ConstBase = unchecked_reference<T, Dims>;
    using ConstBase::ConstBase;
329
    using ConstBase::Dynamic;
330
331
332
public:
    /// Mutable, unchecked access to data at the given indices.
    template <typename... Ix> T& operator()(Ix... index) {
333
334
        static_assert(sizeof...(Ix) == Dims || Dynamic,
                "Invalid number of indices for unchecked array reference");
335
336
337
        return const_cast<T &>(ConstBase::operator()(index...));
    }
    /** Mutable, unchecked access data at the given index; this operator only participates if the
338
339
     * reference is to a 1-dimensional array (or has runtime dimensions).  When present, this is
     * exactly equivalent to `obj(index)`.
340
     */
341
    template <size_t D = Dims, typename = enable_if_t<D == 1 || Dynamic>>
342
    T &operator[](size_t index) { return operator()(index); }
343
344
345

    /// Mutable pointer access to the data at the given indices.
    template <typename... Ix> T *mutable_data(Ix... ix) { return &operator()(size_t(ix)...); }
346
347
};

348
template <typename T, ssize_t Dim>
349
struct type_caster<unchecked_reference<T, Dim>> {
350
    static_assert(Dim == 0 && Dim > 0 /* always fail */, "unchecked array proxy object is not castable");
351
};
352
template <typename T, ssize_t Dim>
353
354
struct type_caster<unchecked_mutable_reference<T, Dim>> : type_caster<unchecked_reference<T, Dim>> {};

355
NAMESPACE_END(detail)
356

357
class dtype : public object {
358
public:
359
    PYBIND11_OBJECT_DEFAULT(dtype, object, detail::npy_api::get().PyArrayDescr_Check_);
Wenzel Jakob's avatar
Wenzel Jakob committed
360

361
    explicit dtype(const buffer_info &info) {
362
        dtype descr(_dtype_from_pep3118()(PYBIND11_STR_TYPE(info.format)));
363
364
        // If info.itemsize == 0, use the value calculated from the format string
        m_ptr = descr.strip_padding(info.itemsize ? info.itemsize : descr.itemsize()).release().ptr();
365
    }
366

367
    explicit dtype(const std::string &format) {
368
        m_ptr = from_args(pybind11::str(format)).release().ptr();
Wenzel Jakob's avatar
Wenzel Jakob committed
369
370
    }

371
    dtype(const char *format) : dtype(std::string(format)) { }
372

373
374
375
376
377
    dtype(list names, list formats, list offsets, size_t itemsize) {
        dict args;
        args["names"] = names;
        args["formats"] = formats;
        args["offsets"] = offsets;
378
        args["itemsize"] = pybind11::int_(itemsize);
379
380
381
        m_ptr = from_args(args).release().ptr();
    }

Ivan Smirnov's avatar
Ivan Smirnov committed
382
    /// This is essentially the same as calling numpy.dtype(args) in Python.
383
384
385
    static dtype from_args(object args) {
        PyObject *ptr = nullptr;
        if (!detail::npy_api::get().PyArray_DescrConverter_(args.release().ptr(), &ptr) || !ptr)
386
            throw error_already_set();
387
        return reinterpret_steal<dtype>(ptr);
388
    }
389

Ivan Smirnov's avatar
Ivan Smirnov committed
390
    /// Return dtype associated with a C++ type.
391
    template <typename T> static dtype of() {
392
        return detail::npy_format_descriptor<typename std::remove_cv<T>::type>::dtype();
393
    }
394

Ivan Smirnov's avatar
Ivan Smirnov committed
395
    /// Size of the data type in bytes.
396
    size_t itemsize() const {
397
        return (size_t) detail::array_descriptor_proxy(m_ptr)->elsize;
Wenzel Jakob's avatar
Wenzel Jakob committed
398
399
    }

Ivan Smirnov's avatar
Ivan Smirnov committed
400
    /// Returns true for structured data types.
401
    bool has_fields() const {
402
        return detail::array_descriptor_proxy(m_ptr)->names != nullptr;
403
404
    }

Ivan Smirnov's avatar
Ivan Smirnov committed
405
    /// Single-character type code.
406
    char kind() const {
407
        return detail::array_descriptor_proxy(m_ptr)->kind;
408
409
410
    }

private:
411
412
413
    static object _dtype_from_pep3118() {
        static PyObject *obj = module::import("numpy.core._internal")
            .attr("_dtype_from_pep3118").cast<object>().release().ptr();
414
        return reinterpret_borrow<object>(obj);
415
    }
416

417
    dtype strip_padding(size_t itemsize) {
418
419
        // Recursively strip all void fields with empty names that are generated for
        // padding fields (as of NumPy v1.11).
420
        if (!has_fields())
421
            return *this;
422

423
        struct field_descr { PYBIND11_STR_TYPE name; object format; pybind11::int_ offset; };
424
425
        std::vector<field_descr> field_descriptors;

426
        for (auto field : attr("fields").attr("items")()) {
427
            auto spec = field.cast<tuple>();
428
            auto name = spec[0].cast<pybind11::str>();
429
            auto format = spec[1].cast<tuple>()[0].cast<dtype>();
430
            auto offset = spec[1].cast<tuple>()[1].cast<pybind11::int_>();
431
            if (!len(name) && format.kind() == 'V')
432
                continue;
433
            field_descriptors.push_back({(PYBIND11_STR_TYPE) name, format.strip_padding(format.itemsize()), offset});
434
435
436
437
        }

        std::sort(field_descriptors.begin(), field_descriptors.end(),
                  [](const field_descr& a, const field_descr& b) {
438
                      return a.offset.cast<int>() < b.offset.cast<int>();
439
440
441
442
                  });

        list names, formats, offsets;
        for (auto& descr : field_descriptors) {
443
444
445
            names.append(descr.name);
            formats.append(descr.format);
            offsets.append(descr.offset);
446
        }
447
        return dtype(names, formats, offsets, itemsize);
448
449
    }
};
450

451
452
class array : public buffer {
public:
453
    PYBIND11_OBJECT_CVT(array, buffer, detail::npy_api::get().PyArray_Check_, raw_array)
454
455

    enum {
456
457
        c_style = detail::npy_api::NPY_ARRAY_C_CONTIGUOUS_,
        f_style = detail::npy_api::NPY_ARRAY_F_CONTIGUOUS_,
458
459
460
        forcecast = detail::npy_api::NPY_ARRAY_FORCECAST_
    };

461
    array() : array({{0}}, static_cast<const double *>(nullptr)) {}
462

463
464
465
466
467
468
469
470
    using ShapeContainer = detail::any_container<Py_intptr_t>;
    using StridesContainer = detail::any_container<Py_intptr_t>;

    // Constructs an array taking shape/strides from arbitrary container types
    array(const pybind11::dtype &dt, ShapeContainer shape, StridesContainer strides,
          const void *ptr = nullptr, handle base = handle()) {

        if (strides->empty())
471
            *strides = default_strides(*shape, dt.itemsize());
472
473
474

        auto ndim = shape->size();
        if (ndim != strides->size())
475
476
            pybind11_fail("NumPy: shape ndim doesn't match strides ndim");
        auto descr = dt;
477
478
479

        int flags = 0;
        if (base && ptr) {
480
            if (isinstance<array>(base))
Wenzel Jakob's avatar
Wenzel Jakob committed
481
                /* Copy flags from base (except ownership bit) */
482
                flags = reinterpret_borrow<array>(base).flags() & ~detail::npy_api::NPY_ARRAY_OWNDATA_;
483
484
485
486
487
            else
                /* Writable by default, easy to downgrade later on if needed */
                flags = detail::npy_api::NPY_ARRAY_WRITEABLE_;
        }

488
        auto &api = detail::npy_api::get();
489
        auto tmp = reinterpret_steal<object>(api.PyArray_NewFromDescr_(
490
            api.PyArray_Type_, descr.release().ptr(), (int) ndim, shape->data(), strides->data(),
491
            const_cast<void *>(ptr), flags, nullptr));
492
493
        if (!tmp)
            pybind11_fail("NumPy: unable to create array!");
494
495
        if (ptr) {
            if (base) {
Jason Rhinelander's avatar
Jason Rhinelander committed
496
                api.PyArray_SetBaseObject_(tmp.ptr(), base.inc_ref().ptr());
497
            } else {
498
                tmp = reinterpret_steal<object>(api.PyArray_NewCopy_(tmp.ptr(), -1 /* any order */));
499
500
            }
        }
501
502
503
        m_ptr = tmp.release().ptr();
    }

504
505
    array(const pybind11::dtype &dt, ShapeContainer shape, const void *ptr = nullptr, handle base = handle())
        : array(dt, std::move(shape), {}, ptr, base) { }
506

507
508
509
    template <typename T, typename = detail::enable_if_t<std::is_integral<T>::value && !std::is_same<bool, T>::value>>
    array(const pybind11::dtype &dt, T count, const void *ptr = nullptr, handle base = handle())
        : array(dt, {{count}}, ptr, base) { }
510

511
512
513
    template <typename T>
    array(ShapeContainer shape, StridesContainer strides, const T *ptr, handle base = handle())
        : array(pybind11::dtype::of<T>(), std::move(shape), std::move(strides), ptr, base) { }
514

515
    template <typename T>
516
517
    array(ShapeContainer shape, const T *ptr, handle base = handle())
        : array(std::move(shape), {}, ptr, base) { }
518

519
520
521
    template <typename T>
    explicit array(size_t count, const T *ptr, handle base = handle()) : array({count}, {}, ptr, base) { }

522
    explicit array(const buffer_info &info)
523
    : array(pybind11::dtype(info), info.shape, info.strides, info.ptr) { }
524

525
526
    /// Array descriptor (dtype)
    pybind11::dtype dtype() const {
527
        return reinterpret_borrow<pybind11::dtype>(detail::array_proxy(m_ptr)->descr);
528
529
530
531
532
533
534
535
536
    }

    /// Total number of elements
    size_t size() const {
        return std::accumulate(shape(), shape() + ndim(), (size_t) 1, std::multiplies<size_t>());
    }

    /// Byte size of a single element
    size_t itemsize() const {
537
        return (size_t) detail::array_descriptor_proxy(detail::array_proxy(m_ptr)->descr)->elsize;
538
539
540
541
542
543
544
545
546
    }

    /// Total number of bytes
    size_t nbytes() const {
        return size() * itemsize();
    }

    /// Number of dimensions
    size_t ndim() const {
547
        return (size_t) detail::array_proxy(m_ptr)->nd;
548
549
    }

550
551
    /// Base object
    object base() const {
552
        return reinterpret_borrow<object>(detail::array_proxy(m_ptr)->base);
553
554
    }

555
556
    /// Dimensions of the array
    const size_t* shape() const {
557
        return reinterpret_cast<const size_t *>(detail::array_proxy(m_ptr)->dimensions);
558
559
560
561
562
    }

    /// Dimension along a given axis
    size_t shape(size_t dim) const {
        if (dim >= ndim())
563
            fail_dim_check(dim, "invalid axis");
564
565
566
567
568
        return shape()[dim];
    }

    /// Strides of the array
    const size_t* strides() const {
569
        return reinterpret_cast<const size_t *>(detail::array_proxy(m_ptr)->strides);
570
571
572
573
574
    }

    /// Stride along a given axis
    size_t strides(size_t dim) const {
        if (dim >= ndim())
575
            fail_dim_check(dim, "invalid axis");
576
577
578
        return strides()[dim];
    }

579
580
    /// Return the NumPy array flags
    int flags() const {
581
        return detail::array_proxy(m_ptr)->flags;
582
583
    }

584
585
    /// If set, the array is writeable (otherwise the buffer is read-only)
    bool writeable() const {
586
        return detail::check_flags(m_ptr, detail::npy_api::NPY_ARRAY_WRITEABLE_);
587
588
589
590
    }

    /// If set, the array owns the data (will be freed when the array is deleted)
    bool owndata() const {
591
        return detail::check_flags(m_ptr, detail::npy_api::NPY_ARRAY_OWNDATA_);
592
593
    }

594
595
    /// Pointer to the contained data. If index is not provided, points to the
    /// beginning of the buffer. May throw if the index would lead to out of bounds access.
596
    template<typename... Ix> const void* data(Ix... index) const {
597
        return static_cast<const void *>(detail::array_proxy(m_ptr)->data + offset_at(index...));
598
599
    }

600
601
602
    /// Mutable pointer to the contained data. If index is not provided, points to the
    /// beginning of the buffer. May throw if the index would lead to out of bounds access.
    /// May throw if the array is not writeable.
603
    template<typename... Ix> void* mutable_data(Ix... index) {
604
        check_writeable();
605
        return static_cast<void *>(detail::array_proxy(m_ptr)->data + offset_at(index...));
606
607
608
609
    }

    /// Byte offset from beginning of the array to a given index (full or partial).
    /// May throw if the index would lead to out of bounds access.
610
    template<typename... Ix> size_t offset_at(Ix... index) const {
611
612
        if (sizeof...(index) > ndim())
            fail_dim_check(sizeof...(index), "too many indices for an array");
613
        return byte_offset(size_t(index)...);
614
615
616
617
618
619
    }

    size_t offset_at() const { return 0; }

    /// Item count from beginning of the array to a given index (full or partial).
    /// May throw if the index would lead to out of bounds access.
620
    template<typename... Ix> size_t index_at(Ix... index) const {
621
        return offset_at(index...) / itemsize();
622
623
    }

624
625
626
627
628
    /** Returns a proxy object that provides access to the array's data without bounds or
     * dimensionality checking.  Will throw if the array is missing the `writeable` flag.  Use with
     * care: the array must not be destroyed or reshaped for the duration of the returned object,
     * and the caller must take care not to access invalid dimensions or dimension indices.
     */
629
630
    template <typename T, ssize_t Dims = -1> detail::unchecked_mutable_reference<T, Dims> mutable_unchecked() {
        if (Dims >= 0 && ndim() != (size_t) Dims)
631
632
            throw std::domain_error("array has incorrect number of dimensions: " + std::to_string(ndim()) +
                    "; expected " + std::to_string(Dims));
633
        return detail::unchecked_mutable_reference<T, Dims>(mutable_data(), shape(), strides(), ndim());
634
635
636
637
638
639
640
641
    }

    /** Returns a proxy object that provides const access to the array's data without bounds or
     * dimensionality checking.  Unlike `mutable_unchecked()`, this does not require that the
     * underlying array have the `writable` flag.  Use with care: the array must not be destroyed or
     * reshaped for the duration of the returned object, and the caller must take care not to access
     * invalid dimensions or dimension indices.
     */
642
643
    template <typename T, ssize_t Dims = -1> detail::unchecked_reference<T, Dims> unchecked() const {
        if (Dims >= 0 && ndim() != (size_t) Dims)
644
645
            throw std::domain_error("array has incorrect number of dimensions: " + std::to_string(ndim()) +
                    "; expected " + std::to_string(Dims));
646
        return detail::unchecked_reference<T, Dims>(data(), shape(), strides(), ndim());
647
648
    }

649
650
651
    /// Return a new view with all of the dimensions of length 1 removed
    array squeeze() {
        auto& api = detail::npy_api::get();
652
        return reinterpret_steal<array>(api.PyArray_Squeeze_(m_ptr));
653
654
    }

655
    /// Ensure that the argument is a NumPy array
656
657
658
659
660
661
    /// In case of an error, nullptr is returned and the Python error is cleared.
    static array ensure(handle h, int ExtraFlags = 0) {
        auto result = reinterpret_steal<array>(raw_array(h.ptr(), ExtraFlags));
        if (!result)
            PyErr_Clear();
        return result;
662
663
    }

664
protected:
665
666
667
668
669
670
671
    template<typename, typename> friend struct detail::npy_format_descriptor;

    void fail_dim_check(size_t dim, const std::string& msg) const {
        throw index_error(msg + ": " + std::to_string(dim) +
                          " (ndim = " + std::to_string(ndim()) + ")");
    }

672
673
    template<typename... Ix> size_t byte_offset(Ix... index) const {
        check_dimensions(index...);
674
        return detail::byte_offset_unsafe(strides(), size_t(index)...);
675
676
    }

677
678
    void check_writeable() const {
        if (!writeable())
679
            throw std::domain_error("array is not writeable");
680
    }
681

682
    static std::vector<Py_intptr_t> default_strides(const std::vector<Py_intptr_t>& shape, size_t itemsize) {
683
        auto ndim = shape.size();
684
        std::vector<Py_intptr_t> strides(ndim);
685
686
687
688
689
690
691
692
        if (ndim) {
            std::fill(strides.begin(), strides.end(), itemsize);
            for (size_t i = 0; i < ndim - 1; i++)
                for (size_t j = 0; j < ndim - 1 - i; j++)
                    strides[j] *= shape[ndim - 1 - i];
        }
        return strides;
    }
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707

    template<typename... Ix> void check_dimensions(Ix... index) const {
        check_dimensions_impl(size_t(0), shape(), size_t(index)...);
    }

    void check_dimensions_impl(size_t, const size_t*) const { }

    template<typename... Ix> void check_dimensions_impl(size_t axis, const size_t* shape, size_t i, Ix... index) const {
        if (i >= *shape) {
            throw index_error(std::string("index ") + std::to_string(i) +
                              " is out of bounds for axis " + std::to_string(axis) +
                              " with size " + std::to_string(*shape));
        }
        check_dimensions_impl(axis + 1, shape + 1, index...);
    }
708
709
710

    /// Create array from any object -- always returns a new reference
    static PyObject *raw_array(PyObject *ptr, int ExtraFlags = 0) {
711
712
        if (ptr == nullptr) {
            PyErr_SetString(PyExc_ValueError, "cannot create a pybind11::array from a nullptr");
713
            return nullptr;
714
        }
715
        return detail::npy_api::get().PyArray_FromAny_(
716
            ptr, nullptr, 0, 0, detail::npy_api::NPY_ARRAY_ENSUREARRAY_ | ExtraFlags, nullptr);
717
    }
Wenzel Jakob's avatar
Wenzel Jakob committed
718
719
};

720
template <typename T, int ExtraFlags = array::forcecast> class array_t : public array {
Wenzel Jakob's avatar
Wenzel Jakob committed
721
public:
722
723
    using value_type = T;

724
    array_t() : array(0, static_cast<const T *>(nullptr)) {}
725
726
    array_t(handle h, borrowed_t) : array(h, borrowed_t{}) { }
    array_t(handle h, stolen_t) : array(h, stolen_t{}) { }
727

728
    PYBIND11_DEPRECATED("Use array_t<T>::ensure() instead")
729
    array_t(handle h, bool is_borrowed) : array(raw_array_t(h.ptr()), stolen_t{}) {
730
731
732
        if (!m_ptr) PyErr_Clear();
        if (!is_borrowed) Py_XDECREF(h.ptr());
    }
733

734
    array_t(const object &o) : array(raw_array_t(o.ptr()), stolen_t{}) {
735
736
        if (!m_ptr) throw error_already_set();
    }
737

738
    explicit array_t(const buffer_info& info) : array(info) { }
739

740
741
    array_t(ShapeContainer shape, StridesContainer strides, const T *ptr = nullptr, handle base = handle())
        : array(std::move(shape), std::move(strides), ptr, base) { }
742

743
744
    explicit array_t(ShapeContainer shape, const T *ptr = nullptr, handle base = handle())
        : array(std::move(shape), ptr, base) { }
745

746
747
748
    explicit array_t(size_t count, const T *ptr = nullptr, handle base = handle())
        : array({count}, {}, ptr, base) { }

749
750
    constexpr size_t itemsize() const {
        return sizeof(T);
751
752
    }

753
    template<typename... Ix> size_t index_at(Ix... index) const {
754
755
756
        return offset_at(index...) / itemsize();
    }

757
    template<typename... Ix> const T* data(Ix... index) const {
758
759
760
        return static_cast<const T*>(array::data(index...));
    }

761
    template<typename... Ix> T* mutable_data(Ix... index) {
762
763
764
765
        return static_cast<T*>(array::mutable_data(index...));
    }

    // Reference to element at a given index
766
    template<typename... Ix> const T& at(Ix... index) const {
767
768
        if (sizeof...(index) != ndim())
            fail_dim_check(sizeof...(index), "index dimension mismatch");
769
        return *(static_cast<const T*>(array::data()) + byte_offset(size_t(index)...) / itemsize());
770
771
772
    }

    // Mutable reference to element at a given index
773
    template<typename... Ix> T& mutable_at(Ix... index) {
774
775
        if (sizeof...(index) != ndim())
            fail_dim_check(sizeof...(index), "index dimension mismatch");
776
        return *(static_cast<T*>(array::mutable_data()) + byte_offset(size_t(index)...) / itemsize());
777
    }
778

779
780
781
782
783
    /** Returns a proxy object that provides access to the array's data without bounds or
     * dimensionality checking.  Will throw if the array is missing the `writeable` flag.  Use with
     * care: the array must not be destroyed or reshaped for the duration of the returned object,
     * and the caller must take care not to access invalid dimensions or dimension indices.
     */
784
    template <ssize_t Dims = -1> detail::unchecked_mutable_reference<T, Dims> mutable_unchecked() {
785
786
787
788
789
790
791
792
793
        return array::mutable_unchecked<T, Dims>();
    }

    /** Returns a proxy object that provides const access to the array's data without bounds or
     * dimensionality checking.  Unlike `unchecked()`, this does not require that the underlying
     * array have the `writable` flag.  Use with care: the array must not be destroyed or reshaped
     * for the duration of the returned object, and the caller must take care not to access invalid
     * dimensions or dimension indices.
     */
794
    template <ssize_t Dims = -1> detail::unchecked_reference<T, Dims> unchecked() const {
795
796
797
        return array::unchecked<T, Dims>();
    }

Jason Rhinelander's avatar
Jason Rhinelander committed
798
799
    /// Ensure that the argument is a NumPy array of the correct dtype (and if not, try to convert
    /// it).  In case of an error, nullptr is returned and the Python error is cleared.
800
801
    static array_t ensure(handle h) {
        auto result = reinterpret_steal<array_t>(raw_array_t(h.ptr()));
802
803
        if (!result)
            PyErr_Clear();
804
        return result;
Wenzel Jakob's avatar
Wenzel Jakob committed
805
    }
806

Wenzel Jakob's avatar
Wenzel Jakob committed
807
    static bool check_(handle h) {
808
809
        const auto &api = detail::npy_api::get();
        return api.PyArray_Check_(h.ptr())
810
               && api.PyArray_EquivTypes_(detail::array_proxy(h.ptr())->descr, dtype::of<T>().ptr());
811
812
813
814
815
    }

protected:
    /// Create array from any object -- always returns a new reference
    static PyObject *raw_array_t(PyObject *ptr) {
816
817
        if (ptr == nullptr) {
            PyErr_SetString(PyExc_ValueError, "cannot create a pybind11::array_t from a nullptr");
818
            return nullptr;
819
        }
820
821
        return detail::npy_api::get().PyArray_FromAny_(
            ptr, dtype::of<T>().release().ptr(), 0, 0,
822
            detail::npy_api::NPY_ARRAY_ENSUREARRAY_ | ExtraFlags, nullptr);
823
    }
Wenzel Jakob's avatar
Wenzel Jakob committed
824
825
};

826
template <typename T>
827
struct format_descriptor<T, detail::enable_if_t<detail::is_pod_struct<T>::value>> {
828
829
830
    static std::string format() {
        return detail::npy_format_descriptor<typename std::remove_cv<T>::type>::format();
    }
831
832
833
};

template <size_t N> struct format_descriptor<char[N]> {
834
    static std::string format() { return std::to_string(N) + "s"; }
835
836
};
template <size_t N> struct format_descriptor<std::array<char, N>> {
837
    static std::string format() { return std::to_string(N) + "s"; }
838
839
};

840
841
842
843
844
845
846
847
template <typename T>
struct format_descriptor<T, detail::enable_if_t<std::is_enum<T>::value>> {
    static std::string format() {
        return format_descriptor<
            typename std::remove_cv<typename std::underlying_type<T>::type>::type>::format();
    }
};

848
NAMESPACE_BEGIN(detail)
849
850
851
852
template <typename T, int ExtraFlags>
struct pyobject_caster<array_t<T, ExtraFlags>> {
    using type = array_t<T, ExtraFlags>;

853
854
855
    bool load(handle src, bool convert) {
        if (!convert && !type::check_(src))
            return false;
856
        value = type::ensure(src);
857
858
859
860
861
862
863
864
865
        return static_cast<bool>(value);
    }

    static handle cast(const handle &src, return_value_policy /* policy */, handle /* parent */) {
        return src.inc_ref();
    }
    PYBIND11_TYPE_CASTER(type, handle_type_name<type>::name());
};

866
867
868
869
870
871
872
template <typename T>
struct compare_buffer_info<T, detail::enable_if_t<detail::is_pod_struct<T>::value>> {
    static bool compare(const buffer_info& b) {
        return npy_api::get().PyArray_EquivTypes_(dtype::of<T>().ptr(), dtype(b).ptr());
    }
};

873
template <typename T> struct npy_format_descriptor<T, enable_if_t<satisfies_any_of<T, std::is_arithmetic, is_complex>::value>> {
874
private:
875
876
877
878
879
880
881
882
883
    // NB: the order here must match the one in common.h
    constexpr static const int values[15] = {
        npy_api::NPY_BOOL_,
        npy_api::NPY_BYTE_,   npy_api::NPY_UBYTE_,   npy_api::NPY_SHORT_,    npy_api::NPY_USHORT_,
        npy_api::NPY_INT_,    npy_api::NPY_UINT_,    npy_api::NPY_LONGLONG_, npy_api::NPY_ULONGLONG_,
        npy_api::NPY_FLOAT_,  npy_api::NPY_DOUBLE_,  npy_api::NPY_LONGDOUBLE_,
        npy_api::NPY_CFLOAT_, npy_api::NPY_CDOUBLE_, npy_api::NPY_CLONGDOUBLE_
    };

884
public:
885
886
    static constexpr int value = values[detail::is_fmt_numeric<T>::index];

887
    static pybind11::dtype dtype() {
888
        if (auto ptr = npy_api::get().PyArray_DescrFromType_(value))
889
            return reinterpret_borrow<pybind11::dtype>(ptr);
890
        pybind11_fail("Unsupported buffer format!");
891
    }
892
893
894
895
896
897
898
899
900
901
902
903
904
    template <typename T2 = T, enable_if_t<std::is_integral<T2>::value, int> = 0>
    static PYBIND11_DESCR name() {
        return _<std::is_same<T, bool>::value>(_("bool"),
            _<std::is_signed<T>::value>("int", "uint") + _<sizeof(T)*8>());
    }
    template <typename T2 = T, enable_if_t<std::is_floating_point<T2>::value, int> = 0>
    static PYBIND11_DESCR name() {
        return _<std::is_same<T, float>::value || std::is_same<T, double>::value>(
                _("float") + _<sizeof(T)*8>(), _("longdouble"));
    }
    template <typename T2 = T, enable_if_t<is_complex<T2>::value, int> = 0>
    static PYBIND11_DESCR name() {
        return _<std::is_same<typename T2::value_type, float>::value || std::is_same<typename T2::value_type, double>::value>(
905
                _("complex") + _<sizeof(typename T2::value_type)*16>(), _("longcomplex"));
906
    }
907
};
908
909

#define PYBIND11_DECL_CHAR_FMT \
910
    static PYBIND11_DESCR name() { return _("S") + _<N>(); } \
911
    static pybind11::dtype dtype() { return pybind11::dtype(std::string("S") + std::to_string(N)); }
912
913
914
template <size_t N> struct npy_format_descriptor<char[N]> { PYBIND11_DECL_CHAR_FMT };
template <size_t N> struct npy_format_descriptor<std::array<char, N>> { PYBIND11_DECL_CHAR_FMT };
#undef PYBIND11_DECL_CHAR_FMT
915

916
917
918
919
920
921
922
923
template<typename T> struct npy_format_descriptor<T, enable_if_t<std::is_enum<T>::value>> {
private:
    using base_descr = npy_format_descriptor<typename std::underlying_type<T>::type>;
public:
    static PYBIND11_DESCR name() { return base_descr::name(); }
    static pybind11::dtype dtype() { return base_descr::dtype(); }
};

924
925
struct field_descriptor {
    const char *name;
926
    size_t offset;
927
    size_t size;
928
    size_t alignment;
929
    std::string format;
930
    dtype descr;
931
932
};

933
934
935
inline PYBIND11_NOINLINE void register_structured_dtype(
    const std::initializer_list<field_descriptor>& fields,
    const std::type_info& tinfo, size_t itemsize,
936
937
    bool (*direct_converter)(PyObject *, void *&)) {

938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
    auto& numpy_internals = get_numpy_internals();
    if (numpy_internals.get_type_info(tinfo, false))
        pybind11_fail("NumPy: dtype is already registered");

    list names, formats, offsets;
    for (auto field : fields) {
        if (!field.descr)
            pybind11_fail(std::string("NumPy: unsupported field dtype: `") +
                            field.name + "` @ " + tinfo.name());
        names.append(PYBIND11_STR_TYPE(field.name));
        formats.append(field.descr);
        offsets.append(pybind11::int_(field.offset));
    }
    auto dtype_ptr = pybind11::dtype(names, formats, offsets, itemsize).release().ptr();

    // 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';
969
        // mark unaligned fields with '^' (unaligned native type)
970
        if (field.offset % field.alignment)
971
            oss << '^';
972
        oss << field.format << ':' << field.name << ':';
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
        offset = field.offset + field.size;
    }
    if (itemsize > offset)
        oss << (itemsize - offset) << 'x';
    oss << '}';
    auto format_str = oss.str();

    // Sanity check: verify that NumPy properly parses our buffer format string
    auto& api = npy_api::get();
    auto arr =  array(buffer_info(nullptr, itemsize, format_str, 1));
    if (!api.PyArray_EquivTypes_(dtype_ptr, arr.dtype().ptr()))
        pybind11_fail("NumPy: invalid buffer descriptor!");

    auto tindex = std::type_index(tinfo);
    numpy_internals.registered_dtypes[tindex] = { dtype_ptr, format_str };
    get_internals().direct_conversions[tindex].push_back(direct_converter);
}

991
992
993
template <typename T, typename SFINAE> struct npy_format_descriptor {
    static_assert(is_pod_struct<T>::value, "Attempt to use a non-POD or unimplemented POD type as a numpy dtype");

994
    static PYBIND11_DESCR name() { return make_caster<T>::name(); }
995

996
    static pybind11::dtype dtype() {
997
        return reinterpret_borrow<pybind11::dtype>(dtype_ptr());
998
999
    }

1000
    static std::string format() {
1001
        static auto format_str = get_numpy_internals().get_type_info<T>(true)->format_str;
1002
        return format_str;
1003
1004
    }

1005
1006
1007
    static void register_dtype(const std::initializer_list<field_descriptor>& fields) {
        register_structured_dtype(fields, typeid(typename std::remove_cv<T>::type),
                                  sizeof(T), &direct_converter);
1008
1009
1010
    }

private:
1011
1012
1013
1014
    static PyObject* dtype_ptr() {
        static PyObject* ptr = get_numpy_internals().get_type_info<T>(true)->dtype_ptr;
        return ptr;
    }
1015

1016
1017
1018
    static bool direct_converter(PyObject *obj, void*& value) {
        auto& api = npy_api::get();
        if (!PyObject_TypeCheck(obj, api.PyVoidArrType_Type_))
1019
            return false;
1020
        if (auto descr = reinterpret_steal<object>(api.PyArray_DescrFromScalar_(obj))) {
1021
            if (api.PyArray_EquivTypes_(dtype_ptr(), descr.ptr())) {
1022
1023
1024
1025
1026
1027
                value = ((PyVoidScalarObject_Proxy *) obj)->obval;
                return true;
            }
        }
        return false;
    }
1028
1029
};

1030
1031
1032
#define PYBIND11_FIELD_DESCRIPTOR_EX(T, Field, Name)                                          \
    ::pybind11::detail::field_descriptor {                                                    \
        Name, offsetof(T, Field), sizeof(decltype(std::declval<T>().Field)),                  \
1033
        alignof(decltype(std::declval<T>().Field)),                                           \
1034
1035
        ::pybind11::format_descriptor<decltype(std::declval<T>().Field)>::format(),           \
        ::pybind11::detail::npy_format_descriptor<decltype(std::declval<T>().Field)>::dtype() \
1036
    }
1037

1038
1039
1040
// Extract name, offset and format descriptor for a struct field
#define PYBIND11_FIELD_DESCRIPTOR(T, Field) PYBIND11_FIELD_DESCRIPTOR_EX(T, Field, #Field)

1041
1042
// The main idea of this macro is borrowed from https://github.com/swansontec/map-macro
// (C) William Swanson, Paul Fultz
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
#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)
1056
#ifdef _MSC_VER // MSVC is not as eager to expand macros, hence this workaround
1057
1058
#define PYBIND11_MAP_LIST_NEXT1(test, next) \
    PYBIND11_EVAL0 (PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0))
1059
#else
1060
1061
#define PYBIND11_MAP_LIST_NEXT1(test, next) \
    PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0)
1062
#endif
1063
1064
1065
1066
1067
1068
#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__)
1069
// PYBIND11_MAP_LIST(f, t, a1, a2, ...) expands to f(t, a1), f(t, a2), ...
1070
1071
#define PYBIND11_MAP_LIST(f, t, ...) \
    PYBIND11_EVAL (PYBIND11_MAP_LIST1 (f, t, __VA_ARGS__, (), 0))
1072

1073
#define PYBIND11_NUMPY_DTYPE(Type, ...) \
1074
    ::pybind11::detail::npy_format_descriptor<Type>::register_dtype \
1075
        ({PYBIND11_MAP_LIST (PYBIND11_FIELD_DESCRIPTOR, Type, __VA_ARGS__)})
1076

1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
#ifdef _MSC_VER
#define PYBIND11_MAP2_LIST_NEXT1(test, next) \
    PYBIND11_EVAL0 (PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0))
#else
#define PYBIND11_MAP2_LIST_NEXT1(test, next) \
    PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0)
#endif
#define PYBIND11_MAP2_LIST_NEXT(test, next) \
    PYBIND11_MAP2_LIST_NEXT1 (PYBIND11_MAP_GET_END test, next)
#define PYBIND11_MAP2_LIST0(f, t, x1, x2, peek, ...) \
    f(t, x1, x2) PYBIND11_MAP2_LIST_NEXT (peek, PYBIND11_MAP2_LIST1) (f, t, peek, __VA_ARGS__)
#define PYBIND11_MAP2_LIST1(f, t, x1, x2, peek, ...) \
    f(t, x1, x2) PYBIND11_MAP2_LIST_NEXT (peek, PYBIND11_MAP2_LIST0) (f, t, peek, __VA_ARGS__)
// PYBIND11_MAP2_LIST(f, t, a1, a2, ...) expands to f(t, a1, a2), f(t, a3, a4), ...
#define PYBIND11_MAP2_LIST(f, t, ...) \
    PYBIND11_EVAL (PYBIND11_MAP2_LIST1 (f, t, __VA_ARGS__, (), 0))

#define PYBIND11_NUMPY_DTYPE_EX(Type, ...) \
    ::pybind11::detail::npy_format_descriptor<Type>::register_dtype \
        ({PYBIND11_MAP2_LIST (PYBIND11_FIELD_DESCRIPTOR_EX, Type, __VA_ARGS__)})

1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
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() {}
1118

1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
    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;
};

1142
template <size_t N> class multi_array_iterator {
1143
1144
1145
public:
    using container_type = std::vector<size_t>;

1146
1147
1148
1149
1150
    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() {

1151
        // Manual copy to avoid conversion warning if using std::copy
1152
        for (size_t i = 0; i < shape.size(); ++i)
1153
1154
1155
            m_shape[i] = static_cast<container_type::value_type>(shape[i]);

        container_type strides(shape.size());
1156
        for (size_t i = 0; i < N; ++i)
1157
1158
1159
1160
1161
1162
1163
1164
1165
            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;
1166
            } else {
1167
1168
1169
1170
1171
1172
                m_index[i] = 0;
            }
        }
        return *this;
    }

1173
    template <size_t K, class T> const T& data() const {
1174
1175
1176
1177
1178
1179
1180
        return *reinterpret_cast<T*>(m_common_iterator[K].data());
    }

private:

    using common_iter = common_iterator;

1181
1182
1183
    void init_common_iterator(const buffer_info &buffer,
                              const std::vector<size_t> &shape,
                              common_iter &iterator, container_type &strides) {
1184
1185
1186
1187
1188
1189
1190
        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)
1191
                *strides_iter = static_cast<size_t>(*buffer_strides_iter);
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
            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) {
1206
        for (auto &iter : m_common_iterator)
1207
1208
1209
1210
1211
1212
1213
1214
            iter.increment(dim);
    }

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

1215
1216
1217
1218
1219
1220
enum class broadcast_trivial { non_trivial, c_trivial, f_trivial };

// Populates the shape and number of dimensions for the set of buffers.  Returns a broadcast_trivial
// enum value indicating whether the broadcast is "trivial"--that is, has each buffer being either a
// singleton or a full-size, C-contiguous (`c_trivial`) or Fortran-contiguous (`f_trivial`) storage
// buffer; returns `non_trivial` otherwise.
1221
template <size_t N>
1222
broadcast_trivial broadcast(const std::array<buffer_info, N> &buffers, size_t &ndim, std::vector<size_t> &shape) {
1223
    ndim = std::accumulate(buffers.begin(), buffers.end(), size_t(0), [](size_t res, const buffer_info& buf) {
1224
1225
1226
        return std::max(res, buf.ndim);
    });

1227
1228
1229
    shape.clear();
    shape.resize(ndim, 1);

1230
1231
    // Figure out the output size, and make sure all input arrays conform (i.e. are either size 1 or
    // the full size).
1232
1233
    for (size_t i = 0; i < N; ++i) {
        auto res_iter = shape.rbegin();
1234
1235
        auto end = buffers[i].shape.rend();
        for (auto shape_iter = buffers[i].shape.rbegin(); shape_iter != end; ++shape_iter, ++res_iter) {
1236
1237
1238
1239
1240
1241
1242
            const auto &dim_size_in = *shape_iter;
            auto &dim_size_out = *res_iter;

            // Each input dimension can either be 1 or `n`, but `n` values must match across buffers
            if (dim_size_out == 1)
                dim_size_out = dim_size_in;
            else if (dim_size_in != 1 && dim_size_in != dim_size_out)
1243
                pybind11_fail("pybind11::vectorize: incompatible size/dimension of inputs!");
1244
1245
        }
    }
1246

1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
    bool trivial_broadcast_c = true;
    bool trivial_broadcast_f = true;
    for (size_t i = 0; i < N && (trivial_broadcast_c || trivial_broadcast_f); ++i) {
        if (buffers[i].size == 1)
            continue;

        // Require the same number of dimensions:
        if (buffers[i].ndim != ndim)
            return broadcast_trivial::non_trivial;

        // Require all dimensions be full-size:
        if (!std::equal(buffers[i].shape.cbegin(), buffers[i].shape.cend(), shape.cbegin()))
            return broadcast_trivial::non_trivial;

        // Check for C contiguity (but only if previous inputs were also C contiguous)
        if (trivial_broadcast_c) {
            size_t expect_stride = buffers[i].itemsize;
            auto end = buffers[i].shape.crend();
            for (auto shape_iter = buffers[i].shape.crbegin(), stride_iter = buffers[i].strides.crbegin();
                    trivial_broadcast_c && shape_iter != end; ++shape_iter, ++stride_iter) {
                if (expect_stride == *stride_iter)
                    expect_stride *= *shape_iter;
                else
                    trivial_broadcast_c = false;
1271
            }
1272
        }
1273

1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
        // Check for Fortran contiguity (if previous inputs were also F contiguous)
        if (trivial_broadcast_f) {
            size_t expect_stride = buffers[i].itemsize;
            auto end = buffers[i].shape.cend();
            for (auto shape_iter = buffers[i].shape.cbegin(), stride_iter = buffers[i].strides.cbegin();
                    trivial_broadcast_f && shape_iter != end; ++shape_iter, ++stride_iter) {
                if (expect_stride == *stride_iter)
                    expect_stride *= *shape_iter;
                else
                    trivial_broadcast_f = false;
            }
1285
1286
        }
    }
1287
1288
1289
1290
1291

    return
        trivial_broadcast_c ? broadcast_trivial::c_trivial :
        trivial_broadcast_f ? broadcast_trivial::f_trivial :
        broadcast_trivial::non_trivial;
1292
1293
}

1294
1295
1296
template <typename Func, typename Return, typename... Args>
struct vectorize_helper {
    typename std::remove_reference<Func>::type f;
1297
    static constexpr size_t N = sizeof...(Args);
1298

1299
    template <typename T>
1300
    explicit vectorize_helper(T&&f) : f(std::forward<T>(f)) { }
Wenzel Jakob's avatar
Wenzel Jakob committed
1301

1302
1303
    object operator()(array_t<Args, array::forcecast>... args) {
        return run(args..., make_index_sequence<N>());
1304
    }
Wenzel Jakob's avatar
Wenzel Jakob committed
1305

1306
    template <size_t ... Index> object run(array_t<Args, array::forcecast>&... args, index_sequence<Index...> index) {
Wenzel Jakob's avatar
Wenzel Jakob committed
1307
1308
1309
1310
        /* Request buffers from all parameters */
        std::array<buffer_info, N> buffers {{ args.request()... }};

        /* Determine dimensions parameters of output array */
1311
        size_t ndim = 0;
1312
        std::vector<size_t> shape(0);
1313
        auto trivial = broadcast(buffers, ndim, shape);
1314

1315
        size_t size = 1;
Wenzel Jakob's avatar
Wenzel Jakob committed
1316
1317
        std::vector<size_t> strides(ndim);
        if (ndim > 0) {
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
            if (trivial == broadcast_trivial::f_trivial) {
                strides[0] = sizeof(Return);
                for (size_t i = 1; i < ndim; ++i) {
                    strides[i] = strides[i - 1] * shape[i - 1];
                    size *= shape[i - 1];
                }
                size *= shape[ndim - 1];
            }
            else {
                strides[ndim-1] = sizeof(Return);
                for (size_t i = ndim - 1; i > 0; --i) {
                    strides[i - 1] = strides[i] * shape[i];
                    size *= shape[i];
                }
                size *= shape[0];
1333
            }
Wenzel Jakob's avatar
Wenzel Jakob committed
1334
1335
        }

1336
        if (size == 1)
1337
            return cast(f(*reinterpret_cast<Args *>(buffers[Index].ptr)...));
Wenzel Jakob's avatar
Wenzel Jakob committed
1338

1339
        array_t<Return> result(shape, strides);
1340
1341
        auto buf = result.request();
        auto output = (Return *) buf.ptr;
1342

1343
1344
1345
1346
        /* Call the function */
        if (trivial == broadcast_trivial::non_trivial) {
            apply_broadcast<Index...>(buffers, buf, index);
        } else {
1347
1348
            for (size_t i = 0; i < size; ++i)
                output[i] = f((reinterpret_cast<Args *>(buffers[Index].ptr)[buffers[Index].size == 1 ? 0 : i])...);
1349
        }
1350
1351

        return result;
1352
    }
1353

1354
    template <size_t... Index>
1355
1356
    void apply_broadcast(const std::array<buffer_info, N> &buffers,
                         buffer_info &output, index_sequence<Index...>) {
1357
1358
1359
1360
1361
1362
        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);

1363
1364
        for (output_iterator iter = array_begin<Return>(output);
             iter != output_end; ++iter, ++input_iter) {
1365
1366
1367
            *iter = f((input_iter.template data<Index, Args>())...);
        }
    }
1368
1369
};

1370
template <typename T, int Flags> struct handle_type_name<array_t<T, Flags>> {
1371
1372
1373
    static PYBIND11_DESCR name() {
        return _("numpy.ndarray[") + npy_format_descriptor<T>::name() + _("]");
    }
1374
1375
};

1376
NAMESPACE_END(detail)
Wenzel Jakob's avatar
Wenzel Jakob committed
1377

1378
template <typename Func, typename Return, typename... Args>
1379
detail::vectorize_helper<Func, Return, Args...>
1380
vectorize(const Func &f, Return (*) (Args ...)) {
1381
    return detail::vectorize_helper<Func, Return, Args...>(f);
Wenzel Jakob's avatar
Wenzel Jakob committed
1382
1383
}

1384
1385
1386
template <typename Return, typename... Args>
detail::vectorize_helper<Return (*) (Args ...), Return, Args...>
vectorize(Return (*f) (Args ...)) {
1387
    return vectorize<Return (*) (Args ...), Return, Args...>(f, f);
Wenzel Jakob's avatar
Wenzel Jakob committed
1388
1389
}

1390
template <typename Func, typename FuncType = typename detail::remove_class<decltype(&std::remove_reference<Func>::type::operator())>::type>
1391
auto vectorize(Func &&f) -> decltype(
1392
1393
        vectorize(std::forward<Func>(f), (FuncType *) nullptr)) {
    return vectorize(std::forward<Func>(f), (FuncType *) nullptr);
Wenzel Jakob's avatar
Wenzel Jakob committed
1394
1395
}

1396
NAMESPACE_END(pybind11)
Wenzel Jakob's avatar
Wenzel Jakob committed
1397
1398
1399
1400

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