1
#define PY_SSIZE_T_CLEAN
2
#include <Python.h>
3
#include "structmember.h"
4

5
#define NPY_NO_DEPRECATED_API NPY_API_VERSION
6
#define _MULTIARRAYMODULE
7
#include "numpy/arrayobject.h"
8
#include "numpy/arrayscalars.h"
9

10
#include "numpy/npy_math.h"
11
#include "numpy/npy_cpu.h"
12

13
#include "npy_config.h"
14

15
#include "npy_pycompat.h"
16

17
#include "multiarraymodule.h"
18
#include "common.h"
19
#include "arrayobject.h"
20
#include "ctors.h"
21
#include "lowlevel_strided_loops.h"
22
#include "array_assign.h"
23

24
#include "npy_sort.h"
25
#include "npy_partition.h"
26
#include "npy_binsearch.h"
27
#include "alloc.h"
28
#include "arraytypes.h"
29
#include "array_coercion.h"
30

31

32
static NPY_GCC_OPT_3 NPY_INLINE int
33 1
npy_fasttake_impl(
34
        char *dest, char *src, const npy_intp *indices,
35
        npy_intp n, npy_intp m, npy_intp max_item,
36
        npy_intp nelem, npy_intp chunk,
37
        NPY_CLIPMODE clipmode, npy_intp itemsize, int needs_refcounting,
38
        PyArray_Descr *dtype, int axis)
39
{
40 1
    NPY_BEGIN_THREADS_DEF;
41 1
    NPY_BEGIN_THREADS_DESCR(dtype);
42 1
    switch (clipmode) {
43
        case NPY_RAISE:
44 1
            for (npy_intp i = 0; i < n; i++) {
45 1
                for (npy_intp j = 0; j < m; j++) {
46 1
                    npy_intp tmp = indices[j];
47 1
                    if (check_and_adjust_index(&tmp, max_item, axis,
48
                                               _save) < 0) {
49 1
                        return -1;
50
                    }
51 1
                    char *tmp_src = src + tmp * chunk;
52 1
                    if (needs_refcounting) {
53 1
                        for (npy_intp k = 0; k < nelem; k++) {
54 1
                            PyArray_Item_INCREF(tmp_src, dtype);
55 1
                            PyArray_Item_XDECREF(dest, dtype);
56 1
                            memmove(dest, tmp_src, itemsize);
57 1
                            dest += itemsize;
58 1
                            tmp_src += itemsize;
59
                        }
60
                    }
61
                    else {
62 1
                        memmove(dest, tmp_src, chunk);
63 1
                        dest += chunk;
64
                    }
65
                }
66 1
                src += chunk*max_item;
67
            }
68
            break;
69
        case NPY_WRAP:
70 1
            for (npy_intp i = 0; i < n; i++) {
71 1
                for (npy_intp j = 0; j < m; j++) {
72 1
                    npy_intp tmp = indices[j];
73 1
                    if (tmp < 0) {
74 1
                        while (tmp < 0) {
75 1
                            tmp += max_item;
76
                        }
77
                    }
78 1
                    else if (tmp >= max_item) {
79 1
                        while (tmp >= max_item) {
80 1
                            tmp -= max_item;
81
                        }
82
                    }
83 1
                    char *tmp_src = src + tmp * chunk;
84 1
                    if (needs_refcounting) {
85 1
                        for (npy_intp k = 0; k < nelem; k++) {
86 1
                            PyArray_Item_INCREF(tmp_src, dtype);
87 1
                            PyArray_Item_XDECREF(dest, dtype);
88 1
                            memmove(dest, tmp_src, itemsize);
89 1
                            dest += itemsize;
90 1
                            tmp_src += itemsize;
91
                        }
92
                    }
93
                    else {
94 1
                        memmove(dest, tmp_src, chunk);
95 1
                        dest += chunk;
96
                    }
97
                }
98 1
                src += chunk*max_item;
99
            }
100
            break;
101
        case NPY_CLIP:
102 1
            for (npy_intp i = 0; i < n; i++) {
103 1
                for (npy_intp j = 0; j < m; j++) {
104 1
                    npy_intp tmp = indices[j];
105 1
                    if (tmp < 0) {
106
                        tmp = 0;
107
                    }
108 1
                    else if (tmp >= max_item) {
109 1
                        tmp = max_item - 1;
110
                    }
111 1
                    char *tmp_src = src + tmp * chunk;
112 1
                    if (needs_refcounting) {
113 1
                        for (npy_intp k = 0; k < nelem; k++) {
114 1
                            PyArray_Item_INCREF(tmp_src, dtype);
115 1
                            PyArray_Item_XDECREF(dest, dtype);
116 1
                            memmove(dest, tmp_src, itemsize);
117 1
                            dest += itemsize;
118 1
                            tmp_src += itemsize;
119
                        }
120
                    }
121
                    else {
122 1
                        memmove(dest, tmp_src, chunk);
123 1
                        dest += chunk;
124
                    }
125
                }
126 1
                src += chunk*max_item;
127
            }
128
            break;
129
    }
130

131 1
    NPY_END_THREADS;
132
    return 0;
133
}
134

135

136
/*
137
 * Helper function instantiating npy_fasttake_impl in different branches
138
 * to allow the compiler to optimize each to the specific itemsize.
139
 */
140
static NPY_GCC_OPT_3 int
141 1
npy_fasttake(
142
        char *dest, char *src, const npy_intp *indices,
143
        npy_intp n, npy_intp m, npy_intp max_item,
144
        npy_intp nelem, npy_intp chunk,
145
        NPY_CLIPMODE clipmode, npy_intp itemsize, int needs_refcounting,
146
        PyArray_Descr *dtype, int axis)
147
{
148 1
    if (!needs_refcounting) {
149 1
        if (chunk == 1) {
150 1
            return npy_fasttake_impl(
151
                    dest, src, indices, n, m, max_item, nelem, chunk,
152
                    clipmode, itemsize, needs_refcounting, dtype, axis);
153
        }
154 1
        if (chunk == 2) {
155 1
            return npy_fasttake_impl(
156
                    dest, src, indices, n, m, max_item, nelem, chunk,
157
                    clipmode, itemsize, needs_refcounting, dtype, axis);
158
        }
159 1
        if (chunk == 4) {
160 1
            return npy_fasttake_impl(
161
                    dest, src, indices, n, m, max_item, nelem, chunk,
162
                    clipmode, itemsize, needs_refcounting, dtype, axis);
163
        }
164 1
        if (chunk == 8) {
165 1
            return npy_fasttake_impl(
166
                    dest, src, indices, n, m, max_item, nelem, chunk,
167
                    clipmode, itemsize, needs_refcounting, dtype, axis);
168
        }
169 1
        if (chunk == 16) {
170 1
            return npy_fasttake_impl(
171
                    dest, src, indices, n, m, max_item, nelem, chunk,
172
                    clipmode, itemsize, needs_refcounting, dtype, axis);
173
        }
174 1
        if (chunk == 32) {
175 1
            return npy_fasttake_impl(
176
                    dest, src, indices, n, m, max_item, nelem, chunk,
177
                    clipmode, itemsize, needs_refcounting, dtype, axis);
178
        }
179
    }
180

181 1
    return npy_fasttake_impl(
182
            dest, src, indices, n, m, max_item, nelem, chunk,
183
            clipmode, itemsize, needs_refcounting, dtype, axis);
184
}
185

186

187
/*NUMPY_API
188
 * Take
189
 */
190
NPY_NO_EXPORT PyObject *
191 1
PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
192
                 PyArrayObject *out, NPY_CLIPMODE clipmode)
193
{
194
    PyArray_Descr *dtype;
195 1
    PyArrayObject *obj = NULL, *self, *indices;
196
    npy_intp nd, i, n, m, max_item, chunk, itemsize, nelem;
197
    npy_intp shape[NPY_MAXDIMS];
198

199
    npy_bool needs_refcounting;
200

201 1
    indices = NULL;
202 1
    self = (PyArrayObject *)PyArray_CheckAxis(self0, &axis,
203
                                    NPY_ARRAY_CARRAY_RO);
204 1
    if (self == NULL) {
205
        return NULL;
206
    }
207 1
    indices = (PyArrayObject *)PyArray_ContiguousFromAny(indices0,
208
                                                         NPY_INTP,
209
                                                         0, 0);
210 1
    if (indices == NULL) {
211
        goto fail;
212
    }
213

214 1
    n = m = chunk = 1;
215 1
    nd = PyArray_NDIM(self) + PyArray_NDIM(indices) - 1;
216 1
    for (i = 0; i < nd; i++) {
217 1
        if (i < axis) {
218 1
            shape[i] = PyArray_DIMS(self)[i];
219 1
            n *= shape[i];
220
        }
221
        else {
222 1
            if (i < axis+PyArray_NDIM(indices)) {
223 1
                shape[i] = PyArray_DIMS(indices)[i-axis];
224 1
                m *= shape[i];
225
            }
226
            else {
227 1
                shape[i] = PyArray_DIMS(self)[i-PyArray_NDIM(indices)+1];
228 1
                chunk *= shape[i];
229
            }
230
        }
231
    }
232 1
    if (!out) {
233 1
        dtype = PyArray_DESCR(self);
234 1
        Py_INCREF(dtype);
235 1
        obj = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self),
236
                                                    dtype,
237
                                                    nd, shape,
238
                                                    NULL, NULL, 0,
239
                                                    (PyObject *)self);
240

241 1
        if (obj == NULL) {
242
            goto fail;
243
        }
244

245
    }
246
    else {
247 1
        int flags = NPY_ARRAY_CARRAY | NPY_ARRAY_WRITEBACKIFCOPY;
248

249 1
        if ((PyArray_NDIM(out) != nd) ||
250 1
            !PyArray_CompareLists(PyArray_DIMS(out), shape, nd)) {
251 0
            PyErr_SetString(PyExc_ValueError,
252
                        "output array does not match result of ndarray.take");
253 0
            goto fail;
254
        }
255

256 1
        if (arrays_overlap(out, self)) {
257 1
            flags |= NPY_ARRAY_ENSURECOPY;
258
        }
259

260 1
        if (clipmode == NPY_RAISE) {
261
            /*
262
             * we need to make sure and get a copy
263
             * so the input array is not changed
264
             * before the error is called
265
             */
266 1
            flags |= NPY_ARRAY_ENSURECOPY;
267
        }
268 1
        dtype = PyArray_DESCR(self);
269 1
        Py_INCREF(dtype);
270 1
        obj = (PyArrayObject *)PyArray_FromArray(out, dtype, flags);
271 1
        if (obj == NULL) {
272
            goto fail;
273
        }
274
    }
275

276 1
    max_item = PyArray_DIMS(self)[axis];
277 1
    nelem = chunk;
278 1
    itemsize = PyArray_ITEMSIZE(obj);
279 1
    chunk = chunk * itemsize;
280 1
    char *src = PyArray_DATA(self);
281 1
    char *dest = PyArray_DATA(obj);
282 1
    needs_refcounting = PyDataType_REFCHK(PyArray_DESCR(self));
283 1
    npy_intp *indices_data = (npy_intp *)PyArray_DATA(indices);
284

285 1
    if ((max_item == 0) && (PyArray_SIZE(obj) != 0)) {
286
        /* Index error, since that is the usual error for raise mode */
287 1
        PyErr_SetString(PyExc_IndexError,
288
                    "cannot do a non-empty take from an empty axes.");
289 1
        goto fail;
290
    }
291

292 1
    if (npy_fasttake(
293
            dest, src, indices_data, n, m, max_item, nelem, chunk,
294
            clipmode, itemsize, needs_refcounting, dtype, axis) < 0) {
295
        goto fail;
296
    }
297

298 1
    Py_XDECREF(indices);
299 1
    Py_XDECREF(self);
300 1
    if (out != NULL && out != obj) {
301 1
        Py_INCREF(out);
302 1
        PyArray_ResolveWritebackIfCopy(obj);
303 1
        Py_DECREF(obj);
304
        obj = out;
305
    }
306
    return (PyObject *)obj;
307

308 1
 fail:
309 1
    PyArray_DiscardWritebackIfCopy(obj);
310 1
    Py_XDECREF(obj);
311 1
    Py_XDECREF(indices);
312 1
    Py_XDECREF(self);
313
    return NULL;
314
}
315

316
/*NUMPY_API
317
 * Put values into an array
318
 */
319
NPY_NO_EXPORT PyObject *
320 1
PyArray_PutTo(PyArrayObject *self, PyObject* values0, PyObject *indices0,
321
              NPY_CLIPMODE clipmode)
