/*NUMPY_API * ArgMax */ NPY_NO_EXPORT PyObject * PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out) { PyArrayObject *ap = NULL, *rp = NULL; PyArray_ArgFunc* arg_func; char *ip; npy_intp *rptr; npy_intp i, n, m; int elsize; NPY_BEGIN_THREADS_DEF; if ((ap = (PyArrayObject *)PyArray_CheckAxis(op, &axis, 0)) == NULL) { return NULL; } /* * We need to permute the array so that axis is placed at the end. * And all other dimensions are shifted left. */ if (axis != PyArray_NDIM(ap)-1) { PyArray_Dims newaxes; npy_intp dims[NPY_MAXDIMS]; int j; newaxes.ptr = dims; newaxes.len = PyArray_NDIM(ap); for (j = 0; j < axis; j++) { dims[j] = j; } for (j = axis; j < PyArray_NDIM(ap) - 1; j++) { dims[j] = j + 1; } dims[PyArray_NDIM(ap) - 1] = axis; op = (PyArrayObject *)PyArray_Transpose(ap, &newaxes); Py_DECREF(ap); if (op == NULL) { return NULL; } } else { op = ap; } /* Will get native-byte order contiguous copy. */ ap = (PyArrayObject *)PyArray_ContiguousFromAny((PyObject *)op, PyArray_DESCR(op)->type_num, 1, 0); Py_DECREF(op); if (ap == NULL) { return NULL; } arg_func = PyArray_DESCR(ap)->f->argmax; if (arg_func == NULL) { PyErr_SetString(PyExc_TypeError, "data type not ordered"); goto fail; } elsize = PyArray_DESCR(ap)->elsize; m = PyArray_DIMS(ap)[PyArray_NDIM(ap)-1]; if (m == 0) { PyErr_SetString(PyExc_ValueError, "attempt to get argmax of an empty sequence"); goto fail; } if (!out) { rp = (PyArrayObject *)PyArray_NewFromDescr( Py_TYPE(ap), PyArray_DescrFromType(NPY_INTP), PyArray_NDIM(ap) - 1, PyArray_DIMS(ap), NULL, NULL, 0, (PyObject *)ap); if (rp == NULL) { goto fail; } } else { if ((PyArray_NDIM(out) != PyArray_NDIM(ap) - 1) || !PyArray_CompareLists(PyArray_DIMS(out), PyArray_DIMS(ap), PyArray_NDIM(out))) { PyErr_SetString(PyExc_ValueError, "output array does not match result of np.argmax."); goto fail; } rp = (PyArrayObject *)PyArray_FromArray(out, PyArray_DescrFromType(NPY_INTP), NPY_ARRAY_CARRAY | NPY_ARRAY_WRITEBACKIFCOPY); if (rp == NULL) { goto fail; } } NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(ap)); n = PyArray_SIZE(ap)/m; rptr = (npy_intp *)PyArray_DATA(rp); for (ip = PyArray_DATA(ap), i = 0; i < n; i++, ip += elsize*m) { arg_func(ip, m, rptr, ap); rptr += 1; } NPY_END_THREADS_DESCR(PyArray_DESCR(ap)); Py_DECREF(ap); /* Trigger the UPDATEIFCOPY/WRTIEBACKIFCOPY if necessary */ if (out != NULL && out != rp) { PyArray_ResolveWritebackIfCopy(rp); Py_DECREF(rp); rp = out; Py_INCREF(rp); } return (PyObject *)rp; fail: Py_DECREF(ap); Py_XDECREF(rp); return NULL; }
/*NUMPY_API * ArgMax */ NPY_NO_EXPORT PyObject * PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out) { PyArrayObject *ap = NULL, *rp = NULL; PyArray_ArgFunc* arg_func; char *ip; intp *rptr; intp i, n, m; int elsize; int copyret = 0; NPY_BEGIN_THREADS_DEF; if ((ap=(PyAO *)_check_axis(op, &axis, 0)) == NULL) { return NULL; } /* * We need to permute the array so that axis is placed at the end. * And all other dimensions are shifted left. */ if (axis != ap->nd-1) { PyArray_Dims newaxes; intp dims[MAX_DIMS]; int i; newaxes.ptr = dims; newaxes.len = ap->nd; for (i = 0; i < axis; i++) dims[i] = i; for (i = axis; i < ap->nd - 1; i++) dims[i] = i + 1; dims[ap->nd - 1] = axis; op = (PyAO *)PyArray_Transpose(ap, &newaxes); Py_DECREF(ap); if (op == NULL) { return NULL; } } else { op = ap; } /* Will get native-byte order contiguous copy. */ ap = (PyArrayObject *) PyArray_ContiguousFromAny((PyObject *)op, op->descr->type_num, 1, 0); Py_DECREF(op); if (ap == NULL) { return NULL; } arg_func = ap->descr->f->argmax; if (arg_func == NULL) { PyErr_SetString(PyExc_TypeError, "data type not ordered"); goto fail; } elsize = ap->descr->elsize; m = ap->dimensions[ap->nd-1]; if (m == 0) { PyErr_SetString(PyExc_ValueError, "attempt to get argmax/argmin "\ "of an empty sequence"); goto fail; } if (!out) { rp = (PyArrayObject *)PyArray_New(ap->ob_type, ap->nd-1, ap->dimensions, PyArray_INTP, NULL, NULL, 0, 0, (PyObject *)ap); if (rp == NULL) { goto fail; } } else { if (PyArray_SIZE(out) != PyArray_MultiplyList(ap->dimensions, ap->nd - 1)) { PyErr_SetString(PyExc_TypeError, "invalid shape for output array."); } rp = (PyArrayObject *)\ PyArray_FromArray(out, PyArray_DescrFromType(PyArray_INTP), NPY_CARRAY | NPY_UPDATEIFCOPY); if (rp == NULL) { goto fail; } if (rp != out) { copyret = 1; } } NPY_BEGIN_THREADS_DESCR(ap->descr); n = PyArray_SIZE(ap)/m; rptr = (intp *)rp->data; for (ip = ap->data, i = 0; i < n; i++, ip += elsize*m) { arg_func(ip, m, rptr, ap); rptr += 1; } NPY_END_THREADS_DESCR(ap->descr); Py_DECREF(ap); if (copyret) { PyArrayObject *obj; obj = (PyArrayObject *)rp->base; Py_INCREF(obj); Py_DECREF(rp); rp = obj; } return (PyObject *)rp; fail: Py_DECREF(ap); Py_XDECREF(rp); return NULL; }
/*NUMPY_API * ArgMin */ NPY_NO_EXPORT PyObject * PyArray_ArgMin(PyArrayObject *op, int axis, PyArrayObject *out) { PyArrayObject *ap = NULL, *rp = NULL; PyArray_ArgFunc* arg_func; char *ip; intp *rptr; intp i, n, m; int elsize; NPY_BEGIN_THREADS_DEF; if ((ap = (PyArrayObject *)PyArray_CheckAxis(op, &axis, 0)) == NULL) { return NULL; } /* * We need to permute the array so that axis is placed at the end. * And all other dimensions are shifted left. */ if (axis != PyArray_NDIM(ap)-1) { PyArray_Dims newaxes; intp dims[MAX_DIMS]; int i; newaxes.ptr = dims; newaxes.len = PyArray_NDIM(ap); for (i = 0; i < axis; i++) { dims[i] = i; } for (i = axis; i < PyArray_NDIM(ap) - 1; i++) { dims[i] = i + 1; } dims[PyArray_NDIM(ap) - 1] = axis; op = (PyArrayObject *)PyArray_Transpose(ap, &newaxes); Py_DECREF(ap); if (op == NULL) { return NULL; } } else { op = ap; } /* Will get native-byte order contiguous copy. */ ap = (PyArrayObject *)PyArray_ContiguousFromAny((PyObject *)op, PyArray_DESCR(op)->type_num, 1, 0); Py_DECREF(op); if (ap == NULL) { return NULL; } arg_func = PyArray_DESCR(ap)->f->argmin; if (arg_func == NULL) { PyErr_SetString(PyExc_TypeError, "data type not ordered"); goto fail; } elsize = PyArray_DESCR(ap)->elsize; m = PyArray_DIMS(ap)[PyArray_NDIM(ap)-1]; if (m == 0) { PyErr_SetString(PyExc_ValueError, "attempt to get argmin of an empty sequence"); goto fail; } if (!out) { rp = (PyArrayObject *)PyArray_New(Py_TYPE(ap), PyArray_NDIM(ap)-1, PyArray_DIMS(ap), PyArray_INTP, NULL, NULL, 0, 0, (PyObject *)ap); if (rp == NULL) { goto fail; } } else { if (PyArray_SIZE(out) != PyArray_MultiplyList(PyArray_DIMS(ap), PyArray_NDIM(ap) - 1)) { PyErr_SetString(PyExc_TypeError, "invalid shape for output array."); } rp = (PyArrayObject *)PyArray_FromArray(out, PyArray_DescrFromType(PyArray_INTP), NPY_ARRAY_CARRAY | NPY_ARRAY_UPDATEIFCOPY); if (rp == NULL) { goto fail; } } NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(ap)); n = PyArray_SIZE(ap)/m; rptr = (intp *)PyArray_DATA(rp); for (ip = PyArray_DATA(ap), i = 0; i < n; i++, ip += elsize*m) { arg_func(ip, m, rptr, ap); rptr += 1; } NPY_END_THREADS_DESCR(PyArray_DESCR(ap)); Py_DECREF(ap); /* Trigger the UPDATEIFCOPY if necessary */ if (out != NULL && out != rp) { Py_DECREF(rp); rp = out; Py_INCREF(rp); } return (PyObject *)rp; fail: Py_DECREF(ap); Py_XDECREF(rp); return NULL; }