1 // Copyright 2010-2019 Keith Goodman
2 // Copyright 2019 Bottleneck Developers
3 #ifndef ITERATORS_H_
4 #define ITERATORS_H_
5
6 #include <numpy/arrayobject.h>
7
8 /*
9 Bottleneck iterators are based on ideas from NumPy's PyArray_IterAllButAxis
10 and PyArray_ITER_NEXT.
11 */
12
13 /* one input array ------------------------------------------------------- */
14
15 /* these iterators are used mainly by reduce functions such as nansum */
16
17 struct _iter {
18 int ndim_m2; /* ndim - 2 */
19 int axis; /* axis to not iterate over */
20 Py_ssize_t length; /* a.shape[axis] */
21 Py_ssize_t astride; /* a.strides[axis] */
22 npy_intp stride; /* element-level stride to take in the array */
23 npy_intp i; /* integer used by some macros */
24 npy_intp its; /* number of iterations completed */
25 npy_intp nits; /* number of iterations iterator plans to make */
26 npy_intp indices[NPY_MAXDIMS]; /* current location of iterator */
27 npy_intp astrides[NPY_MAXDIMS]; /* a.strides, a.strides[axis] removed */
28 npy_intp shape[NPY_MAXDIMS]; /* a.shape, a.shape[axis] removed */
29 char *pa; /* pointer to data corresponding to indices */
30 PyArrayObject *a_ravel; /* NULL or pointer to ravelled input array */
31 };
32 typedef struct _iter iter;
33
34 static inline void
init_iter_one(iter * it,PyArrayObject * a,int axis)35 init_iter_one(iter *it, PyArrayObject *a, int axis)
36 {
37 int i, j = 0;
38 const int ndim = PyArray_NDIM(a);
39 const npy_intp *shape = PyArray_SHAPE(a);
40 const npy_intp *strides = PyArray_STRIDES(a);
41 const npy_intp item_size = PyArray_ITEMSIZE(a);
42
43 it->axis = axis;
44 it->its = 0;
45 it->nits = 1;
46 it->pa = PyArray_BYTES(a);
47
48 it->ndim_m2 = -1;
49 it->length = 1;
50 it->astride = 0;
51
52 if (ndim != 0) {
53 it->ndim_m2 = ndim - 2;
54 for (i = 0; i < ndim; i++) {
55 if (i == axis) {
56 it->astride = strides[i];
57 it->length = shape[i];
58 } else {
59 it->indices[j] = 0;
60 it->astrides[j] = strides[i];
61 it->shape[j] = shape[i];
62 it->nits *= shape[i];
63 j++;
64 }
65 }
66 }
67 it->stride = it->astride / item_size;
68 }
69
70 /*
71 * If both ravel != 0 and it.a_ravel != NULL then you are responsible for
72 * calling Py_DECREF(it.a_ravel) after you are done with the iterator.
73 * See nanargmin for an example.
74 */
75 static inline void
init_iter_all(iter * it,PyArrayObject * a,int ravel,int anyorder)76 init_iter_all(iter *it, PyArrayObject *a, int ravel, int anyorder)
77 {
78 int i, j = 0;
79 const int ndim = PyArray_NDIM(a);
80 const npy_intp *shape = PyArray_SHAPE(a);
81 const npy_intp *strides = PyArray_STRIDES(a);
82 const npy_intp item_size = PyArray_ITEMSIZE(a);
83
84 it->axis = 0;
85 it->its = 0;
86 it->nits = 1;
87 it->a_ravel = NULL;
88
89 /* The fix for relaxed strides checking in numpy and the fix for
90 * issue #183 has left this if..else tree in need of a refactor from the
91 * the ground up */
92 if (ndim == 1) {
93 it->ndim_m2 = -1;
94 it->length = shape[0];
95 it->astride = strides[0];
96 } else if (ndim == 0) {
97 it->ndim_m2 = -1;
98 it->length = 1;
99 it->astride = 0;
100 } else if (C_CONTIGUOUS(a) && !F_CONTIGUOUS(a)) {
101 /* The &&! in the next two else ifs is to deal with relaxed
102 * stride checking introduced in numpy 1.12.0; see gh #161 */
103 it->ndim_m2 = -1;
104 it->axis = ndim - 1;
105 it->length = PyArray_SIZE(a);
106 it->astride = 0;
107 for (i=ndim-1; i > -1; i--) {
108 /* protect against length zero strides such as in
109 * np.ones((2, 2))[..., np.newaxis] */
110 if (strides[i] == 0) {
111 continue;
112 }
113 it->astride = strides[i];
114 break;
115 }
116 } else if (F_CONTIGUOUS(a) && !C_CONTIGUOUS(a)) {
117 if (anyorder || !ravel) {
118 it->ndim_m2 = -1;
119 it->length = PyArray_SIZE(a);
120 it->astride = 0;
121 for (i=0; i < ndim; i++) {
122 /* protect against length zero strides such as in
123 * np.ones((2, 2), order='F')[np.newaxis, ...] */
124 if (strides[i] == 0) {
125 continue;
126 }
127 it->astride = strides[i];
128 break;
129 }
130 } else {
131 it->ndim_m2 = -1;
132 if (anyorder) {
133 a = (PyArrayObject *)PyArray_Ravel(a, NPY_ANYORDER);
134 } else {
135 a = (PyArrayObject *)PyArray_Ravel(a, NPY_CORDER);
136 }
137 it->a_ravel = a;
138 it->length = PyArray_DIM(a, 0);
139 it->astride = PyArray_STRIDE(a, 0);
140 }
141 } else if (ravel) {
142 it->ndim_m2 = -1;
143 if (anyorder) {
144 a = (PyArrayObject *)PyArray_Ravel(a, NPY_ANYORDER);
145 } else {
146 a = (PyArrayObject *)PyArray_Ravel(a, NPY_CORDER);
147 }
148 it->a_ravel = a;
149 it->length = PyArray_DIM(a, 0);
150 it->astride = PyArray_STRIDE(a, 0);
151 } else {
152 it->ndim_m2 = ndim - 2;
153 it->astride = strides[0];
154 for (i = 1; i < ndim; i++) {
155 if (strides[i] < it->astride) {
156 it->astride = strides[i];
157 it->axis = i;
158 }
159 }
160 it->length = shape[it->axis];
161 for (i = 0; i < ndim; i++) {
162 if (i != it->axis) {
163 it->indices[j] = 0;
164 it->astrides[j] = strides[i];
165 it->shape[j] = shape[i];
166 it->nits *= shape[i];
167 j++;
168 }
169 }
170 }
171
172 it->stride = it->astride / item_size;
173 it->pa = PyArray_BYTES(a);
174 }
175
176 #define NEXT \
177 for (it.i = it.ndim_m2; it.i > -1; it.i--) { \
178 if (it.indices[it.i] < it.shape[it.i] - 1) { \
179 it.pa += it.astrides[it.i]; \
180 it.indices[it.i]++; \
181 break; \
182 } \
183 it.pa -= it.indices[it.i] * it.astrides[it.i]; \
184 it.indices[it.i] = 0; \
185 } \
186 it.its++;
187
188 /* two input arrays ------------------------------------------------------ */
189
190 /* this iterator is used mainly by moving window functions such as move_sum */
191
192 struct _iter2 {
193 int ndim_m2;
194 int axis;
195 Py_ssize_t length;
196 Py_ssize_t astride;
197 Py_ssize_t ystride;
198 npy_intp i;
199 npy_intp its;
200 npy_intp nits;
201 npy_intp indices[NPY_MAXDIMS];
202 npy_intp astrides[NPY_MAXDIMS];
203 npy_intp ystrides[NPY_MAXDIMS];
204 npy_intp shape[NPY_MAXDIMS];
205 char *pa;
206 char *py;
207 };
208 typedef struct _iter2 iter2;
209
210 static inline void
init_iter2(iter2 * it,PyArrayObject * a,PyObject * y,int axis)211 init_iter2(iter2 *it, PyArrayObject *a, PyObject *y, int axis)
212 {
213 int i, j = 0;
214 const int ndim = PyArray_NDIM(a);
215 const npy_intp *shape = PyArray_SHAPE(a);
216 const npy_intp *astrides = PyArray_STRIDES(a);
217 const npy_intp *ystrides = PyArray_STRIDES((PyArrayObject *)y);
218
219 /* to avoid compiler warning of uninitialized variables */
220 it->length = 0;
221 it->astride = 0;
222 it->ystride = 0;
223
224 it->ndim_m2 = ndim - 2;
225 it->axis = axis;
226 it->its = 0;
227 it->nits = 1;
228 it->pa = PyArray_BYTES(a);
229 it->py = PyArray_BYTES((PyArrayObject *)y);
230
231 for (i = 0; i < ndim; i++) {
232 if (i == axis) {
233 it->astride = astrides[i];
234 it->ystride = ystrides[i];
235 it->length = shape[i];
236 } else {
237 it->indices[j] = 0;
238 it->astrides[j] = astrides[i];
239 it->ystrides[j] = ystrides[i];
240 it->shape[j] = shape[i];
241 it->nits *= shape[i];
242 j++;
243 }
244 }
245 }
246
247 #define NEXT2 \
248 for (it.i = it.ndim_m2; it.i > -1; it.i--) { \
249 if (it.indices[it.i] < it.shape[it.i] - 1) { \
250 it.pa += it.astrides[it.i]; \
251 it.py += it.ystrides[it.i]; \
252 it.indices[it.i]++; \
253 break; \
254 } \
255 it.pa -= it.indices[it.i] * it.astrides[it.i]; \
256 it.py -= it.indices[it.i] * it.ystrides[it.i]; \
257 it.indices[it.i] = 0; \
258 } \
259 it.its++;
260
261 /* three input arrays ---------------------------------------------------- */
262
263 /* this iterator is used mainly by rankdata and nanrankdata */
264
265 struct _iter3 {
266 int ndim_m2;
267 int axis;
268 Py_ssize_t length;
269 Py_ssize_t astride;
270 Py_ssize_t ystride;
271 Py_ssize_t zstride;
272 npy_intp i;
273 npy_intp its;
274 npy_intp nits;
275 npy_intp indices[NPY_MAXDIMS];
276 npy_intp astrides[NPY_MAXDIMS];
277 npy_intp ystrides[NPY_MAXDIMS];
278 npy_intp zstrides[NPY_MAXDIMS];
279 npy_intp shape[NPY_MAXDIMS];
280 char *pa;
281 char *py;
282 char *pz;
283 };
284 typedef struct _iter3 iter3;
285
286 static inline void
init_iter3(iter3 * it,PyArrayObject * a,PyObject * y,PyObject * z,int axis)287 init_iter3(iter3 *it, PyArrayObject *a, PyObject *y, PyObject *z, int axis)
288 {
289 int i, j = 0;
290 const int ndim = PyArray_NDIM(a);
291 const npy_intp *shape = PyArray_SHAPE(a);
292 const npy_intp *astrides = PyArray_STRIDES(a);
293 const npy_intp *ystrides = PyArray_STRIDES((PyArrayObject *)y);
294 const npy_intp *zstrides = PyArray_STRIDES((PyArrayObject *)z);
295
296 /* to avoid compiler warning of uninitialized variables */
297 it->length = 0;
298 it->astride = 0;
299 it->ystride = 0;
300 it->zstride = 0;
301
302 it->ndim_m2 = ndim - 2;
303 it->axis = axis;
304 it->its = 0;
305 it->nits = 1;
306 it->pa = PyArray_BYTES(a);
307 it->py = PyArray_BYTES((PyArrayObject *)y);
308 it->pz = PyArray_BYTES((PyArrayObject *)z);
309
310 for (i = 0; i < ndim; i++) {
311 if (i == axis) {
312 it->astride = astrides[i];
313 it->ystride = ystrides[i];
314 it->zstride = zstrides[i];
315 it->length = shape[i];
316 } else {
317 it->indices[j] = 0;
318 it->astrides[j] = astrides[i];
319 it->ystrides[j] = ystrides[i];
320 it->zstrides[j] = zstrides[i];
321 it->shape[j] = shape[i];
322 it->nits *= shape[i];
323 j++;
324 }
325 }
326 }
327
328 #define NEXT3 \
329 for (it.i = it.ndim_m2; it.i > -1; it.i--) { \
330 if (it.indices[it.i] < it.shape[it.i] - 1) { \
331 it.pa += it.astrides[it.i]; \
332 it.py += it.ystrides[it.i]; \
333 it.pz += it.zstrides[it.i]; \
334 it.indices[it.i]++; \
335 break; \
336 } \
337 it.pa -= it.indices[it.i] * it.astrides[it.i]; \
338 it.py -= it.indices[it.i] * it.ystrides[it.i]; \
339 it.pz -= it.indices[it.i] * it.zstrides[it.i]; \
340 it.indices[it.i] = 0; \
341 } \
342 it.its++;
343
344 /* macros used with iterators -------------------------------------------- */
345
346 /* most of these macros assume iterator is named `it` */
347
348 #define NDIM it.ndim_m2 + 2
349 #define SHAPE it.shape
350 #define SIZE it.nits * it.length
351 #define LENGTH it.length
352 #define INDEX it.i
353
354 #define WHILE while (it.its < it.nits)
355 #define WHILE0 it.i = 0; while (it.i < min_count - 1)
356 #define WHILE1 while (it.i < window)
357 #define WHILE2 while (it.i < it.length)
358
359 #define FOR for (it.i = 0; it.i < it.length; it.i++)
360 #define FOR_REVERSE for (it.i = it.length - 1; it.i > -1; it.i--)
361
362 #define RESET it.its = 0;
363
364 #define PA(dtype) (npy_##dtype *)(it.pa)
365
366 #define A0(dtype) *(npy_##dtype *)(it.pa)
367 #define AI(dtype) *(npy_##dtype *)(it.pa + it.i * it.astride)
368 #define AX(dtype, x) *(npy_##dtype *)(it.pa + (x) * it.astride)
369 #define AOLD(dtype) *(npy_##dtype *)(it.pa + (it.i - window) * it.astride)
370
371 #define SI(pa) pa[it.i * it.stride]
372
373 #define YPP *py++
374 #define YI(dtype) *(npy_##dtype *)(it.py + it.i++ * it.ystride)
375 #define YX(dtype, x) *(npy_##dtype *)(it.py + (x) * it.ystride)
376
377 #define ZX(dtype, x) *(npy_##dtype *)(it.pz + (x) * it.zstride)
378
379 #define FILL_Y(value) \
380 npy_intp _i; \
381 npy_intp size = PyArray_SIZE((PyArrayObject *)y); \
382 for (_i = 0; _i < size; _i++) { \
383 YPP = value; \
384 }
385
386 #endif // ITERATORS_H_
387