1 // Copyright 2010-2019 Keith Goodman
2 // Copyright 2019 Bottleneck Developers
3 #include "bottleneck.h"
4 #include "iterators.h"
5
6 /* function signatures --------------------------------------------------- */
7
8 /* low-level functions such as move_sum_float64 */
9 #define NRA(name, dtype) \
10 static PyObject * \
11 name##_##dtype(PyArrayObject *a, int axis, int n)
12
13 /* top-level functions such as move_sum */
14 #define NRA_MAIN(name, parse) \
15 static PyObject * \
16 name(PyObject *self, PyObject *args, PyObject *kwds) \
17 { \
18 return nonreducer_axis(#name, \
19 args, \
20 kwds, \
21 name##_float64, \
22 name##_float32, \
23 name##_int64, \
24 name##_int32, \
25 parse); \
26 }
27
28 /* typedefs and prototypes ----------------------------------------------- */
29
30 /* how should input be parsed? */
31 typedef enum {PARSE_PARTITION, PARSE_RANKDATA, PARSE_PUSH} parse_type;
32
33 /* function pointer for functions passed to nonreducer_axis */
34 typedef PyObject *(*nra_t)(PyArrayObject *, int, int);
35
36 static PyObject *
37 nonreducer_axis(char *name,
38 PyObject *args,
39 PyObject *kwds,
40 nra_t,
41 nra_t,
42 nra_t,
43 nra_t,
44 parse_type);
45
46 /* partition ------------------------------------------------------------- */
47
48 #define B(dtype, i) AX(dtype, i) /* used by PARTITION */
49
50 /* dtype = [['float64'], ['float32'], ['int64'], ['int32']] */
NRA(partition,DTYPE0)51 NRA(partition, DTYPE0) {
52 npy_intp i;
53 npy_intp j, l, r, k;
54 iter it;
55
56 a = (PyArrayObject *)PyArray_NewCopy(a, NPY_ANYORDER);
57 init_iter_one(&it, a, axis);
58
59 if (LENGTH == 0) return (PyObject *)a;
60 if (n < 0 || n > LENGTH - 1) {
61 PyErr_Format(PyExc_ValueError,
62 "`n` (=%d) must be between 0 and %zd, inclusive.",
63 n, LENGTH - 1);
64 return NULL;
65 }
66
67 BN_BEGIN_ALLOW_THREADS
68 k = n;
69 WHILE {
70 l = 0;
71 r = LENGTH - 1;
72 PARTITION(DTYPE0)
73 NEXT
74 }
75 BN_END_ALLOW_THREADS
76
77 return (PyObject *)a;
78 }
79 /* dtype end */
80
NRA_MAIN(partition,PARSE_PARTITION)81 NRA_MAIN(partition, PARSE_PARTITION)
82
83
84 /* argpartition ----------------------------------------------------------- */
85
86 #define BUFFER_NEW(dtype) dtype *B = malloc(LENGTH * sizeof(dtype));
87 #define BUFFER_DELETE free(B);
88
89 #define ARGWIRTH(dtype0, dtype1) \
90 x = B[k]; \
91 i = l; \
92 j = r; \
93 do { \
94 while (B[i] < x) i++; \
95 while (x < B[j]) j--; \
96 if (i <= j) { \
97 npy_##dtype0 atmp = B[i]; \
98 B[i] = B[j]; \
99 B[j] = atmp; \
100 ytmp = YX(dtype1, i); \
101 YX(dtype1, i) = YX(dtype1, j); \
102 YX(dtype1, j) = ytmp; \
103 i++; \
104 j--; \
105 } \
106 } while (i <= j); \
107 if (j < k) l = i; \
108 if (k < i) r = j;
109
110 #define ARGPARTITION(dtype0, dtype1) \
111 while (l < r) { \
112 npy_##dtype0 x; \
113 npy_##dtype0 al = B[l]; \
114 npy_##dtype0 ak = B[k]; \
115 npy_##dtype0 ar = B[r]; \
116 npy_##dtype1 ytmp; \
117 if (al > ak) { \
118 if (ak < ar) { \
119 if (al < ar) { \
120 B[k] = al; \
121 B[l] = ak; \
122 ytmp = YX(dtype1, k); \
123 YX(dtype1, k) = YX(dtype1, l); \
124 YX(dtype1, l) = ytmp; \
125 } else { \
126 B[k] = ar; \
127 B[r] = ak; \
128 ytmp = YX(dtype1, k); \
129 YX(dtype1, k) = YX(dtype1, r); \
130 YX(dtype1, r) = ytmp; \
131 } \
132 } \
133 } else { \
134 if (ak > ar) { \
135 if (al > ar) { \
136 B[k] = al; \
137 B[l] = ak; \
138 ytmp = YX(dtype1, k); \
139 YX(dtype1, k) = YX(dtype1, l); \
140 YX(dtype1, l) = ytmp; \
141 } else { \
142 B[k] = ar; \
143 B[r] = ak; \
144 ytmp = YX(dtype1, k); \
145 YX(dtype1, k) = YX(dtype1, r); \
146 YX(dtype1, r) = ytmp; \
147 } \
148 } \
149 } \
150 ARGWIRTH(dtype0, dtype1) \
151 }
152
153 #define ARGPARTSORT(dtype0, dtype1) \
154 for (i = 0; i < LENGTH; i++) { \
155 B[i] = AX(dtype0, i); \
156 YX(dtype1, i) = i; \
157 } \
158 l = 0; \
159 r = LENGTH - 1; \
160 ARGPARTITION(dtype0, dtype1)
161
162 /* dtype = [['float64', 'intp'], ['float32', 'intp'],
163 ['int64', 'intp'], ['int32', 'intp']] */
164 NRA(argpartition, DTYPE0) {
165 npy_intp i;
166 PyObject *y = PyArray_EMPTY(PyArray_NDIM(a), PyArray_SHAPE(a),
167 NPY_DTYPE1, 0);
168 iter2 it;
169 init_iter2(&it, a, y, axis);
170 if (LENGTH == 0) return y;
171 if (n < 0 || n > LENGTH - 1) {
172 PyErr_Format(PyExc_ValueError,
173 "`n` (=%d) must be between 0 and %zd, inclusive.",
174 n, LENGTH - 1);
175 return NULL;
176 }
177 BN_BEGIN_ALLOW_THREADS
178 BUFFER_NEW(npy_DTYPE0)
179 npy_intp j, l, r, k;
180 k = n;
181 WHILE {
182 l = 0;
183 r = LENGTH - 1;
184 ARGPARTSORT(DTYPE0, DTYPE1)
185 NEXT2
186 }
187 BUFFER_DELETE
188 BN_END_ALLOW_THREADS
189 return y;
190 }
191 /* dtype end */
192
NRA_MAIN(argpartition,PARSE_PARTITION)193 NRA_MAIN(argpartition, PARSE_PARTITION)
194
195
196 /* rankdata -------------------------------------------------------------- */
197
198 /* dtype = [['float64', 'float64', 'intp'], ['float32', 'float64', 'intp'],
199 ['int64', 'float64', 'intp'], ['int32', 'float64', 'intp']] */
200 NRA(rankdata, DTYPE0) {
201 Py_ssize_t j=0, k, idx, dupcount=0, i;
202 npy_DTYPE1 old, new, averank, sumranks = 0;
203
204 PyObject *z = PyArray_ArgSort(a, axis, NPY_QUICKSORT);
205 PyObject *y = PyArray_EMPTY(PyArray_NDIM(a),
206 PyArray_SHAPE(a), NPY_DTYPE1, 0);
207
208 iter3 it;
209 init_iter3(&it, a, y, z, axis);
210
211 BN_BEGIN_ALLOW_THREADS
212 if (LENGTH == 0) {
213 Py_ssize_t size = PyArray_SIZE((PyArrayObject *)y);
214 npy_DTYPE1 *py = (npy_DTYPE1 *)PyArray_DATA(a);
215 for (i = 0; i < size; i++) YPP = BN_NAN;
216 } else {
217 WHILE {
218 idx = ZX(DTYPE2, 0);
219 old = AX(DTYPE0, idx);
220 sumranks = 0;
221 dupcount = 0;
222 for (i = 0; i < LENGTH - 1; i++) {
223 sumranks += i;
224 dupcount++;
225 k = i + 1;
226 idx = ZX(DTYPE2, k);
227 new = AX(DTYPE0, idx);
228 if (old != new) {
229 averank = sumranks / dupcount + 1;
230 for (j = k - dupcount; j < k; j++) {
231 idx = ZX(DTYPE2, j);
232 YX(DTYPE1, idx) = averank;
233 }
234 sumranks = 0;
235 dupcount = 0;
236 }
237 old = new;
238 }
239 sumranks += (LENGTH - 1);
240 dupcount++;
241 averank = sumranks / dupcount + 1;
242 for (j = LENGTH - dupcount; j < LENGTH; j++) {
243 idx = ZX(DTYPE2, j);
244 YX(DTYPE1, idx) = averank;
245 }
246 NEXT3
247 }
248 }
249 BN_END_ALLOW_THREADS
250
251 Py_DECREF(z);
252 return y;
253 }
254 /* dtype end */
255
NRA_MAIN(rankdata,PARSE_RANKDATA)256 NRA_MAIN(rankdata, PARSE_RANKDATA)
257
258
259 /* nanrankdata ----------------------------------------------------------- */
260
261 /* dtype = [['float64', 'float64', 'intp'], ['float32', 'float64', 'intp']] */
262 NRA(nanrankdata, DTYPE0) {
263 Py_ssize_t j=0, k, idx, dupcount=0, i;
264 npy_DTYPE1 old, new, averank, sumranks = 0;
265
266 PyObject *z = PyArray_ArgSort(a, axis, NPY_QUICKSORT);
267 PyObject *y = PyArray_EMPTY(PyArray_NDIM(a),
268 PyArray_SHAPE(a), NPY_DTYPE1, 0);
269
270 iter3 it;
271 init_iter3(&it, a, y, z, axis);
272
273 BN_BEGIN_ALLOW_THREADS
274 if (LENGTH == 0) {
275 Py_ssize_t size = PyArray_SIZE((PyArrayObject *)y);
276 npy_DTYPE1 *py = (npy_DTYPE1 *)PyArray_DATA(a);
277 for (i = 0; i < size; i++) YPP = BN_NAN;
278 } else {
279 WHILE {
280 idx = ZX(DTYPE2, 0);
281 old = AX(DTYPE0, idx);
282 sumranks = 0;
283 dupcount = 0;
284 for (i = 0; i < LENGTH - 1; i++) {
285 sumranks += i;
286 dupcount++;
287 k = i + 1;
288 idx = ZX(DTYPE2, k);
289 new = AX(DTYPE0, idx);
290 if (old != new) {
291 if (old == old) {
292 averank = sumranks / dupcount + 1;
293 for (j = k - dupcount; j < k; j++) {
294 idx = ZX(DTYPE2, j);
295 YX(DTYPE1, idx) = averank;
296 }
297 } else {
298 idx = ZX(DTYPE2, i);
299 YX(DTYPE1, idx) = BN_NAN;
300 }
301 sumranks = 0;
302 dupcount = 0;
303 }
304 old = new;
305 }
306 sumranks += (LENGTH - 1);
307 dupcount++;
308 averank = sumranks / dupcount + 1;
309 if (old == old) {
310 for (j = LENGTH - dupcount; j < LENGTH; j++) {
311 idx = ZX(DTYPE2, j);
312 YX(DTYPE1, idx) = averank;
313 }
314 } else {
315 idx = ZX(DTYPE2, LENGTH - 1);
316 YX(DTYPE1, idx) = BN_NAN;
317 }
318 NEXT3
319 }
320 }
321 BN_END_ALLOW_THREADS
322
323 Py_DECREF(z);
324 return y;
325 }
326 /* dtype end */
327
328 static PyObject *
nanrankdata(PyObject * self,PyObject * args,PyObject * kwds)329 nanrankdata(PyObject *self, PyObject *args, PyObject *kwds) {
330 return nonreducer_axis("nanrankdata",
331 args,
332 kwds,
333 nanrankdata_float64,
334 nanrankdata_float32,
335 rankdata_int64,
336 rankdata_int32,
337 PARSE_RANKDATA);
338 }
339
340
341 /* push ------------------------------------------------------------------ */
342
343 /* dtype = [['float64'], ['float32']] */
NRA(push,DTYPE0)344 NRA(push, DTYPE0) {
345 npy_intp index;
346 npy_DTYPE0 ai, ai_last, n_float;
347 PyObject *y = PyArray_Copy(a);
348 iter it;
349 init_iter_one(&it, (PyArrayObject *)y, axis);
350 if (LENGTH == 0 || NDIM == 0) {
351 return y;
352 }
353 n_float = n < 0 ? BN_INFINITY : (npy_DTYPE0)n;
354 BN_BEGIN_ALLOW_THREADS
355 WHILE {
356 index = 0;
357 ai_last = BN_NAN;
358 FOR {
359 ai = AI(DTYPE0);
360 if (ai == ai) {
361 ai_last = ai;
362 index = INDEX;
363 } else {
364 if (INDEX - index <= n_float) {
365 AI(DTYPE0) = ai_last;
366 }
367 }
368 }
369 NEXT
370 }
371 BN_END_ALLOW_THREADS
372 return y;
373 }
374 /* dtype end */
375
376 /* dtype = [['int64'], ['int32']] */
NRA(push,DTYPE0)377 NRA(push, DTYPE0) {
378 PyObject *y = PyArray_Copy(a);
379 return y;
380 }
381 /* dtype end */
382
383 NRA_MAIN(push, PARSE_PUSH)
384
385
386 /* python strings -------------------------------------------------------- */
387
388 PyObject *pystr_a = NULL;
389 PyObject *pystr_n = NULL;
390 PyObject *pystr_kth = NULL;
391 PyObject *pystr_axis = NULL;
392
393 static int
intern_strings(void)394 intern_strings(void) {
395 pystr_a = PyString_InternFromString("a");
396 pystr_n = PyString_InternFromString("n");
397 pystr_kth = PyString_InternFromString("kth");
398 pystr_axis = PyString_InternFromString("axis");
399 return pystr_a && pystr_n && pystr_axis;
400 }
401
402 /* nonreducer_axis ------------------------------------------------------- */
403
404 static inline int
parse_partition(PyObject * args,PyObject * kwds,PyObject ** a,PyObject ** n,PyObject ** axis)405 parse_partition(PyObject *args,
406 PyObject *kwds,
407 PyObject **a,
408 PyObject **n,
409 PyObject **axis) {
410 const Py_ssize_t nargs = PyTuple_GET_SIZE(args);
411 const Py_ssize_t nkwds = kwds == NULL ? 0 : PyDict_Size(kwds);
412 if (nkwds) {
413 int nkwds_found = 0;
414 PyObject *tmp;
415 switch (nargs) {
416 case 2: *n = PyTuple_GET_ITEM(args, 1);
417 case 1: *a = PyTuple_GET_ITEM(args, 0);
418 case 0: break;
419 default:
420 TYPE_ERR("wrong number of arguments");
421 return 0;
422 }
423 switch (nargs) {
424 case 0:
425 *a = PyDict_GetItem(kwds, pystr_a);
426 if (*a == NULL) {
427 TYPE_ERR("Cannot find `a` keyword input");
428 return 0;
429 }
430 nkwds_found += 1;
431 case 1:
432 *n = PyDict_GetItem(kwds, pystr_kth);
433 if (*n == NULL) {
434 TYPE_ERR("Cannot find `kth` keyword input");
435 return 0;
436 }
437 nkwds_found++;
438 case 2:
439 tmp = PyDict_GetItem(kwds, pystr_axis);
440 if (tmp != NULL) {
441 *axis = tmp;
442 nkwds_found++;
443 }
444 break;
445 default:
446 TYPE_ERR("wrong number of arguments");
447 return 0;
448 }
449 if (nkwds_found != nkwds) {
450 TYPE_ERR("wrong number of keyword arguments");
451 return 0;
452 }
453 if (nargs + nkwds_found > 3) {
454 TYPE_ERR("too many arguments");
455 return 0;
456 }
457 } else {
458 switch (nargs) {
459 case 3:
460 *axis = PyTuple_GET_ITEM(args, 2);
461 case 2:
462 *n = PyTuple_GET_ITEM(args, 1);
463 *a = PyTuple_GET_ITEM(args, 0);
464 break;
465 default:
466 TYPE_ERR("wrong number of arguments");
467 return 0;
468 }
469 }
470
471 return 1;
472
473 }
474
475 static inline int
parse_rankdata(PyObject * args,PyObject * kwds,PyObject ** a,PyObject ** axis)476 parse_rankdata(PyObject *args,
477 PyObject *kwds,
478 PyObject **a,
479 PyObject **axis) {
480 const Py_ssize_t nargs = PyTuple_GET_SIZE(args);
481 const Py_ssize_t nkwds = kwds == NULL ? 0 : PyDict_Size(kwds);
482 if (nkwds) {
483 int nkwds_found = 0;
484 PyObject *tmp;
485 switch (nargs) {
486 case 1: *a = PyTuple_GET_ITEM(args, 0);
487 case 0: break;
488 default:
489 TYPE_ERR("wrong number of arguments");
490 return 0;
491 }
492 switch (nargs) {
493 case 0:
494 *a = PyDict_GetItem(kwds, pystr_a);
495 if (*a == NULL) {
496 TYPE_ERR("Cannot find `a` keyword input");
497 return 0;
498 }
499 nkwds_found += 1;
500 case 1:
501 tmp = PyDict_GetItem(kwds, pystr_axis);
502 if (tmp != NULL) {
503 *axis = tmp;
504 nkwds_found++;
505 }
506 break;
507 default:
508 TYPE_ERR("wrong number of arguments");
509 return 0;
510 }
511 if (nkwds_found != nkwds) {
512 TYPE_ERR("wrong number of keyword arguments");
513 return 0;
514 }
515 if (nargs + nkwds_found > 2) {
516 TYPE_ERR("too many arguments");
517 return 0;
518 }
519 } else {
520 switch (nargs) {
521 case 2:
522 *axis = PyTuple_GET_ITEM(args, 1);
523 case 1:
524 *a = PyTuple_GET_ITEM(args, 0);
525 break;
526 default:
527 TYPE_ERR("wrong number of arguments");
528 return 0;
529 }
530 }
531
532 return 1;
533
534 }
535
536 static inline int
parse_push(PyObject * args,PyObject * kwds,PyObject ** a,PyObject ** n,PyObject ** axis)537 parse_push(PyObject *args,
538 PyObject *kwds,
539 PyObject **a,
540 PyObject **n,
541 PyObject **axis) {
542 const Py_ssize_t nargs = PyTuple_GET_SIZE(args);
543 const Py_ssize_t nkwds = kwds == NULL ? 0 : PyDict_Size(kwds);
544 if (nkwds) {
545 int nkwds_found = 0;
546 PyObject *tmp;
547 switch (nargs) {
548 case 2: *n = PyTuple_GET_ITEM(args, 1);
549 case 1: *a = PyTuple_GET_ITEM(args, 0);
550 case 0: break;
551 default:
552 TYPE_ERR("wrong number of arguments");
553 return 0;
554 }
555 switch (nargs) {
556 case 0:
557 *a = PyDict_GetItem(kwds, pystr_a);
558 if (*a == NULL) {
559 TYPE_ERR("Cannot find `a` keyword input");
560 return 0;
561 }
562 nkwds_found += 1;
563 case 1:
564 tmp = PyDict_GetItem(kwds, pystr_n);
565 if (tmp != NULL) {
566 *n = tmp;
567 nkwds_found++;
568 }
569 case 2:
570 tmp = PyDict_GetItem(kwds, pystr_axis);
571 if (tmp != NULL) {
572 *axis = tmp;
573 nkwds_found++;
574 }
575 break;
576 default:
577 TYPE_ERR("wrong number of arguments");
578 return 0;
579 }
580 if (nkwds_found != nkwds) {
581 TYPE_ERR("wrong number of keyword arguments");
582 return 0;
583 }
584 if (nargs + nkwds_found > 3) {
585 TYPE_ERR("too many arguments");
586 return 0;
587 }
588 } else {
589 switch (nargs) {
590 case 3:
591 *axis = PyTuple_GET_ITEM(args, 2);
592 case 2:
593 *n = PyTuple_GET_ITEM(args, 1);
594 case 1:
595 *a = PyTuple_GET_ITEM(args, 0);
596 break;
597 default:
598 TYPE_ERR("wrong number of arguments");
599 return 0;
600 }
601 }
602
603 return 1;
604
605 }
606
607 static PyObject *
nonreducer_axis(char * name,PyObject * args,PyObject * kwds,nra_t nra_float64,nra_t nra_float32,nra_t nra_int64,nra_t nra_int32,parse_type parse)608 nonreducer_axis(char *name,
609 PyObject *args,
610 PyObject *kwds,
611 nra_t nra_float64,
612 nra_t nra_float32,
613 nra_t nra_int64,
614 nra_t nra_int32,
615 parse_type parse) {
616
617 int n;
618 int axis;
619 int dtype;
620
621 PyArrayObject *a;
622 PyObject *y;
623
624 PyObject *a_obj = NULL;
625 PyObject *n_obj = NULL;
626 PyObject *axis_obj = NULL;
627
628 if (parse == PARSE_PARTITION) {
629 if (!parse_partition(args, kwds, &a_obj, &n_obj, &axis_obj)) {
630 return NULL;
631 }
632 } else if (parse == PARSE_RANKDATA) {
633 if (!parse_rankdata(args, kwds, &a_obj, &axis_obj)) {
634 return NULL;
635 }
636 } else if (parse == PARSE_PUSH) {
637 if (!parse_push(args, kwds, &a_obj, &n_obj, &axis_obj)) {
638 return NULL;
639 }
640 } else {
641 RUNTIME_ERR("Unknown parse type; please report error.");
642 }
643
644 /* convert to array if necessary */
645 if (PyArray_Check(a_obj)) {
646 a = (PyArrayObject *)a_obj;
647 Py_INCREF(a);
648 } else {
649 a = (PyArrayObject *)PyArray_FROM_O(a_obj);
650 if (a == NULL) {
651 return NULL;
652 }
653 }
654
655 /* check for byte swapped input array */
656 if (PyArray_ISBYTESWAPPED(a)) {
657 return slow(name, args, kwds);
658 }
659
660 /* defend against the axis of negativity */
661 if (axis_obj == NULL) {
662 if (parse == PARSE_PARTITION || parse == PARSE_PUSH) {
663 axis = PyArray_NDIM(a) - 1;
664 if (axis < 0) {
665 PyErr_Format(PyExc_ValueError,
666 "axis(=%d) out of bounds", axis);
667 goto error;
668 }
669 } else {
670 if (PyArray_NDIM(a) != 1) {
671 a = (PyArrayObject *)PyArray_Ravel(a, NPY_CORDER);
672 }
673 axis = 0;
674 }
675 } else if (axis_obj == Py_None) {
676 if (parse == PARSE_PUSH) {
677 VALUE_ERR("`axis` cannot be None");
678 goto error;
679 }
680 if (PyArray_NDIM(a) != 1) {
681 a = (PyArrayObject *)PyArray_Ravel(a, NPY_CORDER);
682 }
683 axis = 0;
684 } else {
685 axis = PyArray_PyIntAsInt(axis_obj);
686 if (error_converting(axis)) {
687 TYPE_ERR("`axis` must be an integer");
688 goto error;
689 }
690 if (axis < 0) {
691 axis += PyArray_NDIM(a);
692 if (axis < 0) {
693 PyErr_Format(PyExc_ValueError,
694 "axis(=%d) out of bounds", axis);
695 goto error;
696 }
697 } else if (axis >= PyArray_NDIM(a)) {
698 PyErr_Format(PyExc_ValueError, "axis(=%d) out of bounds", axis);
699 goto error;
700 }
701 }
702
703 /* n */
704 if (n_obj == NULL) {
705 n = -1;
706 } else if (parse == PARSE_PUSH && n_obj == Py_None) {
707 n = -1;
708 } else {
709 n = PyArray_PyIntAsInt(n_obj);
710 if (error_converting(n)) {
711 TYPE_ERR("`n` must be an integer");
712 goto error;
713 }
714 if (n < 0 && parse == PARSE_PUSH) {
715 VALUE_ERR("`n` must be nonnegative");
716 goto error;
717 }
718 }
719
720 dtype = PyArray_TYPE(a);
721 if (dtype == NPY_float64) y = nra_float64(a, axis, n);
722 else if (dtype == NPY_float32) y = nra_float32(a, axis, n);
723 else if (dtype == NPY_int64) y = nra_int64(a, axis, n);
724 else if (dtype == NPY_int32) y = nra_int32(a, axis, n);
725 else y = slow(name, args, kwds);
726
727 Py_DECREF(a);
728
729 return y;
730
731 error:
732 Py_DECREF(a);
733 return NULL;
734
735 }
736
737 /* docstrings ------------------------------------------------------------- */
738
739 static char nra_doc[] =
740 "Bottleneck non-reducing functions that operate along an axis.";
741
742 static char partition_doc[] =
743 /* MULTILINE STRING BEGIN
744 partition(a, kth, axis=-1)
745
746 Partition array elements along given axis.
747
748 A 1d array B is partitioned at array index `kth` if three conditions
749 are met: (1) B[kth] is in its sorted position, (2) all elements to the
750 left of `kth` are less than or equal to B[kth], and (3) all elements
751 to the right of `kth` are greater than or equal to B[kth]. Note that
752 the array elements in conditions (2) and (3) are in general unordered.
753
754 Shuffling the input array may change the output. The only guarantee is
755 given by the three conditions above.
756
757 This functions is not protected against NaN. Therefore, you may get
758 unexpected results if the input contains NaN.
759
760 Parameters
761 ----------
762 a : array_like
763 Input array. If `a` is not an array, a conversion is attempted.
764 kth : int
765 The value of the element at index `kth` will be in its sorted
766 position. Smaller (larger) or equal values will be to the left
767 (right) of index `kth`.
768 axis : {int, None}, optional
769 Axis along which the partition is performed. The default
770 (axis=-1) is to partition along the last axis.
771
772 Returns
773 -------
774 y : ndarray
775 A partitioned copy of the input array with the same shape and
776 type of `a`.
777
778 See Also
779 --------
780 bottleneck.argpartition: Indices that would partition an array
781
782 Notes
783 -----
784 Unexpected results may occur if the input array contains NaN.
785
786 Examples
787 --------
788 Create a numpy array:
789
790 >>> a = np.array([1, 0, 3, 4, 2])
791
792 Partition array so that the first 3 elements (indices 0, 1, 2) are the
793 smallest 3 elements (note, as in this example, that the smallest 3
794 elements may not be sorted):
795
796 >>> bn.partition(a, kth=2)
797 array([1, 0, 2, 4, 3])
798
799 Now Partition array so that the last 2 elements are the largest 2
800 elements:
801
802 >>> bn.partition(a, kth=3)
803 array([1, 0, 2, 3, 4])
804
805 MULTILINE STRING END */
806
807 static char argpartition_doc[] =
808 /* MULTILINE STRING BEGIN
809 argpartition(a, kth, axis=-1)
810
811 Return indices that would partition array along the given axis.
812
813 A 1d array B is partitioned at array index `kth` if three conditions
814 are met: (1) B[kth] is in its sorted position, (2) all elements to the
815 left of `kth` are less than or equal to B[kth], and (3) all elements
816 to the right of `kth` are greater than or equal to B[kth]. Note that
817 the array elements in conditions (2) and (3) are in general unordered.
818
819 Shuffling the input array may change the output. The only guarantee is
820 given by the three conditions above.
821
822 This functions is not protected against NaN. Therefore, you may get
823 unexpected results if the input contains NaN.
824
825 Parameters
826 ----------
827 a : array_like
828 Input array. If `a` is not an array, a conversion is attempted.
829 kth : int
830 The value of the element at index `kth` will be in its sorted
831 position. Smaller (larger) or equal values will be to the left
832 (right) of index `kth`.
833 axis : {int, None}, optional
834 Axis along which the partition is performed. The default (axis=-1)
835 is to partition along the last axis.
836
837 Returns
838 -------
839 y : ndarray
840 An array the same shape as the input array containing the indices
841 that partition `a`. The dtype of the indices is numpy.intp.
842
843 See Also
844 --------
845 bottleneck.partition: Partition array elements along given axis.
846
847 Notes
848 -----
849 Unexpected results may occur if the input array contains NaN.
850
851 Examples
852 --------
853 Create a numpy array:
854
855 >>> a = np.array([10, 0, 30, 40, 20])
856
857 Find the indices that partition the array so that the first 3
858 elements are the smallest 3 elements:
859
860 >>> index = bn.argpartition(a, kth=2)
861 >>> index
862 array([0, 1, 4, 3, 2])
863
864 Let's use the indices to partition the array (note, as in this
865 example, that the smallest 3 elements may not be in order):
866
867 >>> a[index]
868 array([10, 0, 20, 40, 30])
869
870 MULTILINE STRING END */
871
872 static char rankdata_doc[] =
873 /* MULTILINE STRING BEGIN
874 rankdata(a, axis=None)
875
876 Ranks the data, dealing with ties appropriately.
877
878 Equal values are assigned a rank that is the average of the ranks that
879 would have been otherwise assigned to all of the values within that set.
880 Ranks begin at 1, not 0.
881
882 Parameters
883 ----------
884 a : array_like
885 Input array. If `a` is not an array, a conversion is attempted.
886 axis : {int, None}, optional
887 Axis along which the elements of the array are ranked. The default
888 (axis=None) is to rank the elements of the flattened array.
889
890 Returns
891 -------
892 y : ndarray
893 An array with the same shape as `a`. The dtype is 'float64'.
894
895 See also
896 --------
897 bottleneck.nanrankdata: Ranks the data dealing with ties and NaNs.
898
899 Examples
900 --------
901 >>> bn.rankdata([0, 2, 2, 3])
902 array([ 1. , 2.5, 2.5, 4. ])
903 >>> bn.rankdata([[0, 2], [2, 3]])
904 array([ 1. , 2.5, 2.5, 4. ])
905 >>> bn.rankdata([[0, 2], [2, 3]], axis=0)
906 array([[ 1., 1.],
907 [ 2., 2.]])
908 >>> bn.rankdata([[0, 2], [2, 3]], axis=1)
909 array([[ 1., 2.],
910 [ 1., 2.]])
911
912 MULTILINE STRING END */
913
914 static char nanrankdata_doc[] =
915 /* MULTILINE STRING BEGIN
916 nanrankdata(a, axis=None)
917
918 Ranks the data, dealing with ties and NaNs appropriately.
919
920 Equal values are assigned a rank that is the average of the ranks that
921 would have been otherwise assigned to all of the values within that set.
922 Ranks begin at 1, not 0.
923
924 NaNs in the input array are returned as NaNs.
925
926 Parameters
927 ----------
928 a : array_like
929 Input array. If `a` is not an array, a conversion is attempted.
930 axis : {int, None}, optional
931 Axis along which the elements of the array are ranked. The default
932 (axis=None) is to rank the elements of the flattened array.
933
934 Returns
935 -------
936 y : ndarray
937 An array with the same shape as `a`. The dtype is 'float64'.
938
939 See also
940 --------
941 bottleneck.rankdata: Ranks the data, dealing with ties and appropriately.
942
943 Examples
944 --------
945 >>> bn.nanrankdata([np.nan, 2, 2, 3])
946 array([ nan, 1.5, 1.5, 3. ])
947 >>> bn.nanrankdata([[np.nan, 2], [2, 3]])
948 array([ nan, 1.5, 1.5, 3. ])
949 >>> bn.nanrankdata([[np.nan, 2], [2, 3]], axis=0)
950 array([[ nan, 1.],
951 [ 1., 2.]])
952 >>> bn.nanrankdata([[np.nan, 2], [2, 3]], axis=1)
953 array([[ nan, 1.],
954 [ 1., 2.]])
955
956 MULTILINE STRING END */
957
958 static char push_doc[] =
959 /* MULTILINE STRING BEGIN
960 push(a, n=None, axis=-1)
961
962 Fill missing values (NaNs) with most recent non-missing values.
963
964 Filling proceeds along the specified axis from small index values to large
965 index values.
966
967 Parameters
968 ----------
969 a : array_like
970 Input array. If `a` is not an array, a conversion is attempted.
971 n : {int, None}, optional
972 How far to push values. If the most recent non-NaN array element is
973 more than `n` index positions away, than a NaN is returned. The default
974 (n = None) is to push the entire length of the slice. If `n` is an integer
975 it must be nonnegative.
976 axis : int, optional
977 Axis along which the elements of the array are pushed. The default
978 (axis=-1) is to push along the last axis of the input array.
979
980 Returns
981 -------
982 y : ndarray
983 An array with the same shape and dtype as `a`.
984
985 See also
986 --------
987 bottleneck.replace: Replace specified value of an array with new value.
988
989 Examples
990 --------
991 >>> a = np.array([5, np.nan, np.nan, 6, np.nan])
992 >>> bn.push(a)
993 array([ 5., 5., 5., 6., 6.])
994 >>> bn.push(a, n=1)
995 array([ 5., 5., nan, 6., 6.])
996 >>> bn.push(a, n=2)
997 array([ 5., 5., 5., 6., 6.])
998
999 MULTILINE STRING END */
1000
1001 /* python wrapper -------------------------------------------------------- */
1002
1003 static PyMethodDef
1004 nra_methods[] = {
1005 {"partition", (PyCFunction)partition, VARKEY, partition_doc},
1006 {"argpartition", (PyCFunction)argpartition, VARKEY, argpartition_doc},
1007 {"rankdata", (PyCFunction)rankdata, VARKEY, rankdata_doc},
1008 {"nanrankdata", (PyCFunction)nanrankdata, VARKEY, nanrankdata_doc},
1009 {"push", (PyCFunction)push, VARKEY, push_doc},
1010 {NULL, NULL, 0, NULL}
1011 };
1012
1013
1014 #if PY_MAJOR_VERSION >= 3
1015 static struct PyModuleDef
1016 nra_def = {
1017 PyModuleDef_HEAD_INIT,
1018 "nonreduce_axis",
1019 nra_doc,
1020 -1,
1021 nra_methods
1022 };
1023 #endif
1024
1025
1026 PyMODINIT_FUNC
1027 #if PY_MAJOR_VERSION >= 3
1028 #define RETVAL m
PyInit_nonreduce_axis(void)1029 PyInit_nonreduce_axis(void)
1030 #else
1031 #define RETVAL
1032 initnonreduce_axis(void)
1033 #endif
1034 {
1035 #if PY_MAJOR_VERSION >=3
1036 PyObject *m = PyModule_Create(&nra_def);
1037 #else
1038 PyObject *m = Py_InitModule3("nonreduce_axis", nra_methods, nra_doc);
1039 #endif
1040 if (m == NULL) return RETVAL;
1041 import_array();
1042 if (!intern_strings()) {
1043 return RETVAL;
1044 }
1045 return RETVAL;
1046 }
1047