322
{
323
    PyArrayObject  *indices, *values;
324
    npy_intp i, chunk, ni, max_item, nv, tmp;
325
    char *src, *dest;
326 1
    int copied = 0;
327 1
    int overlap = 0;
328

329 1
    indices = NULL;
330 1
    values = NULL;
331 1
    if (!PyArray_Check(self)) {
332 0
        PyErr_SetString(PyExc_TypeError,
333
                        "put: first argument must be an array");
334 0
        return NULL;
335
    }
336

337 1
    if (PyArray_FailUnlessWriteable(self, "put: output array") < 0) {
338
        return NULL;
339
    }
340

341 1
    indices = (PyArrayObject *)PyArray_ContiguousFromAny(indices0,
342
                                                         NPY_INTP, 0, 0);
343 1
    if (indices == NULL) {
344
        goto fail;
345
    }
346 1
    ni = PyArray_SIZE(indices);
347 1
    Py_INCREF(PyArray_DESCR(self));
348 1
    values = (PyArrayObject *)PyArray_FromAny(values0, PyArray_DESCR(self), 0, 0,
349
                              NPY_ARRAY_DEFAULT | NPY_ARRAY_FORCECAST, NULL);
350 1
    if (values == NULL) {
351
        goto fail;
352
    }
353 1
    nv = PyArray_SIZE(values);
354 1
    if (nv <= 0) {
355
        goto finish;
356
    }
357

358 1
    overlap = arrays_overlap(self, values) || arrays_overlap(self, indices);
359 1
    if (overlap || !PyArray_ISCONTIGUOUS(self)) {
360
        PyArrayObject *obj;
361 1
        int flags = NPY_ARRAY_CARRAY | NPY_ARRAY_WRITEBACKIFCOPY |
362
                    NPY_ARRAY_ENSURECOPY;
363

364 1
        Py_INCREF(PyArray_DESCR(self));
365 1
        obj = (PyArrayObject *)PyArray_FromArray(self,
366
                                                 PyArray_DESCR(self), flags);
367 1
        if (obj != self) {
368 1
            copied = 1;
369
        }
370
        self = obj;
371
    }
372 1
    max_item = PyArray_SIZE(self);
373 1
    dest = PyArray_DATA(self);
374 1
    chunk = PyArray_DESCR(self)->elsize;
375

376 1
    if (PyDataType_REFCHK(PyArray_DESCR(self))) {
377 1
        switch(clipmode) {
378
        case NPY_RAISE:
379 1
            for (i = 0; i < ni; i++) {
380 1
                src = PyArray_BYTES(values) + chunk*(i % nv);
381 1
                tmp = ((npy_intp *)(PyArray_DATA(indices)))[i];
382 1
                if (check_and_adjust_index(&tmp, max_item, 0, NULL) < 0) {
383
                    goto fail;
384
                }
385 1
                PyArray_Item_INCREF(src, PyArray_DESCR(self));
386 1
                PyArray_Item_XDECREF(dest+tmp*chunk, PyArray_DESCR(self));
387 1
                memmove(dest + tmp*chunk, src, chunk);
388
            }
389
            break;
390
        case NPY_WRAP:
391 0
            for (i = 0; i < ni; i++) {
392 0
                src = PyArray_BYTES(values) + chunk * (i % nv);
393 0
                tmp = ((npy_intp *)(PyArray_DATA(indices)))[i];
394 0
                if (tmp < 0) {
395 0
                    while (tmp < 0) {
396 0
                        tmp += max_item;
397
                    }
398
                }
399 0
                else if (tmp >= max_item) {
400 0
                    while (tmp >= max_item) {
401 0
                        tmp -= max_item;
402
                    }
403
                }
404 0
                PyArray_Item_INCREF(src, PyArray_DESCR(self));
405 0
                PyArray_Item_XDECREF(dest+tmp*chunk, PyArray_DESCR(self));
406 0
                memmove(dest + tmp * chunk, src, chunk);
407
            }
408
            break;
409
        case NPY_CLIP:
410 0
            for (i = 0; i < ni; i++) {
411 0
                src = PyArray_BYTES(values) + chunk * (i % nv);
412 0
                tmp = ((npy_intp *)(PyArray_DATA(indices)))[i];
413 0
                if (tmp < 0) {
414 0
                    tmp = 0;
415
                }
416 0
                else if (tmp >= max_item) {
417 0
                    tmp = max_item - 1;
418
                }
419 0
                PyArray_Item_INCREF(src, PyArray_DESCR(self));
420 0
                PyArray_Item_XDECREF(dest+tmp*chunk, PyArray_DESCR(self));
421 0
                memmove(dest + tmp * chunk, src, chunk);
422
            }
423
            break;
424
        }
425
    }
426
    else {
427 1
        NPY_BEGIN_THREADS_DEF;
428 1
        NPY_BEGIN_THREADS_THRESHOLDED(ni);
429 1
        switch(clipmode) {
430
        case NPY_RAISE:
431 1
            for (i = 0; i < ni; i++) {
432 1
                src = PyArray_BYTES(values) + chunk * (i % nv);
433 1
                tmp = ((npy_intp *)(PyArray_DATA(indices)))[i];
434 1
                if (check_and_adjust_index(&tmp, max_item, 0, _save) < 0) {
435
                    goto fail;
436
                }
437 1
                memmove(dest + tmp * chunk, src, chunk);
438
            }
439
            break;
440
        case NPY_WRAP:
441 0
            for (i = 0; i < ni; i++) {
442 0
                src = PyArray_BYTES(values) + chunk * (i % nv);
443 0
                tmp = ((npy_intp *)(PyArray_DATA(indices)))[i];
444 0
                if (tmp < 0) {
445 0
                    while (tmp < 0) {
446 0
                        tmp += max_item;
447
                    }
448
                }
449 0
                else if (tmp >= max_item) {
450 0
                    while (tmp >= max_item) {
451 0
                        tmp -= max_item;
452
                    }
453
                }
454 0
                memmove(dest + tmp * chunk, src, chunk);
455
            }
456
            break;
457
        case NPY_CLIP:
458 0
            for (i = 0; i < ni; i++) {
459 0
                src = PyArray_BYTES(values) + chunk * (i % nv);
460 0
                tmp = ((npy_intp *)(PyArray_DATA(indices)))[i];
461 0
                if (tmp < 0) {
462 0
                    tmp = 0;
463
                }
464 0
                else if (tmp >= max_item) {
465 0
                    tmp = max_item - 1;
466
                }
467 0
                memmove(dest + tmp * chunk, src, chunk);
468
            }
469
            break;
470
        }
471 1
        NPY_END_THREADS;
472
    }
473

474 1
 finish:
475 1
    Py_XDECREF(values);
476 1
    Py_XDECREF(indices);
477 1
    if (copied) {
478 1
        PyArray_ResolveWritebackIfCopy(self);
479 1
        Py_DECREF(self);
480
    }
481 1
    Py_RETURN_NONE;
482

483 1
 fail:
484 1
    Py_XDECREF(indices);
485 1
    Py_XDECREF(values);
486 1
    if (copied) {
487 0
        PyArray_DiscardWritebackIfCopy(self);
488 0
        Py_XDECREF(self);
489
    }
490
    return NULL;
491
}
492

493

494
static NPY_GCC_OPT_3 NPY_INLINE void
495 1
npy_fastputmask_impl(
496
        char *dest, char *src, const npy_bool *mask_data,
497
        npy_intp ni, npy_intp nv, npy_intp chunk)
498
{
499 1
    if (nv == 1) {
500 1
        for (npy_intp i = 0; i < ni; i++) {
501 1
            if (mask_data[i]) {
502 1
                memmove(dest, src, chunk);
503
            }
504 1
            dest += chunk;
505
        }
506
    }
507
    else {
508
        char *tmp_src = src;
509 1
        for (npy_intp i = 0, j = 0; i < ni; i++, j++) {
510 1
            if (NPY_UNLIKELY(j >= nv)) {
511 0
                j = 0;
512 0
                tmp_src = src;
513
            }
514 1
            if (mask_data[i]) {
515 1
                memmove(dest, tmp_src, chunk);
516
            }
517 1
            dest += chunk;
518 1
            tmp_src += chunk;
519
        }
520
    }
521
}
522

523

524
/*
525
 * Helper function instantiating npy_fastput_impl in different branches
526
 * to allow the compiler to optimize each to the specific itemsize.
527
 */
528
static NPY_GCC_OPT_3 void
529 1
npy_fastputmask(
530
        char *dest, char *src, npy_bool *mask_data,
531
        npy_intp ni, npy_intp nv, npy_intp chunk)
532
{
533 1
    if (chunk == 1) {
534 1
        return npy_fastputmask_impl(dest, src, mask_data, ni, nv, chunk);
535
    }
536 1
    if (chunk == 2) {
537 1
        return npy_fastputmask_impl(dest, src, mask_data, ni, nv, chunk);
538
    }
539 1
    if (chunk == 4) {
540 1
        return npy_fastputmask_impl(dest, src, mask_data, ni, nv, chunk);
541
    }
542 1
    if (chunk == 8) {
543 1
        return npy_fastputmask_impl(dest, src, mask_data, ni, nv, chunk);
544
    }
545 1
    if (chunk == 16) {
546 1
        return npy_fastputmask_impl(dest, src, mask_data, ni, nv, chunk);
547
    }
548 1
    if (chunk == 32) {
549 1
        return npy_fastputmask_impl(dest, src, mask_data, ni, nv, chunk);
550
    }
551

552 1
    return npy_fastputmask_impl(dest, src, mask_data, ni, nv, chunk);
553
}
554

555

556
/*NUMPY_API
557
 * Put values into an array according to a mask.
558
 */
559
NPY_NO_EXPORT PyObject *
560 1
PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
561
{
562
    PyArrayObject *mask, *values;
563
    PyArray_Descr *dtype;
564
    npy_intp chunk, ni, nv;
565
    char *src, *dest;
566
    npy_bool *mask_data;
567 1
    int copied = 0;
568 1
    int overlap = 0;
569

570 1
    mask = NULL;
571 1
    values = NULL;
572 1
    if (!PyArray_Check(self)) {
573 0
        PyErr_SetString(PyExc_TypeError,
574
                        "putmask: first argument must "
575
                        "be an array");
576 0
        return NULL;
577
    }
578

579 1
    mask = (PyArrayObject *)PyArray_FROM_OTF(mask0, NPY_BOOL,
580
                                NPY_ARRAY_CARRAY | NPY_ARRAY_FORCECAST);
581 1
    if (mask == NULL) {
582
        goto fail;
583
    }
584 1
    ni = PyArray_SIZE(mask);
585 1
    if (ni != PyArray_SIZE(self)) {
586 1
        PyErr_SetString(PyExc_ValueError,
587
                        "putmask: mask and data must be "
588
                        "the same size");
589 1
        goto fail;
590
    }
591 1
    mask_data = PyArray_DATA(mask);
592 1
    dtype = PyArray_DESCR(self);
593 1
    Py_INCREF(dtype);
594 1
    values = (PyArrayObject *)PyArray_FromAny(values0, dtype,
595
                                    0, 0, NPY_ARRAY_CARRAY, NULL);
596 1
    if (values == NULL) {
597
        goto fail;
598
    }
599 1
    nv = PyArray_SIZE(values); /* zero if null array */
600 1
    if (nv <= 0) {
601 0
        Py_XDECREF(values);
602 0
        Py_XDECREF(mask);
603 0
        Py_RETURN_NONE;
604
    }
605 1
    src = PyArray_DATA(values);
606

607 1
    overlap = arrays_overlap(self, values) || arrays_overlap(self, mask);
608 1
    if (overlap || !PyArray_ISCONTIGUOUS(self)) {
609 1
        int flags = NPY_ARRAY_CARRAY | NPY_ARRAY_WRITEBACKIFCOPY;
610
        PyArrayObject *obj;
611

612 1
        if (overlap) {
613 1
            flags |= NPY_ARRAY_ENSURECOPY;
614
        }
615

616 1
        dtype = PyArray_DESCR(self);
617 1
        Py_INCREF(dtype);
618 1
        obj = (PyArrayObject *)PyArray_FromArray(self, dtype, flags);
619 1
        if (obj != self) {
620 1
            copied = 1;
621
        }
622
        self = obj;
623
    }
624

625 1
    chunk = PyArray_DESCR(self)->elsize;
626 1
    dest = PyArray_DATA(self);
627

628 1
    if (PyDataType_REFCHK(PyArray_DESCR(self))) {
629 1
        for (npy_intp i = 0, j = 0; i < ni; i++, j++) {
630 1
            if (j >= nv) {
631 1
                j = 0;
632
            }
633 1
            if (mask_data[i]) {
634 1
                char *src_ptr = src + j*chunk;
635 1
                char *dest_ptr = dest + i*chunk;
636

637 1
                PyArray_Item_INCREF(src_ptr, PyArray_DESCR(self));
638 1
                PyArray_Item_XDECREF(dest_ptr, PyArray_DESCR(self));
639 1
                memmove(dest_ptr, src_ptr, chunk);
640
            }
641
        }
642
    }
643
    else {
644 1
        NPY_BEGIN_THREADS_DEF;
645 1
        NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(self));
646 1
        npy_fastputmask(dest, src, mask_data, ni, nv, chunk);
647 1
        NPY_END_THREADS;
648
    }
649

650 1
    Py_XDECREF(values);
651 1
    Py_XDECREF(mask);
652 1
    if (copied) {
653 1
        PyArray_ResolveWritebackIfCopy(self);
654 1
        Py_DECREF(self);
655
    }
656 1
    Py_RETURN_NONE;
657

658 0
 fail:
659 1
    Py_XDECREF(mask);
660 1
    Py_XDECREF(values);
661
    if (copied) {
662
        PyArray_DiscardWritebackIfCopy(self);
663
        Py_XDECREF(self);
664
    }
665
    return NULL;
666
}
667

668
/*NUMPY_API
669
 * Repeat the array.
670
 */
