|
| 1 | +/* |
| 2 | + Backport from pybind11/numpy.h v2.0 prerelease. |
| 3 | +
|
| 4 | + Copyright (c) 2016 Wenzel Jakob <[email protected]> |
| 5 | +
|
| 6 | + All rights reserved. Use of this source code is governed by a |
| 7 | + BSD-style license that can be found in the LICENSE file. |
| 8 | +*/ |
| 9 | +#pragma once |
| 10 | + |
| 11 | +#include <cstddef> |
| 12 | +#include <algorithm> |
| 13 | + |
| 14 | +#include "pybind11/numpy.h" |
| 15 | + |
| 16 | +namespace pybind11 { namespace backport { |
| 17 | + struct PyArray_Proxy { |
| 18 | + PyObject_HEAD |
| 19 | + char *data; |
| 20 | + int nd; |
| 21 | + ssize_t *dimensions; |
| 22 | + ssize_t *strides; |
| 23 | + PyObject *base; |
| 24 | + PyObject *descr; |
| 25 | + int flags; |
| 26 | + }; |
| 27 | + |
| 28 | + struct PyArrayDescr_Proxy { |
| 29 | + PyObject_HEAD |
| 30 | + PyObject *typeobj; |
| 31 | + char kind; |
| 32 | + char type; |
| 33 | + char byteorder; |
| 34 | + char flags; |
| 35 | + int type_num; |
| 36 | + int elsize; |
| 37 | + int alignment; |
| 38 | + char *subarray; |
| 39 | + PyObject *fields; |
| 40 | + PyObject *names; |
| 41 | + }; |
| 42 | + |
| 43 | +#ifndef PyArray_GET_ |
| 44 | +# define PyArray_GET_(ptr, attr) \ |
| 45 | + (reinterpret_cast<::pybind11::backport::PyArray_Proxy*>(ptr)->attr) |
| 46 | +#endif |
| 47 | +#ifndef PyArrayDescr_GET_ |
| 48 | +# define PyArrayDescr_GET_(ptr, attr) \ |
| 49 | + (reinterpret_cast<::pybind11::backport::PyArrayDescr_Proxy*>(ptr)->attr) |
| 50 | +#endif |
| 51 | + |
| 52 | + class array : public pybind11::array { |
| 53 | + public: |
| 54 | + using size_type = std::size_t; |
| 55 | + |
| 56 | + using pybind11::array::array; |
| 57 | + |
| 58 | + array() = default; |
| 59 | + |
| 60 | + template<typename T> array(const std::vector<size_type>& shape, |
| 61 | + const std::vector<size_type>& strides, |
| 62 | + const T* ptr, handle base = {}) |
| 63 | + : array(buffer_info(const_cast<T*>(ptr), sizeof(T), format_descriptor<T>::value, |
| 64 | + shape.size(), shape, strides)) { |
| 65 | + if (base) throw std::runtime_error("array base is not supported yet"); |
| 66 | + } |
| 67 | + |
| 68 | + template<typename T> array(const std::vector<size_type> &shape, |
| 69 | + const T *ptr, handle base = {}) |
| 70 | + : array(shape, default_strides(shape, sizeof(T)), ptr, base) { } |
| 71 | + |
| 72 | + template<typename T> array(size_type size, const T *ptr, handle base) |
| 73 | + : array(std::vector<size_type>{size}, ptr) { } |
| 74 | + |
| 75 | + size_type size() const { |
| 76 | + return std::accumulate(shape(), shape() + ndim(), size_type{1}, std::multiplies<size_type>()); |
| 77 | + } |
| 78 | + |
| 79 | + size_type itemsize() const { |
| 80 | + return static_cast<size_type>(PyArrayDescr_GET_(PyArray_GET_(m_ptr, descr), elsize)); |
| 81 | + } |
| 82 | + |
| 83 | + size_type ndim() const { |
| 84 | + return static_cast<size_type>(PyArray_GET_(m_ptr, nd)); |
| 85 | + } |
| 86 | + |
| 87 | + const size_type* shape() const { |
| 88 | + return reinterpret_cast<const size_type *>(PyArray_GET_(m_ptr, dimensions)); |
| 89 | + } |
| 90 | + |
| 91 | + const size_type* strides() const { |
| 92 | + return reinterpret_cast<const size_type *>(PyArray_GET_(m_ptr, strides)); |
| 93 | + } |
| 94 | + |
| 95 | + template<typename... Ix> void* data() { |
| 96 | + return static_cast<void *>(PyArray_GET_(m_ptr, data)); |
| 97 | + } |
| 98 | + |
| 99 | + template<typename... Ix> void* mutable_data() { |
| 100 | + // check_writeable(); |
| 101 | + return static_cast<void *>(PyArray_GET_(m_ptr, data)); |
| 102 | + } |
| 103 | + |
| 104 | + template<typename... Ix> size_type offset_at(Ix... index) const { |
| 105 | + if (sizeof...(index) > ndim()) |
| 106 | + fail_dim_check(sizeof...(index), "too many indices for an array"); |
| 107 | + return get_byte_offset(index...); |
| 108 | + } |
| 109 | + |
| 110 | + size_type offset_at() const { return 0; } |
| 111 | + |
| 112 | + protected: |
| 113 | + void fail_dim_check(size_type dim, const std::string& msg) const { |
| 114 | + throw index_error(msg + ": " + std::to_string(dim) + |
| 115 | + " (ndim = " + std::to_string(ndim()) + ")"); |
| 116 | + } |
| 117 | + |
| 118 | + template<typename... Ix> size_type get_byte_offset(Ix... index) const { |
| 119 | + const size_type idx[] = { static_cast<size_type>(index)... }; |
| 120 | + if (!std::equal(idx + 0, idx + sizeof...(index), shape(), std::less<size_type>{})) { |
| 121 | + auto mismatch = std::mismatch(idx + 0, idx + sizeof...(index), shape(), std::less<size_type>{}); |
| 122 | + throw index_error(std::string("index ") + std::to_string(*mismatch.first) + |
| 123 | + " is out of bounds for axis " + std::to_string(mismatch.first - idx) + |
| 124 | + " with size " + std::to_string(*mismatch.second)); |
| 125 | + } |
| 126 | + return std::inner_product(idx + 0, idx + sizeof...(index), strides(), size_type{0}); |
| 127 | + } |
| 128 | + |
| 129 | + size_type get_byte_offset() const { return 0; } |
| 130 | + |
| 131 | + static std::vector<size_type> default_strides(const std::vector<size_type>& shape, |
| 132 | + size_type itemsize) { |
| 133 | + auto ndim = shape.size(); |
| 134 | + std::vector<size_type> strides(ndim); |
| 135 | + if (ndim) { |
| 136 | + std::fill(strides.begin(), strides.end(), itemsize); |
| 137 | + for (size_type i = 0; i < ndim - 1; i++) |
| 138 | + for (size_type j = 0; j < ndim - 1 - i; j++) |
| 139 | + strides[j] *= shape[ndim - 1 - i]; |
| 140 | + } |
| 141 | + return strides; |
| 142 | + } |
| 143 | + }; |
| 144 | +}} // namespace pybind11::backport |
0 commit comments