/*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;
}
