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 #define PY_SSIZE_T_CLEAN
11 #include <Python.h>
12 #include "array.hh"
13 #include "arithmetic.hh"
14 #include "functions.hh"
15
16 namespace {
17
dtype_converter(const PyObject * ob,Dtype * dtype)18 int dtype_converter(const PyObject *ob, Dtype *dtype)
19 {
20 if (ob == Py_None) {
21 *dtype = default_dtype;
22 #if PY_MAJOR_VERSION < 3
23 } else if (ob == (PyObject *)(&PyInt_Type) ||
24 ob == (PyObject *)(&PyLong_Type)) {
25 #else
26 } else if (ob == (PyObject *)(&PyLong_Type)) {
27 #endif
28 *dtype = LONG;
29 } else if (ob == (PyObject *)(&PyFloat_Type)) {
30 *dtype = DOUBLE;
31 } else if (ob == (PyObject *)(&PyComplex_Type)) {
32 *dtype = COMPLEX;
33 } else {
34 PyErr_SetString(PyExc_TypeError, "Invalid dtype.");
35 return 0;
36 }
37 return 1;
38 }
39
40 template <typename T>
reconstruct(int ndim,const size_t * shape,const void * src_data_,unsigned size_in_bytes)41 PyObject *reconstruct(int ndim, const size_t *shape,
42 const void *src_data_, unsigned size_in_bytes)
43 {
44 const T *src_data = (const T*)src_data_;
45 size_t size;
46 Array<T> *result = Array<T>::make(ndim, shape, &size);
47 if (!result) return 0;
48 if (size * sizeof(T) != size_in_bytes) {
49 PyErr_SetString(PyExc_ValueError,
50 "Data length mismatch during tinyarray unpickling.");
51 return 0;
52 }
53 T *data = result->data();
54 for (size_t i = 0; i < size; ++i) data[i] = src_data[i];
55 return (PyObject*)result;
56 }
57
58 PyObject *(*reconstruct_dtable[])(int, const size_t*, const void*, unsigned) =
59 DTYPE_DISPATCH(reconstruct);
60
reconstruct(PyObject *,PyObject * args)61 PyObject *reconstruct(PyObject *, PyObject *args)
62 {
63 PyObject *pyshape;
64 Format format;
65 const void *data;
66 Py_ssize_t size_in_bytes;
67 if (!PyArg_ParseTuple(args, "Ois#", &pyshape, &format,
68 &data, &size_in_bytes))
69 return 0;
70
71 Dtype dtype = Dtype(0);
72 while (true) {
73 if (format_by_dtype[int(dtype)] == format) break;
74 dtype = Dtype(int(dtype) + 1);
75 if (dtype == NONE) {
76 if (format < 0 || format > UNKNOWN)
77 format = UNKNOWN;
78 PyErr_Format(PyExc_TypeError, "Cannot unpickle %s.",
79 format_names[format]);
80 return 0;
81 }
82 }
83
84 unsigned long shape_as_ulong[max_ndim];
85 int ndim = load_index_seq_as_ulong(pyshape, shape_as_ulong, max_ndim,
86 "Negative dimensions are not allowed.");
87 if (ndim == -1) return 0;
88
89 size_t shape[max_ndim];
90 for (int d = 0; d < ndim; ++d) shape[d] = shape_as_ulong[d];
91 return reconstruct_dtable[int(dtype)](ndim, shape, data, size_in_bytes);
92 }
93
94 template <typename T>
filled(int ndim,const size_t * shape,int value)95 PyObject *filled(int ndim, const size_t *shape, int value)
96 {
97 size_t size;
98 Array<T> *result = Array<T>::make(ndim, shape, &size);
99 if (!result) return 0;
100 T *data = result->data();
101 for (size_t i = 0; i < size; ++i) data[i] = value;
102 return (PyObject*)result;
103 }
104
105 PyObject *(*filled_dtable[])(int, const size_t*, int) =
106 DTYPE_DISPATCH(filled);
107
filled(PyObject * args,int value)108 PyObject *filled(PyObject *args, int value)
109 {
110 PyObject *pyshape;
111 Dtype dtype = default_dtype;
112 if (!PyArg_ParseTuple(args, "O|O&", &pyshape, dtype_converter, &dtype))
113 return 0;
114
115 unsigned long shape_as_ulong[max_ndim];
116 int ndim = load_index_seq_as_ulong(pyshape, shape_as_ulong, max_ndim,
117 "Negative dimensions are not allowed.");
118 if (ndim == -1) return 0;
119
120 size_t shape[max_ndim];
121 for (int d = 0; d < ndim; ++d) shape[d] = shape_as_ulong[d];
122 return filled_dtable[int(dtype)](ndim, shape, value);
123 }
124
zeros(PyObject *,PyObject * args)125 PyObject *zeros(PyObject *, PyObject *args)
126 {
127 return filled(args, 0);
128 }
129
130 PyDoc_STRVAR(zeros_doc,
131 "zeros(shape, dtype=" DEFAULT_DTYPE ")\n\n\
132 Return an array of given shape and type, filled with zeros.");
133
ones(PyObject *,PyObject * args)134 PyObject *ones(PyObject *, PyObject *args)
135 {
136 return filled(args, 1);
137 }
138
139 PyDoc_STRVAR(ones_doc,
140 "ones(shape, dtype=" DEFAULT_DTYPE ")\n\n\
141 Return an array of given shape and type, filled with ones.");
142
143 template <typename T>
identity(size_t n)144 PyObject *identity(size_t n)
145 {
146 size_t size, shape[] = {n, n};
147 Array<T> *result = Array<T>::make(2, shape, &size);
148 if (!result) return 0;
149
150 T *p = result->data();
151 for (size_t i = 1; i < n; ++i) {
152 *p++ = 1;
153 for (T *e = p + n; p < e; ++p)
154 *p = 0;
155 }
156 if (n) *p = 1;
157
158 return (PyObject*)result;
159 }
160
161 PyObject *(*identity_dtable[])(size_t) = DTYPE_DISPATCH(identity);
162
identity(PyObject *,PyObject * args)163 PyObject *identity(PyObject *, PyObject *args)
164 {
165 long n;
166 Dtype dtype = default_dtype;
167 if (!PyArg_ParseTuple(args, "l|O&", &n, dtype_converter, &dtype))
168 return 0;
169 if (n < 0) {
170 PyErr_SetString(PyExc_ValueError,
171 "Negative dimensions are not allowed.");
172 return 0;
173 }
174
175 return identity_dtable[int(dtype)](n);
176 }
177
178 PyDoc_STRVAR(identity_doc,
179 "identity(n, dtype=" DEFAULT_DTYPE ")\n\n\
180 Return an identity matrix of given size and dtype.");
181
array(PyObject *,PyObject * args)182 PyObject *array(PyObject *, PyObject *args)
183 {
184 PyObject *src;
185 Dtype dtype = NONE;
186 if (!PyArg_ParseTuple(args, "O|O&", &src, dtype_converter, &dtype))
187 return 0;
188 return array_from_arraylike(src, &dtype);
189 }
190
191 PyDoc_STRVAR(array_doc,
192 "array(object, [dtype])\n\n\
193 Create an array from something array-like.\n\
194 Valid inputs are numbers, sequences (of sequences, ...) of numbers, NumPy\n\
195 and tinyarray arrays, and objects supporting the buffer protocol.");
196
matrix(PyObject *,PyObject * args)197 PyObject *matrix(PyObject *, PyObject *args)
198 {
199 PyObject *src;
200 Dtype dtype = NONE;
201 if (!PyArg_ParseTuple(args, "O|O&", &src, dtype_converter, &dtype))
202 return 0;
203 return array_from_arraylike(src, &dtype, Dtype(0), true);
204 }
205
206 PyDoc_STRVAR(matrix_doc,
207 "matrix(object, [dtype])\n\n\
208 Create an 2-d array from something array-like.\n\
209 Valid inputs are the same as for array(), however the input is promoted to 2-d.\n\
210 A ``ValueError`` is raised if the input has more than 2 dimensions.");
211
212 PyObject *(*transpose_dtable[])(PyObject*, PyObject *) =
213 DTYPE_DISPATCH(transpose);
214
transpose(PyObject *,PyObject * args)215 PyObject *transpose(PyObject *, PyObject *args)
216 {
217 PyObject *a;
218 if (!PyArg_ParseTuple(args, "O", &a)) return 0;
219 Dtype dtype = NONE;
220 a = array_from_arraylike(a, &dtype);
221 if (!a) return 0;
222 PyObject *result = transpose_dtable[int(dtype)](a, 0);
223 Py_DECREF(a);
224 return result;
225 }
226
227 PyDoc_STRVAR(transpose_doc,
228 "transpose(a)\n\n\
229 Return a copy of the given array with reversed order of dimensions.");
230
dot(PyObject *,PyObject * args)231 PyObject *dot(PyObject *, PyObject *args)
232 {
233 PyObject *a, *b;
234 if (!PyArg_ParseTuple(args, "OO", &a, &b)) return 0;
235 return dot_product(a, b);
236 }
237
238 PyDoc_STRVAR(dot_doc,
239 "dot(a, b)\n\n\
240 Return the dot product of two arrays.\n\n\
241 For 2-d arrays, the dot product is equivalent to matrix multiplication; for 1-d\n\
242 arrays, to a scalar product of vectors. In the general case, it is equivalent\n\
243 to a sum over the last axis of a and the second-to-last of b, e.g.::\n\n\
244 dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])");
245
246 template <template <typename> class Op>
binary_ufunc(PyObject *,PyObject * args)247 PyObject *binary_ufunc(PyObject *, PyObject *args)
248 {
249 PyObject *a, *b;
250 if (!PyArg_ParseTuple(args, "OO", &a, &b)) return 0;
251 return Binary_op<Op>::apply(a, b);
252 }
253
254 template <template <typename> class Op>
unary_ufunc(PyObject *,PyObject * args)255 PyObject *unary_ufunc(PyObject *, PyObject *args)
256 {
257 static PyObject *(*operation_dtable[])(PyObject*) = {
258 apply_unary_ufunc<Op<long> >,
259 apply_unary_ufunc<Op<double> >,
260 apply_unary_ufunc<Op<Complex> >
261 };
262
263 PyObject *a;
264 if (!PyArg_ParseTuple(args, "O", &a)) return 0;
265 Dtype dtype = NONE;
266 a = array_from_arraylike(a, &dtype);
267 if (!a) return 0;
268 PyObject *result = operation_dtable[int(dtype)](a);
269 Py_DECREF(a);
270 return result;
271 }
272
273 template <typename Kind>
unary_ufunc_round(PyObject *,PyObject * args)274 PyObject *unary_ufunc_round(PyObject *, PyObject *args)
275 {
276 static PyObject *(*operation_dtable[])(PyObject*) = {
277 apply_unary_ufunc<Round<Kind, long> >,
278 apply_unary_ufunc<Round<Kind, double> >,
279 apply_unary_ufunc<Round<Kind, Complex> >
280 };
281
282 PyObject *a;
283 if (!PyArg_ParseTuple(args, "O", &a)) return 0;
284 Dtype dtype = NONE;
285 a = array_from_arraylike(a, &dtype);
286 if (!a) return 0;
287 PyObject *result = operation_dtable[int(dtype)](a);
288 Py_DECREF(a);
289 return result;
290 }
291
292 PyDoc_STRVAR(binary_ufunc_doc,
293 "Operates elementwise on two arrays, returns an array of the same shape.");
294
295 PyDoc_STRVAR(unary_ufunc_doc,
296 "Operates elementwise on an array, returns an array of the same shape.");
297
298 } // Anonymous namespace
299
300 PyMethodDef functions[] = {
301 {"_reconstruct", reconstruct, METH_VARARGS},
302 {"zeros", zeros, METH_VARARGS, zeros_doc},
303 {"ones", ones, METH_VARARGS, ones_doc},
304 {"identity", identity, METH_VARARGS, identity_doc},
305 {"array", array, METH_VARARGS, array_doc},
306 {"matrix", matrix, METH_VARARGS, matrix_doc},
307 {"transpose", transpose, METH_VARARGS, transpose_doc},
308 {"dot", dot, METH_VARARGS, dot_doc},
309
310 {"add", binary_ufunc<Add>, METH_VARARGS, binary_ufunc_doc},
311 {"subtract", binary_ufunc<Subtract>, METH_VARARGS, binary_ufunc_doc},
312 {"multiply", binary_ufunc<Multiply>, METH_VARARGS, binary_ufunc_doc},
313 #if PY_MAJOR_VERSION < 3
314 {"divide", binary_ufunc<Divide>, METH_VARARGS, binary_ufunc_doc},
315 #else
316 {"divide", binary_ufunc<True_divide>, METH_VARARGS, binary_ufunc_doc},
317 #endif
318 {"remainder", binary_ufunc<Remainder>, METH_VARARGS, binary_ufunc_doc},
319 {"floor_divide", binary_ufunc<Floor_divide>, METH_VARARGS,
320 binary_ufunc_doc},
321
322 {"negative", unary_ufunc<Negative>, METH_VARARGS, unary_ufunc_doc},
323 {"abs", unary_ufunc<Absolute>, METH_VARARGS, unary_ufunc_doc},
324 {"absolute", unary_ufunc<Absolute>, METH_VARARGS, unary_ufunc_doc},
325 {"conjugate", unary_ufunc<Conjugate>, METH_VARARGS, unary_ufunc_doc},
326 {"round", unary_ufunc_round<Nearest>, METH_VARARGS, unary_ufunc_doc},
327 {"floor", unary_ufunc_round<Floor>, METH_VARARGS, unary_ufunc_doc},
328 {"ceil", unary_ufunc_round<Ceil>, METH_VARARGS, unary_ufunc_doc},
329
330 {0, 0, 0, 0} // Sentinel
331 };
332