functions.cc 8.42 KB
Newer Older
1
2
3
4
5
6
7
8
// Copyright 2012-2013 Tinyarray authors.
//
// This file is part of Tinyarray.  It is subject to the license terms in the
// LICENSE file found in the top-level directory of this distribution and at
// http://kwant-project.org/tinyarray/license.  A list of Tinyarray authors can
// be found in the README file at the top-level directory of this distribution
// and at http://kwant-project.org/tinyarray.

Christoph Groth's avatar
Christoph Groth committed
9
10
#include <Python.h>
#include "array.hh"
11
#include "arithmetic.hh"
Christoph Groth's avatar
Christoph Groth committed
12
13
#include "functions.hh"

14
namespace {
Christoph Groth's avatar
Christoph Groth committed
15

16
int dtype_converter(const PyObject *ob, Dtype *dtype)
Christoph Groth's avatar
Christoph Groth committed
17
{
18
19
20
21
    if (ob == Py_None) {
        *dtype = default_dtype;
    } else if (ob == (PyObject *)(&PyInt_Type) ||
               ob == (PyObject *)(&PyLong_Type)) {
Christoph Groth's avatar
Christoph Groth committed
22
        *dtype = LONG;
23
    } else if (ob == (PyObject *)(&PyFloat_Type)) {
Christoph Groth's avatar
Christoph Groth committed
24
        *dtype = DOUBLE;
25
    } else if (ob == (PyObject *)(&PyComplex_Type)) {
Christoph Groth's avatar
Christoph Groth committed
26
        *dtype = COMPLEX;
27
28
    } else {
        PyErr_SetString(PyExc_TypeError, "Invalid dtype.");
Christoph Groth's avatar
Christoph Groth committed
29
        return 0;
30
31
    }
    return 1;
Christoph Groth's avatar
Christoph Groth committed
32
33
}

Joseph Weston's avatar
Joseph Weston committed
34
template <typename T>
35
36
PyObject *reconstruct(int ndim, const size_t *shape,
                      const void *src_data_, unsigned size_in_bytes)
Christoph Groth's avatar
Christoph Groth committed
37
{
38
    const T *src_data = (const T*)src_data_;
39
40
41
    size_t size;
    Array<T> *result = Array<T>::make(ndim, shape, &size);
    if (!result) return 0;
42
43
44
45
46
    if (size * sizeof(T) != size_in_bytes) {
        PyErr_SetString(PyExc_ValueError,
                        "Data length mismatch during tinyarray unpickling.");
        return 0;
    }
47
    T *data = result->data();
48
    for (size_t i = 0; i < size; ++i) data[i] = src_data[i];
49
    return (PyObject*)result;
Christoph Groth's avatar
Christoph Groth committed
50
51
}

52
53
PyObject *(*reconstruct_dtable[])(int, const size_t*, const void*, unsigned) =
    DTYPE_DISPATCH(reconstruct);
Christoph Groth's avatar
Christoph Groth committed
54

55
PyObject *reconstruct(PyObject *, PyObject *args)
Christoph Groth's avatar
Christoph Groth committed
56
{
57
    PyObject *pyshape;
58
59
60
61
62
    Format format;
    const void *data;
    unsigned size_in_bytes;
    if (!PyArg_ParseTuple(args, "Ois#", &pyshape, &format,
                          &data, &size_in_bytes))
Christoph Groth's avatar
Christoph Groth committed
63
64
        return 0;

65
66
67
68
    Dtype dtype = Dtype(0);
    while (true) {
        if (format_by_dtype[int(dtype)] == format) break;
        dtype = Dtype(int(dtype) + 1);
Christoph Groth's avatar
Christoph Groth committed
69
70
71
        if (dtype == NONE) {
            if (format < 0 || format > UNKNOWN)
                format = UNKNOWN;
72
73
74
75
76
77
            PyErr_Format(PyExc_TypeError, "Cannot unpickle %s.",
                         format_names[format]);
            return 0;
        }
    }

Christoph Groth's avatar
Christoph Groth committed
78
79
80
    unsigned long shape_as_ulong[max_ndim];
    int ndim = load_index_seq_as_ulong(pyshape, shape_as_ulong, max_ndim,
                                       "Negative dimensions are not allowed.");
81
    if (ndim == -1) return 0;
Christoph Groth's avatar
Christoph Groth committed
82

Christoph Groth's avatar
Christoph Groth committed
83
84
    size_t shape[max_ndim];
    for (int d = 0; d < ndim; ++d) shape[d] = shape_as_ulong[d];
85
    return reconstruct_dtable[int(dtype)](ndim, shape, data, size_in_bytes);
Christoph Groth's avatar
Christoph Groth committed
86
87
}

88
89
90
91
92
93
94
95
96
97
98
99
100
101
template <typename T>
PyObject *filled(int ndim, const size_t *shape, int value)
{
    size_t size;
    Array<T> *result = Array<T>::make(ndim, shape, &size);
    if (!result) return 0;
    T *data = result->data();
    for (size_t i = 0; i < size; ++i) data[i] = value;
    return (PyObject*)result;
}

PyObject *(*filled_dtable[])(int, const size_t*, int) =
    DTYPE_DISPATCH(filled);

Christoph Groth's avatar
Christoph Groth committed
102
PyObject *filled(PyObject *args, int value)
Joseph Weston's avatar
Joseph Weston committed
103
104
105
106
107
108
109
110
111
112
113
114
115
{
    PyObject *pyshape;
    Dtype dtype = default_dtype;
    if (!PyArg_ParseTuple(args, "O|O&", &pyshape, dtype_converter, &dtype))
        return 0;

    unsigned long shape_as_ulong[max_ndim];
    int ndim = load_index_seq_as_ulong(pyshape, shape_as_ulong, max_ndim,
                                       "Negative dimensions are not allowed.");
    if (ndim == -1) return 0;

    size_t shape[max_ndim];
    for (int d = 0; d < ndim; ++d) shape[d] = shape_as_ulong[d];
116
    return filled_dtable[int(dtype)](ndim, shape, value);
Joseph Weston's avatar
Joseph Weston committed
117
118
}

119
PyObject *zeros(PyObject *, PyObject *args)
Christoph Groth's avatar
Christoph Groth committed
120
{
Christoph Groth's avatar
Christoph Groth committed
121
    return filled(args, 0);
122
}
Christoph Groth's avatar
Christoph Groth committed
123

124
125
PyObject *ones(PyObject *, PyObject *args)
{
Christoph Groth's avatar
Christoph Groth committed
126
    return filled(args, 1);
127
}
Christoph Groth's avatar
Christoph Groth committed
128

129
130
131
132
133
134
template <typename T>
PyObject *identity(size_t n)
{
    size_t size, shape[] = {n, n};
    Array<T> *result = Array<T>::make(2, shape, &size);
    if (!result) return 0;
Christoph Groth's avatar
Christoph Groth committed
135

136
137
138
139
140
    T *p = result->data();
    for (size_t i = 1; i < n; ++i) {
        *p++ = 1;
        for (T *e = p + n; p < e; ++p)
            *p = 0;
Christoph Groth's avatar
Christoph Groth committed
141
    }
142
    if (n) *p = 1;
Christoph Groth's avatar
Christoph Groth committed
143

144
    return (PyObject*)result;
Christoph Groth's avatar
Christoph Groth committed
145
146
}

147
148
149
PyObject *(*identity_dtable[])(size_t) = DTYPE_DISPATCH(identity);

PyObject *identity(PyObject *, PyObject *args)
Christoph Groth's avatar
Christoph Groth committed
150
{
151
152
153
154
155
156
157
158
    long n;
    Dtype dtype = default_dtype;
    if (!PyArg_ParseTuple(args, "l|O&", &n, dtype_converter, &dtype))
        return 0;
    if (n < 0) {
        PyErr_SetString(PyExc_ValueError,
                        "Negative dimensions are not allowed.");
        return 0;
Christoph Groth's avatar
Christoph Groth committed
159
    }
160
161

    return identity_dtable[int(dtype)](n);
Christoph Groth's avatar
Christoph Groth committed
162
163
}

164
PyObject *array(PyObject *, PyObject *args)
Christoph Groth's avatar
Christoph Groth committed
165
166
{
    PyObject *src;
Christoph Groth's avatar
Christoph Groth committed
167
    Dtype dtype = NONE;
Christoph Groth's avatar
Christoph Groth committed
168
169
    if (!PyArg_ParseTuple(args, "O|O&", &src, dtype_converter, &dtype))
        return 0;
170
    return array_from_arraylike(src, &dtype);
Christoph Groth's avatar
Christoph Groth committed
171
172
}

Christoph Groth's avatar
Christoph Groth committed
173
174
175
PyObject *matrix(PyObject *, PyObject *args)
{
    PyObject *src;
Christoph Groth's avatar
Christoph Groth committed
176
    Dtype dtype = NONE;
Christoph Groth's avatar
Christoph Groth committed
177
178
179
180
181
    if (!PyArg_ParseTuple(args, "O|O&", &src, dtype_converter, &dtype))
        return 0;
    return matrix_from_arraylike(src, &dtype);
}

182
183
PyObject *(*transpose_dtable[])(PyObject*, PyObject *) =
  DTYPE_DISPATCH(transpose);
Christoph Groth's avatar
Christoph Groth committed
184
185
186
187
188

PyObject *transpose(PyObject *, PyObject *args)
{
    PyObject *a;
    if (!PyArg_ParseTuple(args, "O", &a)) return 0;
Christoph Groth's avatar
Christoph Groth committed
189
    Dtype dtype = NONE;
Christoph Groth's avatar
Christoph Groth committed
190
191
    a = array_from_arraylike(a, &dtype);
    if (!a) return 0;
192
    return transpose_dtable[int(dtype)](a, 0);
Christoph Groth's avatar
Christoph Groth committed
193
194
}

195
PyObject *dot(PyObject *, PyObject *args)
Christoph Groth's avatar
Christoph Groth committed
196
{
197
    PyObject *a, *b;
198
    if (!PyArg_ParseTuple(args, "OO", &a, &b)) return 0;
199
    return dot_product(a, b);
Christoph Groth's avatar
Christoph Groth committed
200
201
}

202
template <template <typename> class Op>
203
PyObject *binary_ufunc(PyObject *, PyObject *args)
204
205
206
207
208
209
{
    PyObject *a, *b;
    if (!PyArg_ParseTuple(args, "OO", &a, &b)) return 0;
    return Binary_op<Op>::apply(a, b);
}

210
211
212
213
214
215
216
217
218
219
220
template <template <typename> class Op>
PyObject *unary_ufunc(PyObject *, PyObject *args)
{
    static PyObject *(*operation_dtable[])(PyObject*) = {
        apply_unary_ufunc<Op<long> >,
        apply_unary_ufunc<Op<double> >,
        apply_unary_ufunc<Op<Complex> >
    };

    PyObject *a;
    if (!PyArg_ParseTuple(args, "O", &a)) return 0;
Christoph Groth's avatar
Christoph Groth committed
221
    Dtype dtype = NONE;
222
223
    a = array_from_arraylike(a, &dtype);
    if (!a) return 0;
224
225
226
227
228
    PyObject *result = operation_dtable[int(dtype)](a);
    Py_DECREF(a);
    return result;
}

229
230
231
232
233
234
235
236
237
238
239
template <typename Kind>
PyObject *unary_ufunc_round(PyObject *, PyObject *args)
{
    static PyObject *(*operation_dtable[])(PyObject*) = {
        apply_unary_ufunc<Round<Kind, long> >,
        apply_unary_ufunc<Round<Kind, double> >,
        apply_unary_ufunc<Round<Kind, Complex> >
    };

    PyObject *a;
    if (!PyArg_ParseTuple(args, "O", &a)) return 0;
Christoph Groth's avatar
Christoph Groth committed
240
    Dtype dtype = NONE;
241
242
243
244
245
246
247
    a = array_from_arraylike(a, &dtype);
    if (!a) return 0;
    PyObject *result = operation_dtable[int(dtype)](a);
    Py_DECREF(a);
    return result;
}

248
249
} // Anonymous namespace

