1 // Copyright 2012-2013 Tinyarray authors.
2 //
3 // This file is part of Tinyarray. It is subject to the license terms in the
4 // file LICENSE.rst found in the top-level directory of this distribution and
5 // at https://gitlab.kwant-project.org/kwant/tinyarray/blob/master/LICENSE.rst.
6 // A list of Tinyarray authors can be found in the README.rst file at the
7 // top-level directory of this distribution and at
8 // https://gitlab.kwant-project.org/kwant/tinyarray.
9
10 #ifndef ARRAY_HH
11 #define ARRAY_HH
12
13 #include <complex>
14 typedef std::complex<double> Complex;
15
16 const int max_ndim = 16;
17
18 // First constant must be 0, the last one must be `NONE'.
19 enum Dtype {LONG = 0, DOUBLE, COMPLEX, NONE};
20 const Dtype default_dtype = DOUBLE;
21 #define DEFAULT_DTYPE "float"
22
23 extern const char *dtype_names[];
24
25 #define DTYPE_DISPATCH(func) {func<long>, func<double>, func<Complex>}
26
27 // First constant must be 0, the last one must be `UNKNOWN'.
28 enum Format {INT32_LE = 0, INT32_BE, INT64_LE, INT64_BE,
29 FLOAT64_LE, FLOAT64_BE, COMPLEX128_LE, COMPLEX128_BE,
30 UNKNOWN};
31
32 extern const char *format_names[];
33
34 extern Format format_by_dtype[];
35
36 // We use the ob_size field in a clever way to encode either the length of a
37 // 1-d array, or the number of dimensions for multi-dimensional arrays. The
38 // following code codifies the conventions.
39 class Array_base {
40 public:
ndim_shape(int * ndim,size_t ** shape) const41 void ndim_shape(int *ndim, size_t **shape) const {
42 const Py_ssize_t ob_size = ob_base.ob_size;
43 if (ob_size >= 0) {
44 if (ndim) *ndim = 1;
45 if (shape) *shape = (size_t*)&ob_base.ob_size;
46 } else if (ob_size < -1) {
47 if (ndim) *ndim = static_cast<int>(-ob_size);
48 if (shape) *shape = (size_t*)((char*)this + sizeof(Array_base));
49 } else {
50 if (ndim) *ndim = 0;
51 if (shape) *shape = 0;
52 }
53 }
54
55 protected:
56 PyVarObject ob_base;
57 };
58
59 // Python 3 support macros for module creation
60 #if PY_MAJOR_VERSION >= 3
61 // module creation
62 #define MOD_DEF(ob, name, methods, doc) \
63 static struct PyModuleDef moduledef = { \
64 PyModuleDef_HEAD_INIT, \
65 name, \
66 doc, \
67 -1, \
68 methods, \
69 }; \
70 ob = PyModule_Create(&moduledef);
71 // init function declaration
72 #define MOD_INIT_FUNC(name) PyMODINIT_FUNC PyInit_##name()
73 // init function declaration for use in Array class def
74 #define MOD_INIT_FRIEND(name) PyObject* PyInit_##name()
75 // init function return values
76 #define MOD_ERROR_VAL NULL
77 #define MOD_SUCCESS_VAL(val) val
78 // Python 2
79 #else
80 // module creation
81 #define MOD_DEF(ob, name, methods, doc) \
82 ob = Py_InitModule3(name, methods, doc);
83 // init function declaration
84 #define MOD_INIT_FUNC(name) PyMODINIT_FUNC init##name()
85 // init function declaration for use in Array class def
86 #define MOD_INIT_FRIEND(name) void init##name()
87 // init function return values
88 #define MOD_ERROR_VAL
89 #define MOD_SUCCESS_VAL(val)
90 #endif
91
92 MOD_INIT_FUNC(tinyarray);
93
94 template <typename T>
95 class Array : public Array_base {
96 public:
data()97 T *data() {
98 if (ob_base.ob_size >= -1) {
99 // ndim == 0 or 1
100 return ob_item;
101 } else {
102 // ndim > 1
103 return ob_item + (-ob_base.ob_size * sizeof(size_t) +
104 sizeof(T) - 1) / sizeof(T);
105 }
106 }
107
108 ssize_t object_size() const;
109
check_exact(PyObject * candidate)110 static bool check_exact(PyObject *candidate) {
111 return (Py_TYPE(candidate) == &pytype);
112 }
113
114 static Array<T> *make(int ndim, size_t size);
115 static Array<T> *make(int ndim, const size_t *shape, size_t *size = 0);
116
117 static const char *pyname, *pyformat;
118 private:
119 T ob_item[1];
120
121 static PyMethodDef methods[];
122 static PySequenceMethods as_sequence;
123 static PyMappingMethods as_mapping;
124 static PyBufferProcs as_buffer;
125 static PyNumberMethods as_number;
126 static PyTypeObject pytype;
127
128 friend Dtype get_dtype(PyObject *obj);
129 friend MOD_INIT_FRIEND(tinyarray);
130 };
131
132 int load_index_seq_as_long(PyObject *obj, long *out, int maxlen);
133 int load_index_seq_as_ulong(PyObject *obj, unsigned long *uout,
134 int maxlen, const char *errmsg = 0);
135
calc_size(int ndim,const size_t * shape)136 inline size_t calc_size(int ndim, const size_t *shape)
137 {
138 if (ndim == 0) return 1;
139 size_t result = shape[0];
140 for (int d = 1; d < ndim; ++d) result *= shape[d];
141 return result;
142 }
143
get_dtype(PyObject * obj)144 inline Dtype get_dtype(PyObject *obj)
145 {
146 PyTypeObject *pytype = Py_TYPE(obj);
147 if (pytype == &Array<long>::pytype) return LONG;
148 if (pytype == &Array<double>::pytype) return DOUBLE;
149 if (pytype == &Array<Complex>::pytype) return COMPLEX;
150 return NONE;
151 }
152
153 PyObject *array_from_arraylike(PyObject *in, Dtype *dtype,
154 Dtype dtype_min = Dtype(0),
155 bool as_matrix = false);
156
157 // Coerced_dtype will contain the common dtype of the coerced arrays.
158 int coerce_to_arrays(PyObject **a, PyObject **b, Dtype *coerced_dtype);
159
160 template <typename T> PyObject *transpose(PyObject *in, PyObject *dummy);
161
162 template <typename T>
object_size() const163 ssize_t Array<T>::object_size() const
164 {
165 int ndim;
166 size_t *shape;
167 ndim_shape(&ndim, &shape);
168
169 // this does the same calculation as in Array<T>::make
170 Py_ssize_t size = calc_size(ndim, shape);
171 if (ndim > 1)
172 // if array is > 1D then the shape is stored in
173 // the start of the data buffer
174 size += (ndim * sizeof(size_t) + sizeof(T) - 1) / sizeof(T);
175 size = size * sizeof(T); // element count -> byte count
176 size += Array<T>::pytype.tp_basicsize; // add Python object overhead
177
178 return size;
179 }
180
181 #endif // !ARRAY_HH
182