671
NPY_NO_EXPORT PyObject *
672 1
PyArray_Repeat(PyArrayObject *aop, PyObject *op, int axis)
673
{
674
    npy_intp *counts;
675
    npy_intp n, n_outer, i, j, k, chunk;
676 1
    npy_intp total = 0;
677 1
    npy_bool broadcast = NPY_FALSE;
678 1
    PyArrayObject *repeats = NULL;
679 1
    PyObject *ap = NULL;
680 1
    PyArrayObject *ret = NULL;
681
    char *new_data, *old_data;
682

683 1
    repeats = (PyArrayObject *)PyArray_ContiguousFromAny(op, NPY_INTP, 0, 1);
684 1
    if (repeats == NULL) {
685
        return NULL;
686
    }
687

688
    /*
689
     * Scalar and size 1 'repeat' arrays broadcast to any shape, for all
690
     * other inputs the dimension must match exactly.
691
     */
692 1
    if (PyArray_NDIM(repeats) == 0 || PyArray_SIZE(repeats) == 1) {
693
        broadcast = NPY_TRUE;
694
    }
695

696 1
    counts = (npy_intp *)PyArray_DATA(repeats);
697

698 1
    if ((ap = PyArray_CheckAxis(aop, &axis, NPY_ARRAY_CARRAY)) == NULL) {
699 0
        Py_DECREF(repeats);
700
        return NULL;
701
    }
702

703 1
    aop = (PyArrayObject *)ap;
704 1
    n = PyArray_DIM(aop, axis);
705

706 1
    if (!broadcast && PyArray_SIZE(repeats) != n) {
707 0
        PyErr_Format(PyExc_ValueError,
708
                     "operands could not be broadcast together "
709
                     "with shape (%zd,) (%zd,)", n, PyArray_DIM(repeats, 0));
710 0
        goto fail;
711
    }
712 1
    if (broadcast) {
713 1
        total = counts[0] * n;
714
    }
715
    else {
716 1
        for (j = 0; j < n; j++) {
717 1
            if (counts[j] < 0) {
718 0
                PyErr_SetString(PyExc_ValueError,
719
                                "repeats may not contain negative values.");
720 0
                goto fail;
721
            }
722 1
            total += counts[j];
723
        }
724
    }
725

726
    /* Construct new array */
727 1
    PyArray_DIMS(aop)[axis] = total;
728 1
    Py_INCREF(PyArray_DESCR(aop));
729 1
    ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(aop),
730
                                                PyArray_DESCR(aop),
731
                                                PyArray_NDIM(aop),
732 1
                                                PyArray_DIMS(aop),
733
                                                NULL, NULL, 0,
734
                                                (PyObject *)aop);
735 1
    PyArray_DIMS(aop)[axis] = n;
736 1
    if (ret == NULL) {
737
        goto fail;
738
    }
739 1
    new_data = PyArray_DATA(ret);
740 1
    old_data = PyArray_DATA(aop);
741

742 1
    chunk = PyArray_DESCR(aop)->elsize;
743 1
    for(i = axis + 1; i < PyArray_NDIM(aop); i++) {
744 1
        chunk *= PyArray_DIMS(aop)[i];
745
    }
746

747
    n_outer = 1;
748 1
    for (i = 0; i < axis; i++) {
749 1
        n_outer *= PyArray_DIMS(aop)[i];
750
    }
751 1
    for (i = 0; i < n_outer; i++) {
752 1
        for (j = 0; j < n; j++) {
753 1
            npy_intp tmp = broadcast ? counts[0] : counts[j];
754 1
            for (k = 0; k < tmp; k++) {
755 1
                memcpy(new_data, old_data, chunk);
756 1
                new_data += chunk;
757
            }
758 1
            old_data += chunk;
759
        }
760
    }
761

762 1
    Py_DECREF(repeats);
763 1
    PyArray_INCREF(ret);
764 1
    Py_XDECREF(aop);
765
    return (PyObject *)ret;
766

767 0
 fail:
768 0
    Py_DECREF(repeats);
769 0
    Py_XDECREF(aop);
770 0
    Py_XDECREF(ret);
771
    return NULL;
772
}
773

774
/*NUMPY_API
775
 */
776
NPY_NO_EXPORT PyObject *
777 1
PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
778
               NPY_CLIPMODE clipmode)
779
{
780 1
    PyArrayObject *obj = NULL;
781
    PyArray_Descr *dtype;
782
    int n, elsize;
783
    npy_intp i;
784
    char *ret_data;
785
    PyArrayObject **mps, *ap;
786 1
    PyArrayMultiIterObject *multi = NULL;
787
    npy_intp mi;
788 1
    ap = NULL;
789

790
    /*
791
     * Convert all inputs to arrays of a common type
792
     * Also makes them C-contiguous
793
     */
794 1
    mps = PyArray_ConvertToCommonType(op, &n);
795 1
    if (mps == NULL) {
796
        return NULL;
797
    }
798 1
    for (i = 0; i < n; i++) {
799 1
        if (mps[i] == NULL) {
800
            goto fail;
801
        }
802
    }
803 1
    ap = (PyArrayObject *)PyArray_FROM_OT((PyObject *)ip, NPY_INTP);
804 1
    if (ap == NULL) {
805
        goto fail;
806
    }
807
    /* Broadcast all arrays to each other, index array at the end. */
808 1
    multi = (PyArrayMultiIterObject *)
809 1
        PyArray_MultiIterFromObjects((PyObject **)mps, n, 1, ap);
810 1
    if (multi == NULL) {
811
        goto fail;
812
    }
813
    /* Set-up return array */
814 1
    if (out == NULL) {
815 1
        dtype = PyArray_DESCR(mps[0]);
816 1
        Py_INCREF(dtype);
817 1
        obj = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(ap),
818
                                                    dtype,
819
                                                    multi->nd,
820 1
                                                    multi->dimensions,
821
                                                    NULL, NULL, 0,
822
                                                    (PyObject *)ap);
823
    }
824
    else {
825 1
        int flags = NPY_ARRAY_CARRAY |
826
                    NPY_ARRAY_WRITEBACKIFCOPY |
827
                    NPY_ARRAY_FORCECAST;
828

829 1
        if ((PyArray_NDIM(out) != multi->nd)
830 1
                    || !PyArray_CompareLists(PyArray_DIMS(out),
831 1
                                             multi->dimensions,
832
                                             multi->nd)) {
833 0
            PyErr_SetString(PyExc_TypeError,
834
                            "choose: invalid shape for output array.");
835 0
            goto fail;
836
        }
837

838 1
        for (i = 0; i < n; i++) {
839 1
            if (arrays_overlap(out, mps[i])) {
840 1
                flags |= NPY_ARRAY_ENSURECOPY;
841
            }
842
        }
843

844 1
        if (clipmode == NPY_RAISE) {
845
            /*
846
             * we need to make sure and get a copy
847
             * so the input array is not changed
848
             * before the error is called
849
             */
850 1
            flags |= NPY_ARRAY_ENSURECOPY;
851
        }
852 1
        dtype = PyArray_DESCR(mps[0]);
853 1
        Py_INCREF(dtype);
854 1
        obj = (PyArrayObject *)PyArray_FromArray(out, dtype, flags);
855
    }
856

857 1
    if (obj == NULL) {
858
        goto fail;
859
    }
860 1
    elsize = PyArray_DESCR(obj)->elsize;
861 1
    ret_data = PyArray_DATA(obj);
862

863 1
    while (PyArray_MultiIter_NOTDONE(multi)) {
864 1
        mi = *((npy_intp *)PyArray_MultiIter_DATA(multi, n));
865 1
        if (mi < 0 || mi >= n) {
866 1
            switch(clipmode) {
867 0
            case NPY_RAISE:
868 0
                PyErr_SetString(PyExc_ValueError,
869
                        "invalid entry in choice "\
870
                        "array");
871 0
                goto fail;
872 1
            case NPY_WRAP:
873 1
                if (mi < 0) {
874 0
                    while (mi < 0) {
875 0
                        mi += n;
876
                    }
877
                }
878
                else {
879 1
                    while (mi >= n) {
880 1
                        mi -= n;
881
                    }
882
                }
883
                break;
884 1
            case NPY_CLIP:
885 1
                if (mi < 0) {
886
                    mi = 0;
887
                }
888 1
                else if (mi >= n) {
889 1
                    mi = n - 1;
890
                }
891
                break;
892
            }
893
        }
894 1
        memmove(ret_data, PyArray_MultiIter_DATA(multi, mi), elsize);
895 1
        ret_data += elsize;
896 1
        PyArray_MultiIter_NEXT(multi);
897
    }
898

899 1
    PyArray_INCREF(obj);
900 1
    Py_DECREF(multi);
901 1
    for (i = 0; i < n; i++) {
902 1
        Py_XDECREF(mps[i]);
903
    }
904 1
    Py_DECREF(ap);
905 1
    npy_free_cache(mps, n * sizeof(mps[0]));
906 1
    if (out != NULL && out != obj) {
907 1
        Py_INCREF(out);
908 1
        PyArray_ResolveWritebackIfCopy(obj);
909 1
        Py_DECREF(obj);
910
        obj = out;
911
    }
912
    return (PyObject *)obj;
913

914 0
 fail:
915 0
    Py_XDECREF(multi);
916 0
    for (i = 0; i < n; i++) {
917 0
        Py_XDECREF(mps[i]);
918
    }
919 0
    Py_XDECREF(ap);
920 0
    npy_free_cache(mps, n * sizeof(mps[0]));
921 0
    PyArray_DiscardWritebackIfCopy(obj);
922 0
    Py_XDECREF(obj);
923
    return NULL;
924
}
925

926
/*
927
 * These algorithms use special sorting.  They are not called unless the
928
 * underlying sort function for the type is available.  Note that axis is
929
 * already valid. The sort functions require 1-d contiguous and well-behaved
930
 * data.  Therefore, a copy will be made of the data if needed before handing
931
 * it to the sorting routine.  An iterator is constructed and adjusted to walk
932
 * over all but the desired sorting axis.
933
 */
934
static int
935 1
_new_sortlike(PyArrayObject *op, int axis, PyArray_SortFunc *sort,
936
              PyArray_PartitionFunc *part, npy_intp const *kth, npy_intp nkth)
937
{
938 1
    npy_intp N = PyArray_DIM(op, axis);
939 1
    npy_intp elsize = (npy_intp)PyArray_ITEMSIZE(op);
940 1
    npy_intp astride = PyArray_STRIDE(op, axis);
941 1
    int swap = PyArray_ISBYTESWAPPED(op);
942 1
    int needcopy = !IsAligned(op) || swap || astride != elsize;
943 1
    int hasrefs = PyDataType_REFCHK(PyArray_DESCR(op));
944

945 1
    PyArray_CopySwapNFunc *copyswapn = PyArray_DESCR(op)->f->copyswapn;
946 1
    char *buffer = NULL;
947

948
    PyArrayIterObject *it;
949
    npy_intp size;
950

951 1
    int ret = 0;
952

953 1
    NPY_BEGIN_THREADS_DEF;
954

955
    /* Check if there is any sorting to do */
956 1
    if (N <= 1 || PyArray_SIZE(op) == 0) {
957
        return 0;
958
    }
959

960 1
    it = (PyArrayIterObject *)PyArray_IterAllButAxis((PyObject *)op, &axis);
961 1
    if (it == NULL) {
962
        return -1;
963
    }
964 1
    size = it->size;
965

966 1
    if (needcopy) {
967 1
        buffer = npy_alloc_cache(N * elsize);
968 1
        if (buffer == NULL) {
969
            ret = -1;
970
            goto fail;
971
        }
972
    }
973

974 1
    NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(op));
975

976 1
    while (size--) {
977 1
        char *bufptr = it->dataptr;
978

979 1
        if (needcopy) {
980 1
            if (hasrefs) {
981
                /*
982
                 * For dtype's with objects, copyswapn Py_XINCREF's src
983
                 * and Py_XDECREF's dst. This would crash if called on
984
                 * an uninitialized buffer, or leak a reference to each
985
                 * object if initialized.
986
                 *
987
                 * So, first do the copy with no refcounting...
988
                 */
989 0
                _unaligned_strided_byte_copy(buffer, elsize,
990
                                             it->dataptr, astride, N, elsize);
991
                /* ...then swap in-place if needed */
992 0
                if (swap) {
993 0
                    copyswapn(buffer, elsize, NULL, 0, N, swap, op);
994
                }
995
            }
996
            else {
997 1
                copyswapn(buffer, elsize, it->dataptr, astride, N, swap, op);
998
            }
999
            bufptr = buffer;
1000
        }
1001
        /*
1002
         * TODO: If the input array is byte-swapped but contiguous and
1003
         * aligned, it could be swapped (and later unswapped) in-place
1004
         * rather than after copying to the buffer. Care would have to
1005
         * be taken to ensure that, if there is an error in the call to
1006
         * sort or part, the unswapping is still done before returning.
1007
         */
1008

1009 1
        if (part == NULL) {
1010 1
            ret = sort(bufptr, N, op);
1011 1
            if (hasrefs && PyErr_Occurred()) {
1012
                ret = -1;
1013
            }
1014 1
            if (ret < 0) {
1015
                goto fail;
1016
            }
1017
        }
1018
        else {
1019
            npy_intp pivots[NPY_MAX_PIVOT_STACK];
1020 1
            npy_intp npiv = 0;
1021
            npy_intp i;
1022 1
            for (i = 0; i < nkth; ++i) {
1023 1
                ret = part(bufptr, N, kth[i], pivots, &npiv, op);
1024 1
                if (hasrefs && PyErr_Occurred()) {
1025
                    ret = -1;
1026
                }
1027 1
                if (ret < 0) {
1028
                    goto fail;
1029
                }
1030
            }
1031
        }
1032

1033 1
        if (needcopy) {
1034 1
            if (hasrefs) {
1035 0
                if (swap) {
1036 0
                    copyswapn(buffer, elsize, NULL, 0, N, swap, op);
1037
                }
1038 0
                _unaligned_strided_byte_copy(it->dataptr, astride,
1039
                                             buffer, elsize, N, elsize);
1040
            }
1041
            else {
1042 1
                copyswapn(it->dataptr, astride, buffer, elsize, N, swap, op);
1043
            }
1044
        }
1045

1046 1
        PyArray_ITER_NEXT(it);
1047
    }