250
251
252
253
254
255
// TODO: once we can require gcc 4.7, use the following type aliases instead of
// the function unary_ufunc_round which is a stopgap.
//
// template <typename T> using Round_nearest = Round<Nearest, T>;
// template <typename T> using Round_floor = Round<Floor, T>;
// template <typename T> using Round_ceil = Round<Ceil, T>;
256

Christoph Groth's avatar
Christoph Groth committed
257
PyMethodDef functions[] = {
258
    {"_reconstruct", reconstruct, METH_VARARGS},
Christoph Groth's avatar
Christoph Groth committed
259
260
261
262
    {"zeros", zeros, METH_VARARGS},
    {"ones", ones, METH_VARARGS},
    {"identity", identity, METH_VARARGS},
    {"array", array, METH_VARARGS},
Christoph Groth's avatar
Christoph Groth committed
263
    {"matrix", matrix, METH_VARARGS},
Christoph Groth's avatar
Christoph Groth committed
264
    {"transpose", transpose, METH_VARARGS},
Christoph Groth's avatar
Christoph Groth committed
265
    {"dot", dot, METH_VARARGS},
266
267
268
269
270
271
272
273
274
275
276

    {"add", binary_ufunc<Add>, METH_VARARGS},
    {"subtract", binary_ufunc<Subtract>, METH_VARARGS},
    {"multiply", binary_ufunc<Multiply>, METH_VARARGS},
    {"divide", binary_ufunc<Divide>, METH_VARARGS},
    {"remainder", binary_ufunc<Remainder>, METH_VARARGS},
    {"floor_divide", binary_ufunc<Floor_divide>, METH_VARARGS},

    {"negative", unary_ufunc<Negative>, METH_VARARGS},
    {"abs", unary_ufunc<Absolute>, METH_VARARGS},
    {"absolute", unary_ufunc<Absolute>, METH_VARARGS},
Christoph Groth's avatar
Christoph Groth committed
277
    {"conjugate", unary_ufunc<Conjugate>, METH_VARARGS},
278
279
280
281
282
283
// {"round", unary_ufunc<Round_nearest>, METH_VARARGS},
// {"floor", unary_ufunc<Round_floor>, METH_VARARGS},
// {"ceil", unary_ufunc<Round_ceil>, METH_VARARGS},
    {"round", unary_ufunc_round<Nearest>, METH_VARARGS},
    {"floor", unary_ufunc_round<Floor>, METH_VARARGS},
    {"ceil", unary_ufunc_round<Ceil>, METH_VARARGS},
284

Christoph Groth's avatar
Christoph Groth committed
285
286
    {0, 0, 0, 0}                // Sentinel
};