array.hh 3.46 KB
Newer Older
Christoph Groth's avatar
Christoph Groth committed
1
2
3
#ifndef ARRAY_HH
#define ARRAY_HH

4
5
#include <complex>
typedef std::complex<double> Complex;
Christoph Groth's avatar
Christoph Groth committed
6

7
8
const int max_ndim = 16;

9
// First constant must be 0, the last one must be `NONE'.
Christoph Groth's avatar
Christoph Groth committed
10
11
enum Dtype {LONG = 0, DOUBLE, COMPLEX, NONE};
const Dtype default_dtype = DOUBLE;
12

Christoph Groth's avatar
Christoph Groth committed
13
14
extern const char *dtype_names[];

15
16
#define DTYPE_DISPATCH(func) {func<long>, func<double>, func<Complex>}

17
18
19
20
21
22
23
24
25
// First constant must be 0, the last one must be `UNKNOWN'.
enum Format {INT32_LE = 0, INT32_BE, INT64_LE, INT64_BE,
             FLOAT64_LE, FLOAT64_BE, COMPLEX128_LE, COMPLEX128_BE,
             UNKNOWN};

extern const char *format_names[];

extern Format format_by_dtype[];

26
27
28
29
30
31
32
33
34
35
36
// We use the ob_size field in a clever way to encode either the length of a
// 1-d array, or the number of dimensions for multi-dimensional arrays.  The
// following code codifies the conventions.
class Array_base {
public:
    void ndim_shape(int *ndim, size_t **shape) {
        const Py_ssize_t ob_size = ob_base.ob_size;
        if (ob_size >= 0) {
            if (ndim) *ndim = 1;
            if (shape) *shape = (size_t*)&ob_base.ob_size;
        } else if (ob_size < -1) {
37
            if (ndim) *ndim = static_cast<int>(-ob_size);
38
39
40
41
42
43
44
45
46
            if (shape) *shape = (size_t*)((char*)this + sizeof(Array_base));
        } else {
            if (ndim) *ndim = 0;
            if (shape) *shape = 0;
        }
    }

protected:
    PyVarObject ob_base;
Christoph Groth's avatar
Christoph Groth committed
47
48
};

49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
extern "C" void inittinyarray();

template <typename T>
class Array : public Array_base {
public:
    T *data() {
        if (ob_base.ob_size >= -1) {
            // ndim == 0 or 1
            return ob_item;
        } else {
            // ndim > 1
            return ob_item + (-ob_base.ob_size * sizeof(size_t) +
                              sizeof(T) - 1) / sizeof(T);
        }
    }

    static bool check_exact(PyObject *candidate) {
        return (Py_TYPE(candidate) == &pytype);
    }

    static Array<T> *make(int ndim, size_t size);
    static Array<T> *make(int ndim, const size_t *shape, size_t *size = 0);

    static const char *pyname, *pyformat;
private:
    T ob_item[1];

Christoph Groth's avatar
Christoph Groth committed
76
    static PyMethodDef methods[];
77
78
79
    static PySequenceMethods as_sequence;
    static PyMappingMethods as_mapping;
    static PyBufferProcs as_buffer;
80
    static PyNumberMethods as_number;
81
82
83
84
85
86
    static PyTypeObject pytype;

    friend Dtype get_dtype(PyObject *obj);
    friend void inittinyarray();
};

87
88
89
int load_index_seq_as_long(PyObject *obj, long *out, int maxlen);
int load_index_seq_as_ulong(PyObject *obj, unsigned long *uout,
                            int maxlen, const char *errmsg = 0);
90
91
92
93
94
95
96
97

inline size_t calc_size(int ndim, const size_t *shape)
{
    if (ndim == 0) return 1;
    size_t result = shape[0];
    for (int d = 1; d < ndim; ++d) result *= shape[d];
    return result;
}
Christoph Groth's avatar
Christoph Groth committed
98

99
100
101
inline Dtype get_dtype(PyObject *obj)
{
    PyTypeObject *pytype = Py_TYPE(obj);
Christoph Groth's avatar
Christoph Groth committed
102
103
104
105
    if (pytype == &Array<long>::pytype) return LONG;
    if (pytype == &Array<double>::pytype) return DOUBLE;
    if (pytype == &Array<Complex>::pytype) return COMPLEX;
    return NONE;
106
}
Christoph Groth's avatar
Christoph Groth committed
107

Christoph Groth's avatar
Christoph Groth committed
108
109
110
111
PyObject *array_from_arraylike(PyObject *in, Dtype *dtype,
                               Dtype dtype_min = Dtype(0));
PyObject *matrix_from_arraylike(PyObject *in, Dtype *dtype,
                               Dtype dtype_min = Dtype(0));
112

113
114
// Coerced_dtype will contain the common dtype of the coerced arrays.
int coerce_to_arrays(PyObject **a, PyObject **b, Dtype *coerced_dtype);
Christoph Groth's avatar
Christoph Groth committed
115

116
template <typename T> PyObject *transpose(PyObject *in, PyObject *dummy);
Christoph Groth's avatar
Christoph Groth committed
117

Christoph Groth's avatar
Christoph Groth committed
118
#endif // !ARRAY_HH