1048

1049 1
fail:
1050 1
    NPY_END_THREADS_DESCR(PyArray_DESCR(op));
1051 1
    npy_free_cache(buffer, N * elsize);
1052 1
    if (ret < 0 && !PyErr_Occurred()) {
1053
        /* Out of memory during sorting or buffer creation */
1054 0
        PyErr_NoMemory();
1055
    }
1056 1
    Py_DECREF(it);
1057

1058
    return ret;
1059
}
1060

1061
static PyObject*
1062 1
_new_argsortlike(PyArrayObject *op, int axis, PyArray_ArgSortFunc *argsort,
1063
                 PyArray_ArgPartitionFunc *argpart,
1064
                 npy_intp const *kth, npy_intp nkth)
1065
{
1066 1
    npy_intp N = PyArray_DIM(op, axis);
1067 1
    npy_intp elsize = (npy_intp)PyArray_ITEMSIZE(op);
1068 1
    npy_intp astride = PyArray_STRIDE(op, axis);
1069 1
    int swap = PyArray_ISBYTESWAPPED(op);
1070 1
    int needcopy = !IsAligned(op) || swap || astride != elsize;
1071 1
    int hasrefs = PyDataType_REFCHK(PyArray_DESCR(op));
1072
    int needidxbuffer;
1073

1074 1
    PyArray_CopySwapNFunc *copyswapn = PyArray_DESCR(op)->f->copyswapn;
1075 1
    char *valbuffer = NULL;
1076 1
    npy_intp *idxbuffer = NULL;
1077

1078
    PyArrayObject *rop;
1079
    npy_intp rstride;
1080

1081
    PyArrayIterObject *it, *rit;
1082
    npy_intp size;
1083

1084 1
    int ret = 0;
1085

1086 1
    NPY_BEGIN_THREADS_DEF;
1087

1088 1
    rop = (PyArrayObject *)PyArray_NewFromDescr(
1089 1
            Py_TYPE(op), PyArray_DescrFromType(NPY_INTP),
1090 1
            PyArray_NDIM(op), PyArray_DIMS(op), NULL, NULL,
1091
            0, (PyObject *)op);
1092 1
    if (rop == NULL) {
1093
        return NULL;
1094
    }
1095 1
    rstride = PyArray_STRIDE(rop, axis);
1096 1
    needidxbuffer = rstride != sizeof(npy_intp);
1097

1098
    /* Check if there is any argsorting to do */
1099 1
    if (N <= 1 || PyArray_SIZE(op) == 0) {
1100 1
        memset(PyArray_DATA(rop), 0, PyArray_NBYTES(rop));
1101 1
        return (PyObject *)rop;
1102
    }
1103

1104 1
    it = (PyArrayIterObject *)PyArray_IterAllButAxis((PyObject *)op, &axis);
1105 1
    rit = (PyArrayIterObject *)PyArray_IterAllButAxis((PyObject *)rop, &axis);
1106 1
    if (it == NULL || rit == NULL) {
1107
        ret = -1;
1108
        goto fail;
1109
    }
1110 1
    size = it->size;
1111

1112 1
    if (needcopy) {
1113 1
        valbuffer = npy_alloc_cache(N * elsize);
1114 1
        if (valbuffer == NULL) {
1115
            ret = -1;
1116
            goto fail;
1117
        }
1118
    }
1119

1120 1
    if (needidxbuffer) {
1121 1
        idxbuffer = (npy_intp *)npy_alloc_cache(N * sizeof(npy_intp));
1122 1
        if (idxbuffer == NULL) {
1123
            ret = -1;
1124
            goto fail;
1125
        }
1126
    }
1127

1128 1
    NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(op));
1129

1130 1
    while (size--) {
1131 1
        char *valptr = it->dataptr;
1132 1
        npy_intp *idxptr = (npy_intp *)rit->dataptr;
1133
        npy_intp *iptr, i;
1134

1135 1
        if (needcopy) {
1136 1
            if (hasrefs) {
1137
                /*
1138
                 * For dtype's with objects, copyswapn Py_XINCREF's src
1139
                 * and Py_XDECREF's dst. This would crash if called on
1140
                 * an uninitialized valbuffer, or leak a reference to
1141
                 * each object item if initialized.
1142
                 *
1143
                 * So, first do the copy with no refcounting...
1144
                 */
1145 0
                 _unaligned_strided_byte_copy(valbuffer, elsize,
1146
                                              it->dataptr, astride, N, elsize);
1147
                /* ...then swap in-place if needed */
1148 0
                if (swap) {
1149 0
                    copyswapn(valbuffer, elsize, NULL, 0, N, swap, op);
1150
                }
1151
            }
1152
            else {
1153 1
                copyswapn(valbuffer, elsize,
1154
                          it->dataptr, astride, N, swap, op);
1155
            }
1156
            valptr = valbuffer;
1157
        }
1158

1159 1
        if (needidxbuffer) {
1160 1
            idxptr = idxbuffer;
1161
        }
1162

1163 1
        iptr = idxptr;
1164 1
        for (i = 0; i < N; ++i) {
1165 1
            *iptr++ = i;
1166
        }
1167

1168 1
        if (argpart == NULL) {
1169 1
            ret = argsort(valptr, idxptr, N, op);
1170
            /* Object comparisons may raise an exception in Python 3 */
1171 1
            if (hasrefs && PyErr_Occurred()) {
1172
                ret = -1;
1173
            }
1174 1
            if (ret < 0) {
1175
                goto fail;
1176
            }
1177
        }
1178
        else {
1179
            npy_intp pivots[NPY_MAX_PIVOT_STACK];
1180 1
            npy_intp npiv = 0;
1181

1182 1
            for (i = 0; i < nkth; ++i) {
1183 1
                ret = argpart(valptr, idxptr, N, kth[i], pivots, &npiv, op);
1184
                /* Object comparisons may raise an exception in Python 3 */
1185 1
                if (hasrefs && PyErr_Occurred()) {
1186
                    ret = -1;
1187
                }
1188 1
                if (ret < 0) {
1189
                    goto fail;
1190
                }
1191
            }
1192
        }
1193

1194 1
        if (needidxbuffer) {
1195 1
            char *rptr = rit->dataptr;
1196 1
            iptr = idxbuffer;
1197

1198 1
            for (i = 0; i < N; ++i) {
1199 1
                *(npy_intp *)rptr = *iptr++;
1200 1
                rptr += rstride;
1201
            }
1202
        }
1203

1204 1
        PyArray_ITER_NEXT(it);
1205 1
        PyArray_ITER_NEXT(rit);
1206
    }
1207

1208 1
fail:
1209 1
    NPY_END_THREADS_DESCR(PyArray_DESCR(op));
1210 1
    npy_free_cache(valbuffer, N * elsize);
1211 1
    npy_free_cache(idxbuffer, N * sizeof(npy_intp));
1212 1
    if (ret < 0) {
1213 0
        if (!PyErr_Occurred()) {
1214
            /* Out of memory during sorting or buffer creation */
1215 0
            PyErr_NoMemory();
1216
        }
1217 0
        Py_XDECREF(rop);
1218
        rop = NULL;
1219
    }
1220 1
    Py_XDECREF(it);
1221 1
    Py_XDECREF(rit);
1222

1223
    return (PyObject *)rop;
1224
}
1225

1226

1227
/*NUMPY_API
1228
 * Sort an array in-place
1229
 */
1230
NPY_NO_EXPORT int
1231 1
PyArray_Sort(PyArrayObject *op, int axis, NPY_SORTKIND which)
1232
{
1233 1
    PyArray_SortFunc *sort = NULL;
1234 1
    int n = PyArray_NDIM(op);
1235

1236 1
    if (check_and_adjust_axis(&axis, n) < 0) {
1237
        return -1;
1238
    }
1239

1240 1
    if (PyArray_FailUnlessWriteable(op, "sort array") < 0) {
1241
        return -1;
1242
    }
1243

1244 1
    if (which < 0 || which >= NPY_NSORTS) {
1245 0
        PyErr_SetString(PyExc_ValueError, "not a valid sort kind");
1246 0
        return -1;
1247
    }
1248

1249 1
    sort = PyArray_DESCR(op)->f->sort[which];
1250

1251 1
    if (sort == NULL) {
1252 1
        if (PyArray_DESCR(op)->f->compare) {
1253 1
            switch (which) {
1254
                default:
1255
                case NPY_QUICKSORT:
1256
                    sort = npy_quicksort;
1257
                    break;
1258 1
                case NPY_HEAPSORT:
1259 1
                    sort = npy_heapsort;
1260 1
                    break;
1261 1
                case NPY_STABLESORT:
1262 1
                    sort = npy_timsort;
1263 1
                    break;
1264
            }
1265
        }
1266
        else {
1267 0
            PyErr_SetString(PyExc_TypeError,
1268
                            "type does not have compare function");
1269 0
            return -1;
1270
        }
1271
    }
1272

1273 1
    return _new_sortlike(op, axis, sort, NULL, NULL, 0);
1274
}
1275

1276

1277
/*
1278
 * make kth array positive, ravel and sort it
1279
 */
1280
static PyArrayObject *
1281 1
partition_prep_kth_array(PyArrayObject * ktharray,
1282
                         PyArrayObject * op,
1283
                         int axis)
1284
{
1285 1
    const npy_intp * shape = PyArray_SHAPE(op);
1286
    PyArrayObject * kthrvl;
1287
    npy_intp * kth;
1288
    npy_intp nkth, i;
1289

1290 1
    if (!PyArray_CanCastSafely(PyArray_TYPE(ktharray), NPY_INTP)) {
1291 1
        PyErr_Format(PyExc_TypeError, "Partition index must be integer");
1292 1
        return NULL;
1293
    }
1294

1295 1
    if (PyArray_NDIM(ktharray) > 1) {
1296 0
        PyErr_Format(PyExc_ValueError, "kth array must have dimension <= 1");
1297 0
        return NULL;
1298
    }
1299 1
    kthrvl = (PyArrayObject *)PyArray_Cast(ktharray, NPY_INTP);
1300

1301 1
    if (kthrvl == NULL)
1302
        return NULL;
1303

1304 1
    kth = PyArray_DATA(kthrvl);
1305 1
    nkth = PyArray_SIZE(kthrvl);
1306

1307 1
    for (i = 0; i < nkth; i++) {
1308 1
        if (kth[i] < 0) {
1309 1
            kth[i] += shape[axis];
1310
        }
1311 1
        if (PyArray_SIZE(op) != 0 &&
1312 1
                    (kth[i] < 0 || kth[i] >= shape[axis])) {
1313 1
            PyErr_Format(PyExc_ValueError, "kth(=%zd) out of bounds (%zd)",
1314 1
                         kth[i], shape[axis]);
1315 1
            Py_XDECREF(kthrvl);
1316
            return NULL;
1317
        }
1318
    }
1319

1320
    /*
1321
     * sort the array of kths so the partitions will
1322
     * not trample on each other
1323
     */
1324 1
    if (PyArray_SIZE(kthrvl) > 1) {
1325 1
        PyArray_Sort(kthrvl, -1, NPY_QUICKSORT);
1326
    }
1327

1328
    return kthrvl;
1329
}
1330

1331

1332
/*NUMPY_API
1333
 * Partition an array in-place
1334
 */
1335
NPY_NO_EXPORT int
1336 1
PyArray_Partition(PyArrayObject *op, PyArrayObject * ktharray, int axis,
1337
                  NPY_SELECTKIND which)
1338
{
1339
    PyArrayObject *kthrvl;
1340
    PyArray_PartitionFunc *part;
1341
    PyArray_SortFunc *sort;
1342 1
    int n = PyArray_NDIM(op);
1343
    int ret;
1344

1345 1
    if (check_and_adjust_axis(&axis, n) < 0) {
1346
        return -1;
1347
    }
1348

1349 1
    if (PyArray_FailUnlessWriteable(op, "partition array") < 0) {
1350
        return -1;
1351
    }
1352

1353 1
    if (which < 0 || which >= NPY_NSELECTS) {
1354 0
        PyErr_SetString(PyExc_ValueError, "not a valid partition kind");
1355 0
        return -1;
1356
    }
1357 1
    part = get_partition_func(PyArray_TYPE(op), which);
1358 1
    if (part == NULL) {
1359
        /* Use sorting, slower but equivalent */
1360 1
        if (PyArray_DESCR(op)->f->compare) {
1361
            sort = npy_quicksort;
1362
        }
1363
        else {
1364 0
            PyErr_SetString(PyExc_TypeError,
1365
                            "type does not have compare function");
1366 0
            return -1;
1367
        }
1368
    }
1369

1370
    /* Process ktharray even if using sorting to do bounds checking */
1371 1
    kthrvl = partition_prep_kth_array(ktharray, op, axis);
1372 1
    if (kthrvl == NULL) {
1373
        return -1;
1374
    }
1375

1376 1
    ret = _new_sortlike(op, axis, sort, part,
1377 1
                        PyArray_DATA(kthrvl), PyArray_SIZE(kthrvl));
1378

1379 1
    Py_DECREF(kthrvl);
1380

1381
    return ret;
1382
}
1383

