xref: /petsc/src/binding/petsc4py/src/petsc4py/PETSc/arraynpy.pxi (revision f2300f31d01ffc99ca7a05a43941a14c1c1d9061) 
1# --------------------------------------------------------------------
2
3cdef extern from "<petsc4py/numpy.h>":
4
5    int import_array "_import_array" () except -1
6
7    ctypedef long npy_intp
8
9    ctypedef extern class numpy.dtype [object PyArray_Descr]:
10        pass
11
12    ctypedef extern class numpy.ndarray [object PyArrayObject]:
13        pass
14
15    void* PyArray_DATA(ndarray)
16    npy_intp PyArray_SIZE(ndarray)
17    int PyArray_NDIM(ndarray)
18    npy_intp* PyArray_DIMS(ndarray)
19    npy_intp PyArray_DIM(ndarray, int)
20    npy_intp PyArray_MultiplyList(const npy_intp*, int)
21
22    enum: NPY_INTP
23    dtype   PyArray_DescrFromType(int)
24    object  PyArray_TypeObjectFromType(int)
25
26    enum: NPY_ARRAY_ALIGNED
27    enum: NPY_ARRAY_WRITEABLE
28    enum: NPY_ARRAY_NOTSWAPPED
29    enum: NPY_ARRAY_CARRAY
30    enum: NPY_ARRAY_CARRAY_RO
31    enum: NPY_ARRAY_FARRAY
32    enum: NPY_ARRAY_FARRAY_RO
33
34    ndarray PyArray_FROM_O(object)
35    ndarray PyArray_FROM_OT(object, int)
36    ndarray PyArray_FROM_OTF(object, int, int)
37
38    ndarray PyArray_Copy(ndarray)
39    ndarray PyArray_ArangeObj(object, object, object, dtype)
40    ndarray PyArray_EMPTY(int, npy_intp[], int, int)
41    ndarray PyArray_ZEROS(int, npy_intp[], int, int)
42
43    bint PyArray_ISCONTIGUOUS(ndarray)
44    bint PyArray_ISFORTRAN(ndarray)
45    ctypedef enum NPY_ORDER:
46        NPY_ANYORDER
47        NPY_CORDER
48        NPY_FORTRANORDER
49    ndarray PyArray_NewCopy(ndarray, NPY_ORDER)
50
51    ctypedef struct PyObject
52    ctypedef struct PyTypeObject
53    ndarray PyArray_New(PyTypeObject*, int, npy_intp[], int, npy_intp[], void*, int, int, PyObject*)
54    ndarray PyArray_SimpleNewFromData(int, npy_intp[], int, void*)
55
56    int PyArray_SetBaseObject(ndarray, object) except -1
57
58cdef extern from "<petsc4py/numpy.h>":
59
60    enum: NPY_INT
61    enum: NPY_DOUBLE
62
63    enum: NPY_PETSC_BOOL
64    enum: NPY_PETSC_INT
65    enum: NPY_PETSC_REAL
66    enum: NPY_PETSC_SCALAR
67    enum: NPY_PETSC_COMPLEX
68
69
70# --------------------------------------------------------------------
71
72cdef inline ndarray asarray(object ob):
73    return PyArray_FROM_O(ob)
74
75cdef inline ndarray arange(start, stop, stride):
76    cdef dtype descr = <dtype> PyArray_DescrFromType(NPY_PETSC_INT)
77    return PyArray_ArangeObj(start, stop, stride, descr)
78
79# --------------------------------------------------------------------
80
81cdef inline ndarray empty_b(PetscInt size):
82    cdef npy_intp s = <npy_intp> size
83    return PyArray_EMPTY(1, &s, NPY_PETSC_BOOL, 0)
84
85cdef inline ndarray empty_i(PetscInt size):
86    cdef npy_intp s = <npy_intp> size
87    return PyArray_EMPTY(1, &s, NPY_PETSC_INT, 0)
88
89cdef inline ndarray empty_r(PetscInt size):
90    cdef npy_intp s = <npy_intp> size
91    return PyArray_EMPTY(1, &s, NPY_PETSC_REAL, 0)
92
93cdef inline ndarray empty_s(PetscInt size):
94    cdef npy_intp s = <npy_intp> size
95    return PyArray_EMPTY(1, &s, NPY_PETSC_SCALAR, 0)
96
97cdef inline ndarray empty_c(PetscInt size):
98    cdef npy_intp s = <npy_intp> size
99    return PyArray_EMPTY(1, &s, NPY_PETSC_COMPLEX, 0)
100
101cdef inline ndarray empty_p(PetscInt size):
102    cdef npy_intp s = <npy_intp> size
103    return PyArray_EMPTY(1, &s, NPY_INTP, 0)
104
105# --------------------------------------------------------------------
106
107cdef inline ndarray array_i(PetscInt size, const PetscInt* data):
108    cdef npy_intp s = <npy_intp> size
109    cdef ndarray ary = PyArray_EMPTY(1, &s, NPY_PETSC_INT, 0)
110    if data != NULL:
111        memcpy(PyArray_DATA(ary), data, <size_t>size*sizeof(PetscInt))
112    return ary
113
114cdef inline ndarray array_r(PetscInt size, const PetscReal* data):
115    cdef npy_intp s = <npy_intp> size
116    cdef ndarray ary = PyArray_EMPTY(1, &s, NPY_PETSC_REAL, 0)
117    if data != NULL:
118        memcpy(PyArray_DATA(ary), data, <size_t>size*sizeof(PetscReal))
119    return ary
120
121cdef inline ndarray array_b(PetscInt size, const PetscBool* data):
122    cdef npy_intp s = <npy_intp> size
123    cdef ndarray ary = PyArray_EMPTY(1, &s, NPY_PETSC_BOOL, 0)
124    if data != NULL:
125        memcpy(PyArray_DATA(ary), data, <size_t>size*sizeof(PetscBool))
126    return ary
127
128cdef inline ndarray array_s(PetscInt size, const PetscScalar* data):
129    cdef npy_intp s = <npy_intp> size
130    cdef ndarray ary = PyArray_EMPTY(1, &s, NPY_PETSC_SCALAR, 0)
131    if data != NULL:
132        memcpy(PyArray_DATA(ary), data, <size_t>size*sizeof(PetscScalar))
133    return ary
134
135# --------------------------------------------------------------------
136
137cdef inline ndarray iarray(object ob, int typenum):
138    cdef ndarray ary = PyArray_FROM_OTF(
139        ob, typenum, NPY_ARRAY_ALIGNED|NPY_ARRAY_NOTSWAPPED)
140    if PyArray_ISCONTIGUOUS(ary): return ary
141    if PyArray_ISFORTRAN(ary):    return ary
142    return PyArray_Copy(ary)
143
144cdef inline ndarray iarray_i(object ob, PetscInt* size, PetscInt** data):
145    cdef ndarray ary = iarray(ob, NPY_PETSC_INT)
146    if size != NULL: size[0] = <PetscInt> PyArray_SIZE(ary)
147    if data != NULL: data[0] = <PetscInt*> PyArray_DATA(ary)
148    return ary
149
150cdef inline ndarray iarray_r(object ob, PetscInt* size, PetscReal** data):
151    cdef ndarray ary = iarray(ob, NPY_PETSC_REAL)
152    if size != NULL: size[0] = <PetscInt> PyArray_SIZE(ary)
153    if data != NULL: data[0] = <PetscReal*> PyArray_DATA(ary)
154    return ary
155
156cdef inline ndarray iarray_b(object ob, PetscInt* size, PetscBool** data):
157    cdef ndarray ary = iarray(ob, NPY_PETSC_BOOL)
158    if size != NULL: size[0] = <PetscInt> PyArray_SIZE(ary)
159    if data != NULL: data[0] = <PetscBool*> PyArray_DATA(ary)
160    return ary
161
162cdef inline ndarray iarray_s(object ob, PetscInt* size, PetscScalar** data):
163    cdef ndarray ary = iarray(ob, NPY_PETSC_SCALAR)
164    if size != NULL: size[0] = <PetscInt> PyArray_SIZE(ary)
165    if data != NULL: data[0] = <PetscScalar*> PyArray_DATA(ary)
166    return ary
167
168# --------------------------------------------------------------------
169
170cdef inline ndarray oarray(object ob, int typenum):
171    cdef ndarray ary = PyArray_FROM_OTF(
172        ob, typenum, NPY_ARRAY_ALIGNED|NPY_ARRAY_WRITEABLE|NPY_ARRAY_NOTSWAPPED)
173    if PyArray_ISCONTIGUOUS(ary): return ary
174    if PyArray_ISFORTRAN(ary):    return ary
175    return PyArray_Copy(ary)
176
177cdef inline ndarray oarray_b(object ob, PetscInt* size, PetscBool** data):
178    cdef ndarray ary = oarray(ob, NPY_PETSC_BOOL)
179    cdef Py_ssize_t ssize = PyArray_SIZE(ary)
180    if size != NULL: size[0] = <PetscInt> ssize
181    if data != NULL: data[0] = <PetscBool*> PyArray_DATA(ary)
182    if ssize == 0 and data != NULL: data[0] = NULL
183    return ary
184
185cdef inline ndarray oarray_i(object ob, PetscInt* size, PetscInt** data):
186    cdef ndarray ary = oarray(ob, NPY_PETSC_INT)
187    cdef Py_ssize_t ssize = PyArray_SIZE(ary)
188    if size != NULL: size[0] = <PetscInt> ssize
189    if data != NULL: data[0] = <PetscInt*> PyArray_DATA(ary)
190    if ssize == 0 and data != NULL: data[0] = NULL
191    return ary
192
193cdef inline ndarray oarray_r(object ob, PetscInt* size, PetscReal** data):
194    cdef ndarray ary = oarray(ob, NPY_PETSC_REAL)
195    cdef Py_ssize_t ssize = PyArray_SIZE(ary)
196    if size != NULL: size[0] = <PetscInt> ssize
197    if data != NULL: data[0] = <PetscReal*> PyArray_DATA(ary)
198    if ssize == 0 and data != NULL: data[0] = NULL
199    return ary
200
201cdef inline ndarray oarray_s(object ob, PetscInt* size, PetscScalar** data):
202    cdef ndarray ary = oarray(ob, NPY_PETSC_SCALAR)
203    cdef Py_ssize_t ssize = PyArray_SIZE(ary)
204    if size != NULL: size[0] = <PetscInt> ssize
205    if data != NULL: data[0] = <PetscScalar*> PyArray_DATA(ary)
206    if ssize == 0 and data != NULL: data[0] = NULL
207    return ary
208
209cdef inline ndarray oarray_p(object ob, PetscInt* size, void** data):
210    cdef ndarray ary = oarray(ob, NPY_INTP)
211    cdef Py_ssize_t ssize = PyArray_SIZE(ary)
212    if size != NULL: size[0] = <PetscInt> ssize
213    if data != NULL: data[0] = <void*> PyArray_DATA(ary)
214    if ssize == 0 and data != NULL: data[0] = NULL
215    return ary
216
217# --------------------------------------------------------------------
218
219cdef inline ndarray ocarray_s(object ob, PetscInt* size, PetscScalar** data):
220    cdef ndarray ary = PyArray_FROM_OTF(
221        ob, NPY_PETSC_SCALAR, NPY_ARRAY_CARRAY|NPY_ARRAY_NOTSWAPPED)
222    cdef Py_ssize_t ssize = PyArray_SIZE(ary)
223    if size != NULL: size[0] = <PetscInt> ssize
224    if data != NULL: data[0] = <PetscScalar*> PyArray_DATA(ary)
225    return ary
226
227cdef inline ndarray ofarray_s(object ob, PetscInt* size, PetscScalar** data):
228    cdef ndarray ary = PyArray_FROM_OTF(
229        ob, NPY_PETSC_SCALAR, NPY_ARRAY_FARRAY|NPY_ARRAY_NOTSWAPPED)
230    cdef Py_ssize_t ssize = PyArray_SIZE(ary)
231    if size != NULL: size[0] = <PetscInt> ssize
232    if data != NULL: data[0] = <PetscScalar*> PyArray_DATA(ary)
233    return ary
234
235# --------------------------------------------------------------------
236