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
21    enum: NPY_INTP
22    dtype   PyArray_DescrFromType(int)
23    object  PyArray_TypeObjectFromType(int)
24
25    enum: NPY_ARRAY_ALIGNED
26    enum: NPY_ARRAY_WRITEABLE
27    enum: NPY_ARRAY_NOTSWAPPED
28    enum: NPY_ARRAY_CARRAY
29    enum: NPY_ARRAY_FARRAY
30
31    ndarray PyArray_FROM_O(object)
32    ndarray PyArray_FROM_OT(object,int)
33    ndarray PyArray_FROM_OTF(object,int,int)
34
35    ndarray PyArray_Copy(ndarray)
36    ndarray PyArray_ArangeObj(object,object,object,dtype)
37    ndarray PyArray_EMPTY(int,npy_intp[],int,int)
38    ndarray PyArray_ZEROS(int,npy_intp[],int,int)
39
40    bint PyArray_ISCONTIGUOUS(ndarray)
41    bint PyArray_ISFORTRAN(ndarray)
42    ctypedef enum NPY_ORDER:
43        NPY_ANYORDER
44        NPY_CORDER
45        NPY_FORTRANORDER
46    ndarray PyArray_NewCopy(ndarray,NPY_ORDER)
47
48    ctypedef struct PyObject
49    ctypedef struct PyTypeObject
50    PyObject* PyArray_New(PyTypeObject*,int,npy_intp[],int,npy_intp[],void*,int,int,PyObject*)
51    PyObject* PyArray_SimpleNewFromData(int,npy_intp[],int,void*)
52
53
54cdef extern from "petsc4py/numpy.h":
55
56    enum: NPY_INT
57    enum: NPY_DOUBLE
58
59    enum: NPY_PETSC_INT
60    enum: NPY_PETSC_REAL
61    enum: NPY_PETSC_SCALAR
62    enum: NPY_PETSC_COMPLEX
63
64
65# --------------------------------------------------------------------
66
67cdef inline ndarray asarray(object ob):
68    return PyArray_FROM_O(ob)
69
70cdef inline ndarray arange(start, stop, stride):
71    cdef dtype descr = <dtype> PyArray_DescrFromType(NPY_PETSC_INT)
72    return PyArray_ArangeObj(start, stop, stride, descr)
73
74# --------------------------------------------------------------------
75
76cdef inline ndarray empty_i(PetscInt size):
77    cdef npy_intp s = <npy_intp> size
78    return PyArray_EMPTY(1, &s, NPY_PETSC_INT, 0)
79
80cdef inline ndarray empty_r(PetscInt size):
81    cdef npy_intp s = <npy_intp> size
82    return PyArray_EMPTY(1, &s, NPY_PETSC_REAL, 0)
83
84cdef inline ndarray empty_s(PetscInt size):
85    cdef npy_intp s = <npy_intp> size
86    return PyArray_EMPTY(1, &s, NPY_PETSC_SCALAR, 0)
87
88cdef inline ndarray empty_c(PetscInt size):
89    cdef npy_intp s = <npy_intp> size
90    return PyArray_EMPTY(1, &s, NPY_PETSC_COMPLEX, 0)
91
92cdef inline ndarray empty_p(PetscInt size):
93    cdef npy_intp s = <npy_intp> size
94    return PyArray_EMPTY(1, &s, NPY_INTP, 0)
95
96# --------------------------------------------------------------------
97
98cdef inline ndarray array_i(PetscInt size, const PetscInt* data):
99    cdef npy_intp s = <npy_intp> size
100    cdef ndarray ary = PyArray_EMPTY(1, &s, NPY_PETSC_INT, 0)
101    if data != NULL:
102        memcpy(PyArray_DATA(ary), data, <size_t>size*sizeof(PetscInt))
103    return ary
104
105cdef inline ndarray array_r(PetscInt size, const PetscReal* data):
106    cdef npy_intp s = <npy_intp> size
107    cdef ndarray ary = PyArray_EMPTY(1, &s, NPY_PETSC_REAL, 0)
108    if data != NULL:
109        memcpy(PyArray_DATA(ary), data, <size_t>size*sizeof(PetscReal))
110    return ary
111
112cdef inline ndarray array_s(PetscInt size, const PetscScalar* data):
113    cdef npy_intp s = <npy_intp> size
114    cdef ndarray ary = PyArray_EMPTY(1, &s, NPY_PETSC_SCALAR, 0)
115    if data != NULL:
116        memcpy(PyArray_DATA(ary), data, <size_t>size*sizeof(PetscScalar))
117    return ary
118
119# --------------------------------------------------------------------
120
121cdef inline ndarray iarray(object ob, int typenum):
122    cdef ndarray ary = PyArray_FROM_OTF(
123        ob, typenum, NPY_ARRAY_ALIGNED|NPY_ARRAY_NOTSWAPPED)
124    if PyArray_ISCONTIGUOUS(ary): return ary
125    if PyArray_ISFORTRAN(ary):    return ary
126    return PyArray_Copy(ary)
127
128cdef inline ndarray iarray_i(object ob, PetscInt* size, PetscInt** data):
129    cdef ndarray ary = iarray(ob, NPY_PETSC_INT)
130    if size != NULL: size[0] = <PetscInt>  PyArray_SIZE(ary)
131    if data != NULL: data[0] = <PetscInt*> PyArray_DATA(ary)
132    return ary
133
134cdef inline ndarray iarray_r(object ob, PetscInt* size, PetscReal** data):
135    cdef ndarray ary = iarray(ob, NPY_PETSC_REAL)
136    if size != NULL: size[0] = <PetscInt>   PyArray_SIZE(ary)
137    if data != NULL: data[0] = <PetscReal*> PyArray_DATA(ary)
138    return ary
139
140cdef inline ndarray iarray_s(object ob, PetscInt* size, PetscScalar** data):
141    cdef ndarray ary = iarray(ob, NPY_PETSC_SCALAR)
142    if size != NULL: size[0] = <PetscInt>     PyArray_SIZE(ary)
143    if data != NULL: data[0] = <PetscScalar*> PyArray_DATA(ary)
144    return ary
145
146# --------------------------------------------------------------------
147
148cdef inline ndarray oarray(object ob, int typenum):
149    cdef ndarray ary = PyArray_FROM_OTF(
150        ob, typenum, NPY_ARRAY_ALIGNED|NPY_ARRAY_WRITEABLE|NPY_ARRAY_NOTSWAPPED)
151    if PyArray_ISCONTIGUOUS(ary): return ary
152    if PyArray_ISFORTRAN(ary):    return ary
153    return PyArray_Copy(ary)
154
155cdef inline ndarray oarray_i(object ob, PetscInt* size, PetscInt** data):
156    cdef ndarray ary = oarray(ob, NPY_PETSC_INT)
157    if size != NULL: size[0] = <PetscInt>  PyArray_SIZE(ary)
158    if data != NULL: data[0] = <PetscInt*> PyArray_DATA(ary)
159    return ary
160
161cdef inline ndarray oarray_r(object ob, PetscInt* size, PetscReal** data):
162    cdef ndarray ary = oarray(ob, NPY_PETSC_REAL)
163    if size != NULL: size[0] = <PetscInt>   PyArray_SIZE(ary)
164    if data != NULL: data[0] = <PetscReal*> PyArray_DATA(ary)
165    return ary
166
167cdef inline ndarray oarray_s(object ob, PetscInt* size, PetscScalar** data):
168    cdef ndarray ary = oarray(ob, NPY_PETSC_SCALAR)
169    if size != NULL: size[0] = <PetscInt>     PyArray_SIZE(ary)
170    if data != NULL: data[0] = <PetscScalar*> PyArray_DATA(ary)
171    return ary
172
173cdef inline ndarray oarray_p(object ob, PetscInt* size, void** data):
174    cdef ndarray ary = oarray(ob, NPY_INTP)
175    if size != NULL: size[0] = <PetscInt> PyArray_SIZE(ary)
176    if data != NULL: data[0] = <void*>    PyArray_DATA(ary)
177    return ary
178
179# --------------------------------------------------------------------
180
181cdef inline ndarray ocarray_s(object ob, PetscInt* size, PetscScalar** data):
182    cdef ndarray ary = PyArray_FROM_OTF(
183        ob, NPY_PETSC_SCALAR, NPY_ARRAY_CARRAY|NPY_ARRAY_NOTSWAPPED)
184    if size != NULL: size[0] = <PetscInt>     PyArray_SIZE(ary)
185    if data != NULL: data[0] = <PetscScalar*> PyArray_DATA(ary)
186    return ary
187
188cdef inline ndarray ofarray_s(object ob, PetscInt* size, PetscScalar** data):
189    cdef ndarray ary = PyArray_FROM_OTF(
190        ob, NPY_PETSC_SCALAR, NPY_ARRAY_FARRAY|NPY_ARRAY_NOTSWAPPED)
191    if size != NULL: size[0] = <PetscInt>     PyArray_SIZE(ary)
192    if data != NULL: data[0] = <PetscScalar*> PyArray_DATA(ary)
193    return ary
194
195# --------------------------------------------------------------------
196