1384

1385
/*NUMPY_API
1386
 * ArgSort an array
1387
 */
1388
NPY_NO_EXPORT PyObject *
1389 1
PyArray_ArgSort(PyArrayObject *op, int axis, NPY_SORTKIND which)
1390
{
1391
    PyArrayObject *op2;
1392 1
    PyArray_ArgSortFunc *argsort = NULL;
1393
    PyObject *ret;
1394

1395 1
    argsort = PyArray_DESCR(op)->f->argsort[which];
1396

1397 1
    if (argsort == NULL) {
1398 1
        if (PyArray_DESCR(op)->f->compare) {
1399 1
            switch (which) {
1400
                default:
1401
                case NPY_QUICKSORT:
1402
                    argsort = npy_aquicksort;
1403
                    break;
1404 1
                case NPY_HEAPSORT:
1405 1
                    argsort = npy_aheapsort;
1406 1
                    break;
1407 1
                case NPY_STABLESORT:
1408 1
                    argsort = npy_atimsort;
1409 1
                    break;
1410
            }
1411
        }
1412
        else {
1413 0
            PyErr_SetString(PyExc_TypeError,
1414
                            "type does not have compare function");
1415 0
            return NULL;
1416
        }
1417
    }
1418

1419 1
    op2 = (PyArrayObject *)PyArray_CheckAxis(op, &axis, 0);
1420 1
    if (op2 == NULL) {
1421
        return NULL;
1422
    }
1423

1424 1
    ret = _new_argsortlike(op2, axis, argsort, NULL, NULL, 0);
1425

1426 1
    Py_DECREF(op2);
1427
    return ret;
1428
}
1429

1430

1431
/*NUMPY_API
1432
 * ArgPartition an array
1433
 */
1434
NPY_NO_EXPORT PyObject *
1435 1
PyArray_ArgPartition(PyArrayObject *op, PyArrayObject *ktharray, int axis,
1436
                     NPY_SELECTKIND which)
1437
{
1438
    PyArrayObject *op2, *kthrvl;
1439
    PyArray_ArgPartitionFunc *argpart;
1440
    PyArray_ArgSortFunc *argsort;
1441
    PyObject *ret;
1442

1443
    /*
1444
     * As a C-exported function, enum NPY_SELECTKIND loses its enum property
1445
     * Check the values to make sure they are in range
1446
     */
1447 1
    if ((int)which < 0 || (int)which >= NPY_NSELECTS) {
1448 0
        PyErr_SetString(PyExc_ValueError,
1449
                        "not a valid partition kind");
1450 0
        return NULL;
1451
    }
1452

1453 1
    argpart = get_argpartition_func(PyArray_TYPE(op), which);
1454 1
    if (argpart == NULL) {
1455
        /* Use sorting, slower but equivalent */
1456 1
        if (PyArray_DESCR(op)->f->compare) {
1457
            argsort = npy_aquicksort;
1458
        }
1459
        else {
1460 0
            PyErr_SetString(PyExc_TypeError,
1461
                            "type does not have compare function");
1462 0
            return NULL;
1463
        }
1464
    }
1465

1466 1
    op2 = (PyArrayObject *)PyArray_CheckAxis(op, &axis, 0);
1467 1
    if (op2 == NULL) {
1468
        return NULL;
1469
    }
1470

1471
    /* Process ktharray even if using sorting to do bounds checking */
1472 1
    kthrvl = partition_prep_kth_array(ktharray, op2, axis);
1473 1
    if (kthrvl == NULL) {
1474 1
        Py_DECREF(op2);
1475
        return NULL;
1476
    }
1477

1478 1
    ret = _new_argsortlike(op2, axis, argsort, argpart,
1479 1
                           PyArray_DATA(kthrvl), PyArray_SIZE(kthrvl));
1480

1481 1
    Py_DECREF(kthrvl);
1482 1
    Py_DECREF(op2);
1483

1484
    return ret;
1485
}
1486

1487

1488
/*NUMPY_API
1489
 *LexSort an array providing indices that will sort a collection of arrays
1490
 *lexicographically.  The first key is sorted on first, followed by the second key
1491
 *-- requires that arg"merge"sort is available for each sort_key
1492
 *
1493
 *Returns an index array that shows the indexes for the lexicographic sort along
1494
 *the given axis.
1495
 */
1496
NPY_NO_EXPORT PyObject *
1497 1
PyArray_LexSort(PyObject *sort_keys, int axis)
1498
{
1499
    PyArrayObject **mps;
1500
    PyArrayIterObject **its;
1501 1
    PyArrayObject *ret = NULL;
1502 1
    PyArrayIterObject *rit = NULL;
1503
    npy_intp n, N, size, i, j;
1504
    npy_intp astride, rstride, *iptr;
1505
    int nd;
1506 1
    int needcopy = 0;
1507
    int elsize;
1508
    int maxelsize;
1509 1
    int object = 0;
1510
    PyArray_ArgSortFunc *argsort;
1511 1
    NPY_BEGIN_THREADS_DEF;
1512

1513 1
    if (!PySequence_Check(sort_keys)
1514 1
           || ((n = PySequence_Size(sort_keys)) <= 0)) {
1515 0
        PyErr_SetString(PyExc_TypeError,
1516
                "need sequence of keys with len > 0 in lexsort");
1517 0
        return NULL;
1518
    }
1519 1
    mps = (PyArrayObject **) PyArray_malloc(n * sizeof(PyArrayObject *));
1520 1
    if (mps == NULL) {
1521 0
        return PyErr_NoMemory();
1522
    }
1523 1
    its = (PyArrayIterObject **) PyArray_malloc(n * sizeof(PyArrayIterObject *));
1524 1
    if (its == NULL) {
1525 0
        PyArray_free(mps);
1526 0
        return PyErr_NoMemory();
1527
    }
1528 1
    for (i = 0; i < n; i++) {
1529 1
        mps[i] = NULL;
1530 1
        its[i] = NULL;
1531
    }
1532 1
    for (i = 0; i < n; i++) {
1533
        PyObject *obj;
1534 1
        obj = PySequence_GetItem(sort_keys, i);
1535 1
        if (obj == NULL) {
1536
            goto fail;
1537
        }
1538 1
        mps[i] = (PyArrayObject *)PyArray_FROM_O(obj);
1539 1
        Py_DECREF(obj);
1540 1
        if (mps[i] == NULL) {
1541
            goto fail;
1542
        }
1543 1
        if (i > 0) {
1544 1
            if ((PyArray_NDIM(mps[i]) != PyArray_NDIM(mps[0]))
1545 1
                || (!PyArray_CompareLists(PyArray_DIMS(mps[i]),
1546 1
                                       PyArray_DIMS(mps[0]),
1547
                                       PyArray_NDIM(mps[0])))) {
1548 0
                PyErr_SetString(PyExc_ValueError,
1549
                                "all keys need to be the same shape");
1550 0
                goto fail;
1551
            }
1552
        }
1553 1
        if (!PyArray_DESCR(mps[i])->f->argsort[NPY_STABLESORT]
1554 1
                && !PyArray_DESCR(mps[i])->f->compare) {
1555 0
            PyErr_Format(PyExc_TypeError,
1556
                         "item %zd type does not have compare function", i);
1557 0
            goto fail;
1558
        }
1559 1
        if (!object
1560 1
            && PyDataType_FLAGCHK(PyArray_DESCR(mps[i]), NPY_NEEDS_PYAPI)) {
1561 1
            object = 1;
1562
        }
1563
    }
1564

1565
    /* Now we can check the axis */
1566 1
    nd = PyArray_NDIM(mps[0]);
1567
    /*
1568
    * Special case letting axis={-1,0} slip through for scalars,
1569
    * for backwards compatibility reasons.
1570
    */
1571 1
    if (nd == 0 && (axis == 0 || axis == -1)) {
1572
        /* TODO: can we deprecate this? */
1573
    }
1574 1
    else if (check_and_adjust_axis(&axis, nd) < 0) {
1575
        goto fail;
1576
    }
1577 1
    if ((nd == 0) || (PyArray_SIZE(mps[0]) <= 1)) {
1578
        /* empty/single element case */
1579 1
        ret = (PyArrayObject *)PyArray_NewFromDescr(
1580
            &PyArray_Type, PyArray_DescrFromType(NPY_INTP),
1581 1
            PyArray_NDIM(mps[0]), PyArray_DIMS(mps[0]), NULL, NULL,
1582
            0, NULL);
1583

1584 1
        if (ret == NULL) {
1585
            goto fail;
1586
        }
1587 1
        if (PyArray_SIZE(mps[0]) > 0) {
1588 1
            *((npy_intp *)(PyArray_DATA(ret))) = 0;
1589
        }
1590
        goto finish;
1591
    }
1592

1593 1
    for (i = 0; i < n; i++) {
1594 1
        its[i] = (PyArrayIterObject *)PyArray_IterAllButAxis(
1595 1
                (PyObject *)mps[i], &axis);
1596 1
        if (its[i] == NULL) {
1597
            goto fail;
1598
        }
1599
    }
1600

1601
    /* Now do the sorting */
1602 1
    ret = (PyArrayObject *)PyArray_NewFromDescr(
1603
            &PyArray_Type, PyArray_DescrFromType(NPY_INTP),
1604 1
            PyArray_NDIM(mps[0]), PyArray_DIMS(mps[0]), NULL, NULL,
1605
            0, NULL);
1606 1
    if (ret == NULL) {
1607
        goto fail;
1608
    }
1609 1
    rit = (PyArrayIterObject *)
1610
            PyArray_IterAllButAxis((PyObject *)ret, &axis);
1611 1
    if (rit == NULL) {
1612
        goto fail;
1613
    }
1614 1
    if (!object) {
1615 1
        NPY_BEGIN_THREADS;
1616
    }
1617 1
    size = rit->size;
1618 1
    N = PyArray_DIMS(mps[0])[axis];
1619 1
    rstride = PyArray_STRIDE(ret, axis);
1620 1
    maxelsize = PyArray_DESCR(mps[0])->elsize;
1621 1
    needcopy = (rstride != sizeof(npy_intp));
1622 1
    for (j = 0; j < n; j++) {
1623 1
        needcopy = needcopy
1624 1
            || PyArray_ISBYTESWAPPED(mps[j])
1625 1
            || !(PyArray_FLAGS(mps[j]) & NPY_ARRAY_ALIGNED)
1626 1
            || (PyArray_STRIDES(mps[j])[axis] != (npy_intp)PyArray_DESCR(mps[j])->elsize);
1627 1
        if (PyArray_DESCR(mps[j])->elsize > maxelsize) {
1628 1
            maxelsize = PyArray_DESCR(mps[j])->elsize;
1629
        }
1630
    }
1631

1632 1
    if (needcopy) {
1633
        char *valbuffer, *indbuffer;
1634
        int *swaps;
1635

1636
        assert(N > 0);  /* Guaranteed and assumed by indbuffer */
1637 1
        npy_intp valbufsize = N * maxelsize;
1638 1
        if (NPY_UNLIKELY(valbufsize) == 0) {
1639 0
            valbufsize = 1;  /* Ensure allocation is not empty */
1640
        }
1641

1642 1
        valbuffer = PyDataMem_NEW(valbufsize);
1643 1
        if (valbuffer == NULL) {
1644
            goto fail;
1645
        }
1646 1
        indbuffer = PyDataMem_NEW(N * sizeof(npy_intp));
1647 1
        if (indbuffer == NULL) {
1648 0
            PyDataMem_FREE(valbuffer);
1649 0
            goto fail;
1650
        }
1651 1
        swaps = malloc(NPY_LIKELY(n > 0) ? n * sizeof(int) : 1);
1652 1
        if (swaps == NULL) {
1653 0
            PyDataMem_FREE(valbuffer);
1654 0
            PyDataMem_FREE(indbuffer);
1655 0
            goto fail;
1656
        }
1657

1658 1
        for (j = 0; j < n; j++) {
1659 1
            swaps[j] = PyArray_ISBYTESWAPPED(mps[j]);
1660
        }
1661 1
        while (size--) {
1662
            iptr = (npy_intp *)indbuffer;
1663 1
            for (i = 0; i < N; i++) {
1664 1
                *iptr++ = i;
1665
            }
1666 1
            for (j = 0; j < n; j++) {
1667
                int rcode;
1668 1
                elsize = PyArray_DESCR(mps[j])->elsize;
1669 1
                astride = PyArray_STRIDES(mps[j])[axis];
1670 1
                argsort = PyArray_DESCR(mps[j])->f->argsort[NPY_STABLESORT];
1671 1
                if(argsort == NULL) {
1672 0
                    argsort = npy_atimsort;
1673
                }
1674 1
                _unaligned_strided_byte_copy(valbuffer, (npy_intp) elsize,
1675 1
                                             its[j]->dataptr, astride, N, elsize);
1676 1
                if (swaps[j]) {
1677 0
                    _strided_byte_swap(valbuffer, (npy_intp) elsize, N, elsize);
1678
                }
1679 1
                rcode = argsort(valbuffer, (npy_intp *)indbuffer, N, mps[j]);
1680 1
                if (rcode < 0 || (PyDataType_REFCHK(PyArray_DESCR(mps[j]))
1681 0
                            && PyErr_Occurred())) {
1682 0
                    PyDataMem_FREE(valbuffer);
1683 0
                    PyDataMem_FREE(indbuffer);
1684 0
                    free(swaps);
1685 0
                    goto fail;
1686
                }
1687 1
                PyArray_ITER_NEXT(its[j]);
1688
            }
1689 1
            _unaligned_strided_byte_copy(rit->dataptr, rstride, indbuffer,
1690
                                         sizeof(npy_intp), N, sizeof(npy_intp));
1691 1
            PyArray_ITER_NEXT(rit);
1692
        }
1693 1
        PyDataMem_FREE(valbuffer);
1694 1
        PyDataMem_FREE(indbuffer);
1695 1
        free(swaps);
1696
    }
1697
    else {
1698 1
        while (size--) {
1699 1
            iptr = (npy_intp *)rit->dataptr;
1700 1
            for (i = 0; i < N; i++) {
1701 1
                *iptr++ = i;
1702
            }
1703 1
            for (j = 0; j < n; j++) {
1704
                int rcode;
1705 1
                argsort = PyArray_DESCR(mps[j])->f->argsort[NPY_STABLESORT];
1706 1
                if(argsort == NULL) {
1707 1
                    argsort = npy_atimsort;
1708
                }
1709 1
                rcode = argsort(its[j]->dataptr,
1710 1
                        (npy_intp *)rit->dataptr, N, mps[j]);
1711 1
                if (rcode < 0 || (PyDataType_REFCHK(PyArray_DESCR(mps[j]))
1712 1
                            && PyErr_Occurred())) {
1713
                    goto fail;
1714
                }
1715 1
                PyArray_ITER_NEXT(its[j]);
1716
            }
1717 1
            PyArray_ITER_NEXT(rit);
1718
        }
1719
    }
1720

1721 1
    if (!object) {
1722 1
        NPY_END_THREADS;
1723
    }
1724

1725 1
 finish:
1726 1
    for (i = 0; i < n; i++) {
1727 1
        Py_XDECREF(mps[i]);
1728 1
        Py_XDECREF(its[i]);
1729
    }
1730 1
    Py_XDECREF(rit);
1731 1
    PyArray_free(mps);
1732 1
    PyArray_free(its);
1733 1
    return (PyObject *)ret;
1734

1735 1
 fail:
1736 1
    NPY_END_THREADS;
1737 1
    if (!PyErr_Occurred()) {
1738
        /* Out of memory during sorting or buffer creation */
1739 0
        PyErr_NoMemory();
1740
    }
1741 1
    Py_XDECREF(rit);
1742 1
    Py_XDECREF(ret);
1743 1
    for (i = 0; i < n; i++) {
1744 1
        Py_XDECREF(mps[i]);
1745 1
        Py_XDECREF(its[i]);
1746
    }
1747 1
    PyArray_free(mps);
1748 1
    PyArray_free(its);
1749 1
    return NULL;
1750
}
1751

1752

1753
/*NUMPY_API
1754
 *
1755
 * Search the sorted array op1 for the location of the items in op2. The
1756
 * result is an array of indexes, one for each element in op2, such that if
1757
 * the item were to be inserted in op1 just before that index the array
1758
 * would still be in sorted order.
1759
 *
1760
 * Parameters
1761
 * ----------
1762
 * op1 : PyArrayObject *
1763
 *     Array to be searched, must be 1-D.
1764
 * op2 : PyObject *
1765
 *     Array of items whose insertion indexes in op1 are wanted
1766
 * side : {NPY_SEARCHLEFT, NPY_SEARCHRIGHT}
1767
 *     If NPY_SEARCHLEFT, return first valid insertion indexes
1768
 *     If NPY_SEARCHRIGHT, return last valid insertion indexes
1769
 * perm : PyObject *
1770
 *     Permutation array that sorts op1 (optional)
1771
 *
1772
 * Returns
1773
 * -------
1774
 * ret : PyObject *
1775
 *   New reference to npy_intp array containing indexes where items in op2
1776
 *   could be validly inserted into op1. NULL on error.
1777
 *
1778
 * Notes
1779
 * -----
1780
 * Binary search is used to find the indexes.
1781
 */
1782
NPY_NO_EXPORT PyObject *
1783 1
PyArray_SearchSorted(PyArrayObject *op1, PyObject *op2,
1784
                     NPY_SEARCHSIDE side, PyObject *perm)
1785
{
1786 1
    PyArrayObject *ap1 = NULL;
1787 1
    PyArrayObject *ap2 = NULL;
1788 1
    PyArrayObject *ap3 = NULL;
1789 1
    PyArrayObject *sorter = NULL;
1790 1
    PyArrayObject *ret = NULL;
1791
    PyArray_Descr *dtype;
1792 1
    int ap1_flags = NPY_ARRAY_NOTSWAPPED | NPY_ARRAY_ALIGNED;
1793 1
    PyArray_BinSearchFunc *binsearch = NULL;
1794 1
    PyArray_ArgBinSearchFunc *argbinsearch = NULL;
1795 1
    NPY_BEGIN_THREADS_DEF;
1796

1797
    /* Find common type */
1798 1
    dtype = PyArray_DescrFromObject((PyObject *)op2, PyArray_DESCR(op1));
1799 1
    if (dtype == NULL) {
1800
        return NULL;
1801
    }
1802
    /* refs to dtype we own = 1 */
1803

1804
    /* Look for binary search function */
1805 1
    if (perm) {
1806 1
        argbinsearch = get_argbinsearch_func(dtype, side);
1807
    }
1808
    else {
1809 1
        binsearch = get_binsearch_func(dtype, side);
1810
    }
1811 1
    if (binsearch == NULL && argbinsearch == NULL) {
1812 0
        PyErr_SetString(PyExc_TypeError, "compare not supported for type");
1813
        /* refs to dtype we own = 1 */
1814 0
        Py_DECREF(dtype);
1815
        /* refs to dtype we own = 0 */
1816
        return NULL;
1817
    }
1818

1819
    /* need ap2 as contiguous array and of right type */
1820
    /* refs to dtype we own = 1 */
1821 1
    Py_INCREF(dtype);
1822
    /* refs to dtype we own = 2 */
1823 1
    ap2 = (PyArrayObject *)PyArray_CheckFromAny(op2, dtype,
1824
                                0, 0,
1825
                                NPY_ARRAY_CARRAY_RO | NPY_ARRAY_NOTSWAPPED,
1826
                                NULL);
1827
    /* refs to dtype we own = 1, array creation steals one even on failure */
1828 1
    if (ap2 == NULL) {
1829 0
        Py_DECREF(dtype);
1830
        /* refs to dtype we own = 0 */
1831
        return NULL;
1832
    }
1833

1834
    /*
1835
     * If the needle (ap2) is larger than the haystack (op1) we copy the
1836
     * haystack to a contiguous array for improved cache utilization.
1837
     */
1838 1
    if (PyArray_SIZE(ap2) > PyArray_SIZE(op1)) {
1839 1
        ap1_flags |= NPY_ARRAY_CARRAY_RO;
1840
    }
1841 1
    ap1 = (PyArrayObject *)PyArray_CheckFromAny((PyObject *)op1, dtype,
1842
                                1, 1, ap1_flags, NULL);
1843
    /* refs to dtype we own = 0, array creation steals one even on failure */
1844 1
    if (ap1 == NULL) {
1845
        goto fail;
1846
    }
1847

1848 1
    if (perm) {
1849
        /* need ap3 as a 1D aligned, not swapped, array of right type */
1850 1
        ap3 = (PyArrayObject *)PyArray_CheckFromAny(perm, NULL,
1851
                                    1, 1,
1852
                                    NPY_ARRAY_ALIGNED | NPY_ARRAY_NOTSWAPPED,
1853
                                    NULL);
1854 1
        if (ap3 == NULL) {
1855 0
            PyErr_SetString(PyExc_TypeError,
1856
                        "could not parse sorter argument");
1857 0
            goto fail;
1858
        }
1859 1
        if (!PyArray_ISINTEGER(ap3)) {
1860 1
            PyErr_SetString(PyExc_TypeError,
1861
                        "sorter must only contain integers");
1862 1
            goto fail;
1863
        }
1864
        /* convert to known integer size */
1865 1
        sorter = (PyArrayObject *)PyArray_FromArray(ap3,
1866
                                    PyArray_DescrFromType(NPY_INTP),
1867
                                    NPY_ARRAY_ALIGNED | NPY_ARRAY_NOTSWAPPED);
1868 1
        if (sorter == NULL) {
1869 0
            PyErr_SetString(PyExc_ValueError,
1870
                        "could not parse sorter argument");
1871 0
            goto fail;
1872
        }
1873 1
        if (PyArray_SIZE(sorter) != PyArray_SIZE(ap1)) {
1874 1
            PyErr_SetString(PyExc_ValueError,
1875
                        "sorter.size must equal a.size");
1876 1
            goto fail;
1877
        }
1878
    }
1879

1880
    /* ret is a contiguous array of intp type to hold returned indexes */
1881 1
    ret = (PyArrayObject *)PyArray_NewFromDescr(
1882
            &PyArray_Type, PyArray_DescrFromType(NPY_INTP),
1883 1
            PyArray_NDIM(ap2), PyArray_DIMS(ap2), NULL, NULL,
1884
            0, (PyObject *)ap2);
1885 1
    if (ret == NULL) {
1886
        goto fail;
1887
    }
1888

1889 1
    if (ap3 == NULL) {
1890
        /* do regular binsearch */
1891 1
        NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(ap2));
1892 1
        binsearch((const char *)PyArray_DATA(ap1),
1893 1
                  (const char *)PyArray_DATA(ap2),
1894 1
                  (char *)PyArray_DATA(ret),
1895 1
                  PyArray_SIZE(ap1), PyArray_SIZE(ap2),
1896 1
                  PyArray_STRIDES(ap1)[0], PyArray_DESCR(ap2)->elsize,
1897
                  NPY_SIZEOF_INTP, ap2);
1898 1
        NPY_END_THREADS_DESCR(PyArray_DESCR(ap2));
1899
    }
1900
    else {
1901
        /* do binsearch with a sorter array */
1902 1
        int error = 0;
1903 1
        NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(ap2));
1904 1
        error = argbinsearch((const char *)PyArray_DATA(ap1),
1905 1
                             (const char *)PyArray_DATA(ap2),
1906 1
                             (const char *)PyArray_DATA(sorter),
1907 1
                             (char *)PyArray_DATA(ret),
1908 1
                             PyArray_SIZE(ap1), PyArray_SIZE(ap2),
1909 1
                             PyArray_STRIDES(ap1)[0],
1910 1
                             PyArray_DESCR(ap2)->elsize,
1911 1
                             PyArray_STRIDES(sorter)[0], NPY_SIZEOF_INTP, ap2);
1912 1
        NPY_END_THREADS_DESCR(PyArray_DESCR(ap2));
1913 1
        if (error < 0) {
1914 1
            PyErr_SetString(PyExc_ValueError,
1915
                        "Sorter index out of range.");
1916 1
            goto fail;
1917
        }
1918 1
        Py_DECREF(ap3);
1919 1
        Py_DECREF(sorter);
1920
    }
1921 1
    Py_DECREF(ap1);
1922 1
    Py_DECREF(ap2);
1923
    return (PyObject *)ret;
1924

1925 0
 fail:
1926 1
    Py_XDECREF(ap1);
1927 1
    Py_XDECREF(ap2);
1928 1
    Py_XDECREF(ap3);
1929 1
    Py_XDECREF(sorter);
1930 1
    Py_XDECREF(ret);
1931
    return NULL;
1932
}
1933

1934
/*NUMPY_API
1935
 * Diagonal
1936
 *
1937
 * In NumPy versions prior to 1.7,  this function always returned a copy of
1938
 * the diagonal array. In 1.7, the code has been updated to compute a view
1939
 * onto 'self', but it still copies this array before returning, as well as
1940
 * setting the internal WARN_ON_WRITE flag. In a future version, it will
1941
 * simply return a view onto self.
1942
 */
1943
NPY_NO_EXPORT PyObject *
1944 1
PyArray_Diagonal(PyArrayObject *self, int offset, int axis1, int axis2)
1945
{
1946 1
    int i, idim, ndim = PyArray_NDIM(self);
1947
    npy_intp *strides;
1948
    npy_intp stride1, stride2, offset_stride;
1949
    npy_intp *shape, dim1, dim2;
1950

1951
    char *data;
1952
    npy_intp diag_size;
1953
    PyArray_Descr *dtype;
1954
    PyObject *ret;
1955
    npy_intp ret_shape[NPY_MAXDIMS], ret_strides[NPY_MAXDIMS];
1956

1957 1
    if (ndim < 2) {
1958 1
        PyErr_SetString(PyExc_ValueError,
1959
                        "diag requires an array of at least two dimensions");
1960 1
        return NULL;
1961
    }
1962

1963
    /* Handle negative axes with standard Python indexing rules */
1964 1
    if (check_and_adjust_axis_msg(&axis1, ndim, npy_ma_str_axis1) < 0) {
1965
        return NULL;
1966
    }
1967 1
    if (check_and_adjust_axis_msg(&axis2, ndim, npy_ma_str_axis2) < 0) {
1968
        return NULL;
1969
    }
1970 1
    if (axis1 == axis2) {
1971 1
        PyErr_SetString(PyExc_ValueError,
1972
                    "axis1 and axis2 cannot be the same");
1973 1
        return NULL;
1974
    }
1975

1976
    /* Get the shape and strides of the two axes */
1977 1
    shape = PyArray_SHAPE(self);
1978 1
    dim1 = shape[axis1];
1979 1
    dim2 = shape[axis2];
1980 1
    strides = PyArray_STRIDES(self);
1981 1
    stride1 = strides[axis1];
1982 1
    stride2 = strides[axis2];
1983

1984
    /* Compute the data pointers and diag_size for the view */
1985 1
    data = PyArray_DATA(self);
1986 1
    if (offset >= 0) {
1987 1
        offset_stride = stride2;
1988 1
        dim2 -= offset;
1989
    }
1990
    else {
1991 1
        offset = -offset;
1992 1
        offset_stride = stride1;
1993 1
        dim1 -= offset;
1994
    }
1995 1
    diag_size = dim2 < dim1 ? dim2 : dim1;
1996 1
    if (diag_size < 0) {
1997
        diag_size = 0;
1998
    }
1999
    else {
2000 1
        data += offset * offset_stride;
2001
    }
2002

2003
    /* Build the new shape and strides for the main data */
2004 1
    i = 0;
2005 1
    for (idim = 0; idim < ndim; ++idim) {
2006 1
        if (idim != axis1 && idim != axis2) {
2007 1
            ret_shape[i] = shape[idim];
2008 1
            ret_strides[i] = strides[idim];
2009 1
            ++i;
2010
        }
2011
    }
2012 1
    ret_shape[ndim-2] = diag_size;
2013 1
    ret_strides[ndim-2] = stride1 + stride2;
2014

2015
    /* Create the diagonal view */
2016 1
    dtype = PyArray_DTYPE(self);
2017 1
    Py_INCREF(dtype);
2018 1
    ret = PyArray_NewFromDescrAndBase(
2019 1
            Py_TYPE(self), dtype,
2020
            ndim-1, ret_shape, ret_strides, data,
2021
            PyArray_FLAGS(self), (PyObject *)self, (PyObject *)self);
2022 1
    if (ret == NULL) {
2023
        return NULL;
2024
    }
2025

2026
    /*
2027
     * For numpy 1.9 the diagonal view is not writeable.
2028
     * This line needs to be removed in 1.10.
2029
     */
2030 1
    PyArray_CLEARFLAGS((PyArrayObject *)ret, NPY_ARRAY_WRITEABLE);
2031

2032 1
    return ret;
2033
}
2034

2035
/*NUMPY_API
2036
 * Compress
2037
 */
2038
NPY_NO_EXPORT PyObject *
2039 1
PyArray_Compress(PyArrayObject *self, PyObject *condition, int axis,
2040
                 PyArrayObject *out)
2041
{
2042
    PyArrayObject *cond;
2043
    PyObject *res, *ret;
2044

2045 1
    if (PyArray_Check(condition)) {
2046 1
        cond = (PyArrayObject *)condition;
2047 1
        Py_INCREF(cond);
2048
    }
2049
    else {
2050 1
        PyArray_Descr *dtype = PyArray_DescrFromType(NPY_BOOL);
2051 1
        if (dtype == NULL) {
2052
            return NULL;
2053
        }
2054 1
        cond = (PyArrayObject *)PyArray_FromAny(condition, dtype,
2055
                                    0, 0, 0, NULL);
2056 1
        if (cond == NULL) {
2057
            return NULL;
2058
        }
2059
    }
2060

2061 1
    if (PyArray_NDIM(cond) != 1) {
2062 0
        Py_DECREF(cond);
2063 0
        PyErr_SetString(PyExc_ValueError,
2064
                        "condition must be a 1-d array");
2065 0
        return NULL;
2066
    }
2067

2068 1
    res = PyArray_Nonzero(cond);
2069 1
    Py_DECREF(cond);
2070 1
    if (res == NULL) {
2071
        return res;
2072
    }
2073 1
    ret = PyArray_TakeFrom(self, PyTuple_GET_ITEM(res, 0), axis,
2074
                           out, NPY_RAISE);
2075 1
    Py_DECREF(res);
2076
    return ret;
2077
}
2078

2079
/*
2080
 * count number of nonzero bytes in 48 byte block
2081
 * w must be aligned to 8 bytes
2082
 *
2083
 * even though it uses 64 bit types its faster than the bytewise sum on 32 bit
2084
 * but a 32 bit type version would make it even faster on these platforms
2085
 */
2086
static NPY_INLINE npy_intp
2087 1
count_nonzero_bytes_384(const npy_uint64 * w)
2088
{
2089 1
    const npy_uint64 w1 = w[0];
2090 1
    const npy_uint64 w2 = w[1];
2091 1
    const npy_uint64 w3 = w[2];
2092 1
    const npy_uint64 w4 = w[3];
2093 1
    const npy_uint64 w5 = w[4];
2094 1
    const npy_uint64 w6 = w[5];
2095
    npy_intp r;
2096

2097
    /*
2098
     * last part of sideways add popcount, first three bisections can be
2099
     * skipped as we are dealing with bytes.
2100
     * multiplication equivalent to (x + (x>>8) + (x>>16) + (x>>24)) & 0xFF
2101
     * multiplication overflow well defined for unsigned types.
2102
     * w1 + w2 guaranteed to not overflow as we only have 0 and 1 data.
2103
     */
2104 1
    r = ((w1 + w2 + w3 + w4 + w5 + w6) * 0x0101010101010101ULL) >> 56ULL;
2105

2106
    /*
2107
     * bytes not exclusively 0 or 1, sum them individually.
2108
     * should only happen if one does weird stuff with views or external
2109
     * buffers.
2110
     * Doing this after the optimistic computation allows saving registers and
2111
     * better pipelining
2112
     */
2113 1
    if (NPY_UNLIKELY(
2114
             ((w1 | w2 | w3 | w4 | w5 | w6) & 0xFEFEFEFEFEFEFEFEULL) != 0)) {
2115
        /* reload from pointer to avoid a unnecessary stack spill with gcc */
2116
        const char * c = (const char *)w;
2117
        npy_uintp i, count = 0;
2118 0
        for (i = 0; i < 48; i++) {
2119 0
            count += (c[i] != 0);
2120
        }
2121 0
        return count;
2122
    }
2123

2124
    return r;
2125
}
2126

2127
/*
2128
 * Counts the number of True values in a raw boolean array. This
2129
 * is a low-overhead function which does no heap allocations.
2130
 *
2131
 * Returns -1 on error.
2132
 */
2133
NPY_NO_EXPORT npy_intp
2134 1
count_boolean_trues(int ndim, char *data, npy_intp const *ashape, npy_intp const *astrides)
2135
{
2136
    int idim;
2137
    npy_intp shape[NPY_MAXDIMS], strides[NPY_MAXDIMS];
2138
    npy_intp i, coord[NPY_MAXDIMS];
2139 1
    npy_intp count = 0;
2140 1
    NPY_BEGIN_THREADS_DEF;
2141

2142
    /* Use raw iteration with no heap memory allocation */
2143 1
    if (PyArray_PrepareOneRawArrayIter(
2144
                    ndim, ashape,
2145
                    data, astrides,
2146
                    &ndim, shape,
2147
                    &data, strides) < 0) {
2148
        return -1;
2149
    }
2150

2151
    /* Handle zero-sized array */
2152 1
    if (shape[0] == 0) {
2153
        return 0;
2154
    }
2155

2156 1
    NPY_BEGIN_THREADS_THRESHOLDED(shape[0]);
2157

2158
    /* Special case for contiguous inner loop */
2159 1
    if (strides[0] == 1) {
2160 1
        NPY_RAW_ITER_START(idim, ndim, coord, shape) {
2161
            /* Process the innermost dimension */
2162 1
            const char *d = data;
2163 1
            const char *e = data + shape[0];
2164
            if (NPY_CPU_HAVE_UNALIGNED_ACCESS ||
2165
                    npy_is_aligned(d, sizeof(npy_uint64))) {
2166 1
                npy_uintp stride = 6 * sizeof(npy_uint64);
2167 1
                for (; d < e - (shape[0] % stride); d += stride) {
2168 1
                    count += count_nonzero_bytes_384((const npy_uint64 *)d);
2169
                }
2170
            }
2171 1
            for (; d < e; ++d) {
2172 1
                count += (*d != 0);
2173
            }
2174 0
        } NPY_RAW_ITER_ONE_NEXT(idim, ndim, coord, shape, data, strides);
2175
    }
2176
    /* General inner loop */
2177
    else {
2178 1
        NPY_RAW_ITER_START(idim, ndim, coord, shape) {
2179 1
            char *d = data;
2180
            /* Process the innermost dimension */
2181 1
            for (i = 0; i < shape[0]; ++i, d += strides[0]) {
2182 1
                count += (*d != 0);
2183
            }
2184 0
        } NPY_RAW_ITER_ONE_NEXT(idim, ndim, coord, shape, data, strides);
2185
    }
2186

2187 1
    NPY_END_THREADS;
2188

2189
    return count;
2190
}
2191

2192
/*NUMPY_API
2193
 * Counts the number of non-zero elements in the array.
2194
 *
2195
 * Returns -1 on error.
2196
 */
2197
NPY_NO_EXPORT npy_intp
2198 1
PyArray_CountNonzero(PyArrayObject *self)
2199
{
2200
    PyArray_NonzeroFunc *nonzero;
2201
    char *data;
2202
    npy_intp stride, count;
2203 1
    npy_intp nonzero_count = 0;
2204 1
    int needs_api = 0;
2205
    PyArray_Descr *dtype;
2206

2207
    NpyIter *iter;
2208
    NpyIter_IterNextFunc *iternext;
2209
    char **dataptr;
2210
    npy_intp *strideptr, *innersizeptr;
2211 1
    NPY_BEGIN_THREADS_DEF;
2212

2213
    /* Special low-overhead version specific to the boolean type */
2214 1
    dtype = PyArray_DESCR(self);
2215 1
    if (dtype->type_num == NPY_BOOL) {
2216 1
        return count_boolean_trues(PyArray_NDIM(self), PyArray_DATA(self),
2217 1
                        PyArray_DIMS(self), PyArray_STRIDES(self));
2218
    }
2219 1
    nonzero = PyArray_DESCR(self)->f->nonzero;
2220

2221
    /* If it's a trivial one-dimensional loop, don't use an iterator */
2222 1
    if (PyArray_TRIVIALLY_ITERABLE(self)) {
2223 1
        needs_api = PyDataType_FLAGCHK(dtype, NPY_NEEDS_PYAPI);
2224 1
        PyArray_PREPARE_TRIVIAL_ITERATION(self, count, data, stride);
2225

2226 1
        if (needs_api){
2227 1
            while (count--) {
2228 1
                if (nonzero(data, self)) {
2229 1
                    ++nonzero_count;
2230
                }
2231 1
                if (PyErr_Occurred()) {
2232
                    return -1;
2233
                }
2234 1
                data += stride;
2235
            }
2236
        }
2237
        else {
2238 1
            NPY_BEGIN_THREADS_THRESHOLDED(count);
2239 1
            while (count--) {
2240 1
                if (nonzero(data, self)) {
2241 1
                    ++nonzero_count;
2242
                }
2243 1
                data += stride;
2244
            }
2245 1
            NPY_END_THREADS;
2246
        }
2247

2248
        return nonzero_count;
2249
    }
2250

2251
    /*
2252
     * If the array has size zero, return zero (the iterator rejects
2253
     * size zero arrays)
2254
     */
2255 1
    if (PyArray_SIZE(self) == 0) {
2256
        return 0;
2257
    }
2258

2259
    /*
2260
     * Otherwise create and use an iterator to count the nonzeros.
2261
     */
2262 1
    iter = NpyIter_New(self, NPY_ITER_READONLY |
2263
                             NPY_ITER_EXTERNAL_LOOP |
2264
                             NPY_ITER_REFS_OK,
2265
                        NPY_KEEPORDER, NPY_NO_CASTING,
2266
                        NULL);
2267 1
    if (iter == NULL) {
2268
        return -1;
2269
    }
2270 1
    needs_api = NpyIter_IterationNeedsAPI(iter);
2271

2272
    /* Get the pointers for inner loop iteration */
2273 1
    iternext = NpyIter_GetIterNext(iter, NULL);
2274 1
    if (iternext == NULL) {
2275 0
        NpyIter_Deallocate(iter);
2276 0
        return -1;
2277
    }
2278

2279 1
    NPY_BEGIN_THREADS_NDITER(iter);
2280

2281 1
    dataptr = NpyIter_GetDataPtrArray(iter);
2282 1
    strideptr = NpyIter_GetInnerStrideArray(iter);
2283 1
    innersizeptr = NpyIter_GetInnerLoopSizePtr(iter);
2284

2285
    /* Iterate over all the elements to count the nonzeros */
2286
    do {
2287 1
        data = *dataptr;
2288 1
        stride = *strideptr;
2289 1
        count = *innersizeptr;
2290

2291 1
        while (count--) {
2292 1
            if (nonzero(data, self)) {
2293 1
                ++nonzero_count;
2294
            }
2295 1
            if (needs_api && PyErr_Occurred()) {
2296
                nonzero_count = -1;
2297
                goto finish;
2298
            }
2299 1
            data += stride;
2300
        }
2301

2302 1
    } while(iternext(iter));
2303

2304 1
finish:
2305 1
    NPY_END_THREADS;
2306

2307 1
    NpyIter_Deallocate(iter);
2308

2309 1
    return nonzero_count;
2310
}
2311

2312
/*NUMPY_API
2313
 * Nonzero
2314
 *
2315
 * TODO: In NumPy 2.0, should make the iteration order a parameter.
2316
 */
2317
NPY_NO_EXPORT PyObject *
2318 1
PyArray_Nonzero(PyArrayObject *self)
2319
{
2320 1
    int i, ndim = PyArray_NDIM(self);
2321 1
    PyArrayObject *ret = NULL;
2322
    PyObject *ret_tuple;
2323
    npy_intp ret_dims[2];
2324

2325
    PyArray_NonzeroFunc *nonzero;
2326
    PyArray_Descr *dtype;
2327

2328
    npy_intp nonzero_count;
2329 1
    npy_intp added_count = 0;
2330
    int needs_api;
2331
    int is_bool;
2332

2333
    NpyIter *iter;
2334
    NpyIter_IterNextFunc *iternext;
2335
    NpyIter_GetMultiIndexFunc *get_multi_index;
2336
    char **dataptr;
2337

2338 1
    dtype = PyArray_DESCR(self);
2339 1
    nonzero = dtype->f->nonzero;
2340 1
    needs_api = PyDataType_FLAGCHK(dtype, NPY_NEEDS_PYAPI);
2341

2342
    /* Special case - nonzero(zero_d) is nonzero(atleast_1d(zero_d)) */
2343 1
    if (ndim == 0) {
2344
        char const* msg;
2345 1
        if (PyArray_ISBOOL(self)) {
2346
            msg =
2347
                "Calling nonzero on 0d arrays is deprecated, as it behaves "
2348
                "surprisingly. Use `atleast_1d(cond).nonzero()` if the old "
2349
                "behavior was intended. If the context of this warning is of "
2350
                "the form `arr[nonzero(cond)]`, just use `arr[cond]`.";
2351
        }
2352
        else {
2353 1
            msg =
2354
                "Calling nonzero on 0d arrays is deprecated, as it behaves "
2355
                "surprisingly. Use `atleast_1d(arr).nonzero()` if the old "
2356
                "behavior was intended.";
2357
        }
2358 1
        if (DEPRECATE(msg) < 0) {
2359
            return NULL;
2360
        }
2361

2362
        static npy_intp const zero_dim_shape[1] = {1};
2363
        static npy_intp const zero_dim_strides[1] = {0};
2364

2365 1
        Py_INCREF(PyArray_DESCR(self));  /* array creation steals reference */
2366 1
        PyArrayObject *self_1d = (PyArrayObject *)PyArray_NewFromDescrAndBase(
2367 1
            Py_TYPE(self), PyArray_DESCR(self),
2368 1
            1, zero_dim_shape, zero_dim_strides, PyArray_BYTES(self),
2369
            PyArray_FLAGS(self), (PyObject *)self, (PyObject *)self);
2370 1
        if (self_1d == NULL) {
2371
            return NULL;
2372
        }
2373 1
        ret_tuple = PyArray_Nonzero(self_1d);
2374 1
        Py_DECREF(self_1d);
2375
        return ret_tuple;
2376
    }
2377

2378
    /*
2379
     * First count the number of non-zeros in 'self'.
2380
     */
2381 1
    nonzero_count = PyArray_CountNonzero(self);
2382 1
    if (nonzero_count < 0) {
2383
        return NULL;
2384
    }
2385

2386 1
    is_bool = PyArray_ISBOOL(self);
2387

2388
    /* Allocate the result as a 2D array */
2389 1
    ret_dims[0] = nonzero_count;
2390 1
    ret_dims[1] = ndim;
2391 1
    ret = (PyArrayObject *)PyArray_NewFromDescr(
2392
            &PyArray_Type, PyArray_DescrFromType(NPY_INTP),
2393
            2, ret_dims, NULL, NULL,
2394
            0, NULL);
2395 1
    if (ret == NULL) {
2396
        return NULL;
2397
    }
2398

2399
    /* If it's a one-dimensional result, don't use an iterator */
2400 1
    if (ndim == 1) {
2401 1
        npy_intp * multi_index = (npy_intp *)PyArray_DATA(ret);
2402 1
        char * data = PyArray_BYTES(self);
2403 1
        npy_intp stride = PyArray_STRIDE(self, 0);
2404 1
        npy_intp count = PyArray_DIM(self, 0);
2405 1
        NPY_BEGIN_THREADS_DEF;
2406

2407
        /* nothing to do */
2408 1
        if (nonzero_count == 0) {
2409
            goto finish;
2410
        }
2411

2412 1
        if (!needs_api) {
2413 1
            NPY_BEGIN_THREADS_THRESHOLDED(count);
2414
        }
2415

2416
        /* avoid function call for bool */
2417 1
        if (is_bool) {
2418
            /*
2419
             * use fast memchr variant for sparse data, see gh-4370
2420
             * the fast bool count is followed by this sparse path is faster
2421
             * than combining the two loops, even for larger arrays
2422
             */
2423 1
            if (((double)nonzero_count / count) <= 0.1) {
2424
                npy_intp subsize;
2425
                npy_intp j = 0;
2426
                while (1) {
2427 1
                    npy_memchr(data + j * stride, 0, stride, count - j,
2428
                               &subsize, 1);
2429 1
                    j += subsize;
2430 1
                    if (j >= count) {
2431
                        break;
2432
                    }
2433 1
                    *multi_index++ = j++;
2434
                }
2435
            }
2436
            else {
2437
                npy_intp j;
2438 1
                for (j = 0; j < count; ++j) {
2439 1
                    if (*data != 0) {
2440 1
                        *multi_index++ = j;
2441
                    }
2442 1
                    data += stride;
2443
                }
2444
            }
2445
        }
2446
        else {
2447
            npy_intp j;
2448 1
            for (j = 0; j < count; ++j) {
2449 1
                if (nonzero(data, self)) {
2450 1
                    if (++added_count > nonzero_count) {
2451
                        break;
2452
                    }
2453 1
                    *multi_index++ = j;
2454
                }
2455 1
                if (needs_api && PyErr_Occurred()) {
2456
                    break;
2457
                }
2458 1
                data += stride;
2459
            }
2460
        }
2461

2462 1
        NPY_END_THREADS;
2463

2464
        goto finish;
2465
    }
2466

2467
    /*
2468
     * Build an iterator tracking a multi-index, in C order.
2469
     */
2470 1
    iter = NpyIter_New(self, NPY_ITER_READONLY |
2471
                             NPY_ITER_MULTI_INDEX |
2472
                             NPY_ITER_ZEROSIZE_OK |
2473
                             NPY_ITER_REFS_OK,
2474
                        NPY_CORDER, NPY_NO_CASTING,
2475
                        NULL);
2476

2477 1
    if (iter == NULL) {
2478 0
        Py_DECREF(ret);
2479
        return NULL;
2480
    }
2481

2482 1
    if (NpyIter_GetIterSize(iter) != 0) {
2483
        npy_intp * multi_index;
2484 1
        NPY_BEGIN_THREADS_DEF;
2485
        /* Get the pointers for inner loop iteration */
2486 1
        iternext = NpyIter_GetIterNext(iter, NULL);
2487 1
        if (iternext == NULL) {
2488 0
            NpyIter_Deallocate(iter);
2489 0
            Py_DECREF(ret);
2490
            return NULL;
2491
        }
2492 1
        get_multi_index = NpyIter_GetGetMultiIndex(iter, NULL);
2493 1
        if (get_multi_index == NULL) {
2494 0
            NpyIter_Deallocate(iter);
2495 0
            Py_DECREF(ret);
2496
            return NULL;
2497
        }
2498

2499 1
        needs_api = NpyIter_IterationNeedsAPI(iter);
2500

2501 1
        NPY_BEGIN_THREADS_NDITER(iter);
2502

2503 1
        dataptr = NpyIter_GetDataPtrArray(iter);
2504

2505 1
        multi_index = (npy_intp *)PyArray_DATA(ret);
2506

2507
        /* Get the multi-index for each non-zero element */
2508 1
        if (is_bool) {
2509
            /* avoid function call for bool */
2510
            do {
2511 1
                if (**dataptr != 0) {
2512 1
                    get_multi_index(iter, multi_index);
2513 1
                    multi_index += ndim;
2514
                }
2515 1
            } while(iternext(iter));
2516
        }
2517
        else {
2518
            do {
2519 1
                if (nonzero(*dataptr, self)) {
2520 1
                    if (++added_count > nonzero_count) {
2521
                        break;
2522
                    }
2523 1
                    get_multi_index(iter, multi_index);
2524 1
                    multi_index += ndim;
2525
                }
2526 1
                if (needs_api && PyErr_Occurred()) {
2527
                    break;
2528
                }
2529 1
            } while(iternext(iter));
2530
        }
2531

2532 1
        NPY_END_THREADS;
2533
    }
2534

2535 1
    NpyIter_Deallocate(iter);
2536

2537 1
finish:
2538 1
    if (PyErr_Occurred()) {
2539 1
        Py_DECREF(ret);
2540
        return NULL;
2541
    }
2542

2543
    /* if executed `nonzero()` check for miscount due to side-effect */
2544 1
    if (!is_bool && added_count != nonzero_count) {
2545 1
        PyErr_SetString(PyExc_RuntimeError,
2546
            "number of non-zero array elements "
2547
            "changed during function execution.");
2548 1
        Py_DECREF(ret);
2549
        return NULL;
2550
    }
2551

2552 1
    ret_tuple = PyTuple_New(ndim);
2553 1
    if (ret_tuple == NULL) {
2554 0
        Py_DECREF(ret);
2555
        return NULL;
2556
    }
2557

2558
    /* Create views into ret, one for each dimension */
2559 1
    for (i = 0; i < ndim; ++i) {
2560 1
        npy_intp stride = ndim * NPY_SIZEOF_INTP;
2561
        /* the result is an empty array, the view must point to valid memory */
2562 1
        npy_intp data_offset = nonzero_count == 0 ? 0 : i * NPY_SIZEOF_INTP;
2563

2564 1
        PyArrayObject *view = (PyArrayObject *)PyArray_NewFromDescrAndBase(
2565 1
            Py_TYPE(ret), PyArray_DescrFromType(NPY_INTP),
2566 1
            1, &nonzero_count, &stride, PyArray_BYTES(ret) + data_offset,
2567
            PyArray_FLAGS(ret), (PyObject *)ret, (PyObject *)ret);
2568 1
        if (view == NULL) {
2569 0
            Py_DECREF(ret);
2570 0
            Py_DECREF(ret_tuple);
2571 0
            return NULL;
2572
        }
2573 1
        PyTuple_SET_ITEM(ret_tuple, i, (PyObject *)view);
2574
    }
2575 1
    Py_DECREF(ret);
2576

2577
    return ret_tuple;
2578
}
2579

2580
/*
2581
 * Gets a single item from the array, based on a single multi-index
2582
 * array of values, which must be of length PyArray_NDIM(self).
2583
 */
2584
NPY_NO_EXPORT PyObject *
2585 1
PyArray_MultiIndexGetItem(PyArrayObject *self, const npy_intp *multi_index)
2586
{
2587 1
    int idim, ndim = PyArray_NDIM(self);
2588 1
    char *data = PyArray_DATA(self);
2589 1
    npy_intp *shape = PyArray_SHAPE(self);
2590 1
    npy_intp *strides = PyArray_STRIDES(self);
2591

2592
    /* Get the data pointer */
2593 1
    for (idim = 0; idim < ndim; ++idim) {
2594 1
        npy_intp shapevalue = shape[idim];
2595 1
        npy_intp ind = multi_index[idim];
2596

2597 1
        if (check_and_adjust_index(&ind, shapevalue, idim, NULL) < 0) {
2598 1
            return NULL;
2599
        }
2600 1
        data += ind * strides[idim];
2601
    }
2602

2603 1
    return PyArray_GETITEM(self, data);
2604
}
2605

2606
/*
2607
 * Sets a single item in the array, based on a single multi-index
2608
 * array of values, which must be of length PyArray_NDIM(self).
2609
 *
2610
 * Returns 0 on success, -1 on failure.
2611
 */
2612
NPY_NO_EXPORT int
2613 1
PyArray_MultiIndexSetItem(PyArrayObject *self, const npy_intp *multi_index,
2614
                                                PyObject *obj)
2615
{
2616 1
    int idim, ndim = PyArray_NDIM(self);
2617 1
    char *data = PyArray_DATA(self);
2618 1
    npy_intp *shape = PyArray_SHAPE(self);
2619 1
    npy_intp *strides = PyArray_STRIDES(self);
2620

2621
    /* Get the data pointer */
2622 1
    for (idim = 0; idim < ndim; ++idim) {
2623 1
        npy_intp shapevalue = shape[idim];
2624 1
        npy_intp ind = multi_index[idim];
2625

2626 1
        if (check_and_adjust_index(&ind, shapevalue, idim, NULL) < 0) {
2627 1
            return -1;
2628
        }
2629 1
        data += ind * strides[idim];
2630
    }
2631

2632 1
    return PyArray_Pack(PyArray_DESCR(self), data, obj);
2633
}

Read our documentation on viewing source code .

Loading