path: root/numpy/core/src/umath/reduction.c

/*
 * This file implements generic methods for computing reductions on arrays.
 *
 * Written by Mark Wiebe (mwwiebe@gmail.com)
 * Copyright (c) 2011 by Enthought, Inc.
 *
 * See LICENSE.txt for the license.
 */
#define _UMATHMODULE
#define _MULTIARRAYMODULE
#define NPY_NO_DEPRECATED_API NPY_API_VERSION

#define PY_SSIZE_T_CLEAN
#include <Python.h>

#include "npy_config.h"
#include <numpy/arrayobject.h>

#include "npy_config.h"
#include "npy_pycompat.h"

#include "numpy/ufuncobject.h"
#include "lowlevel_strided_loops.h"
#include "reduction.h"
#include "extobj.h"  /* for _check_ufunc_fperr */

/*
 * Allocates a result array for a reduction operation, with
 * dimensions matching 'arr' except set to 1 with 0 stride
 * wherever axis_flags is True. Dropping the reduction axes
 * from the result must be done later by the caller once the
 * computation is complete.
 *
 * This function always allocates a base class ndarray.
 *
 * If 'dtype' isn't NULL, this function steals its reference.
 */
static PyArrayObject *
allocate_reduce_result(PyArrayObject *arr, npy_bool *axis_flags,
                        PyArray_Descr *dtype, int subok)
{
    npy_intp strides[NPY_MAXDIMS], stride;
    npy_intp shape[NPY_MAXDIMS], *arr_shape = PyArray_DIMS(arr);
    npy_stride_sort_item strideperm[NPY_MAXDIMS];
    int idim, ndim = PyArray_NDIM(arr);

    if (dtype == NULL) {
        dtype = PyArray_DTYPE(arr);
        Py_INCREF(dtype);
    }

    PyArray_CreateSortedStridePerm(PyArray_NDIM(arr),
                                    PyArray_STRIDES(arr), strideperm);

    /* Build the new strides and shape */
    stride = dtype->elsize;
    memcpy(shape, arr_shape, ndim * sizeof(shape[0]));
    for (idim = ndim-1; idim >= 0; --idim) {
        npy_intp i_perm = strideperm[idim].perm;
        if (axis_flags[i_perm]) {
            strides[i_perm] = 0;
            shape[i_perm] = 1;
        }
        else {
            strides[i_perm] = stride;
            stride *= shape[i_perm];
        }
    }

    /* Finally, allocate the array */
    return (PyArrayObject *)PyArray_NewFromDescr(
                                    subok ? Py_TYPE(arr) : &PyArray_Type,
                                    dtype, ndim, shape, strides,
                                    NULL, 0, subok ? (PyObject *)arr : NULL);
}

/*
 * Conforms an output parameter 'out' to have 'ndim' dimensions
 * with dimensions of size one added in the appropriate places
 * indicated by 'axis_flags'.
 *
 * The return value is a view into 'out'.
 */
static PyArrayObject *
conform_reduce_result(int ndim, npy_bool *axis_flags,
                      PyArrayObject *out, int keepdims, const char *funcname,
                      int need_copy)
{
    npy_intp strides[NPY_MAXDIMS], shape[NPY_MAXDIMS];
    npy_intp *strides_out = PyArray_STRIDES(out);
    npy_intp *shape_out = PyArray_DIMS(out);
    int idim, idim_out, ndim_out = PyArray_NDIM(out);
    PyArray_Descr *dtype;
    PyArrayObject_fields *ret;

    /*
     * If the 'keepdims' parameter is true, do a simpler validation and
     * return a new reference to 'out'.
     */
    if (keepdims) {
        if (PyArray_NDIM(out) != ndim) {
            PyErr_Format(PyExc_ValueError,
                    "output parameter for reduction operation %s "
                    "has the wrong number of dimensions (must match "
                    "the operand's when keepdims=True)", funcname);
            return NULL;
        }

        for (idim = 0; idim < ndim; ++idim) {
            if (axis_flags[idim]) {
                if (shape_out[idim] != 1) {
                    PyErr_Format(PyExc_ValueError,
                            "output parameter for reduction operation %s "
                            "has a reduction dimension not equal to one "
                            "(required when keepdims=True)", funcname);
                    return NULL;
                }
            }
        }

        Py_INCREF(out);
        return out;
    }

    /* Construct the strides and shape */
    idim_out = 0;
    for (idim = 0; idim < ndim; ++idim) {
        if (axis_flags[idim]) {
            strides[idim] = 0;
            shape[idim] = 1;
        }
        else {
            if (idim_out >= ndim_out) {
                PyErr_Format(PyExc_ValueError,
                        "output parameter for reduction operation %s "
                        "does not have enough dimensions", funcname);
                return NULL;
            }
            strides[idim] = strides_out[idim_out];
            shape[idim] = shape_out[idim_out];
            ++idim_out;
        }
    }

    if (idim_out != ndim_out) {
        PyErr_Format(PyExc_ValueError,
                "output parameter for reduction operation %s "
                "has too many dimensions", funcname);
        return NULL;
    }

    /* Allocate the view */
    dtype = PyArray_DESCR(out);
    Py_INCREF(dtype);

    /* TODO: use PyArray_NewFromDescrAndBase here once multiarray and umath
     *       are merged
     */
    ret = (PyArrayObject_fields *)PyArray_NewFromDescr(
            &PyArray_Type, dtype,
            ndim, shape, strides, PyArray_DATA(out),
            PyArray_FLAGS(out), NULL);
    if (ret == NULL) {
        return NULL;
    }

    Py_INCREF(out);
    if (PyArray_SetBaseObject((PyArrayObject *)ret, (PyObject *)out) < 0) {
        Py_DECREF(ret);
        return NULL;
    }

    if (need_copy) {
        PyArrayObject *ret_copy;

        ret_copy = (PyArrayObject *)PyArray_NewLikeArray(
            (PyArrayObject *)ret, NPY_ANYORDER, NULL, 0);
        if (ret_copy == NULL) {
            Py_DECREF(ret);
            return NULL;
        }

        if (PyArray_CopyInto(ret_copy, (PyArrayObject *)ret) != 0) {
            Py_DECREF(ret);
            Py_DECREF(ret_copy);
            return NULL;
        }

        if (PyArray_SetWritebackIfCopyBase(ret_copy, (PyArrayObject *)ret) < 0) {
            Py_DECREF(ret);
            Py_DECREF(ret_copy);
            return NULL;
        }

        return ret_copy;
    }
    else {
        return (PyArrayObject *)ret;
    }
}

/*
 * Creates a result for reducing 'operand' along the axes specified
 * in 'axis_flags'. If 'dtype' isn't NULL, this function steals a
 * reference to 'dtype'.
 *
 * If 'out' isn't NULL, this function creates a view conforming
 * to the number of dimensions of 'operand', adding a singleton dimension
 * for each reduction axis specified. In this case, 'dtype' is ignored
 * (but its reference is still stolen), and the caller must handle any
 * type conversion/validity check for 'out'
 *
 * If 'subok' is true, creates a result with the subtype of 'operand',
 * otherwise creates on with the base ndarray class.
 *
 * If 'out' is NULL, it allocates a new array whose shape matches that of
 * 'operand', except for at the reduction axes. If 'dtype' is NULL, the dtype
 * of 'operand' is used for the result.
 */
NPY_NO_EXPORT PyArrayObject *
PyArray_CreateReduceResult(PyArrayObject *operand, PyArrayObject *out,
                           PyArray_Descr *dtype, npy_bool *axis_flags,
                           int keepdims, int subok,
                           const char *funcname)
{
    PyArrayObject *result;

    if (out == NULL) {
        /* This function steals the reference to 'dtype' */
        result = allocate_reduce_result(operand, axis_flags, dtype, subok);
    }
    else {
        int need_copy = 0;

        if (solve_may_share_memory(operand, out, 1) != 0) {
            need_copy = 1;
        }

        /* Steal the dtype reference */
        Py_XDECREF(dtype);
        result = conform_reduce_result(PyArray_NDIM(operand), axis_flags,
                                       out, keepdims, funcname, need_copy);
    }

    return result;
}

/*
 * Count the number of dimensions selected in 'axis_flags'
 */
static int
count_axes(int ndim, npy_bool *axis_flags)
{
    int idim;
    int naxes = 0;

    for (idim = 0; idim < ndim; ++idim) {
        if (axis_flags[idim]) {
            naxes++;
        }
    }
    return naxes;
}

/*
 * This function initializes a result array for a reduction operation
 * which has no identity. This means it needs to copy the first element
 * it sees along the reduction axes to result, then return a view of
 * the operand which excludes that element.
 *
 * If a reduction has an identity, such as 0 or 1, the result should be
 * initialized by calling PyArray_AssignZero(result, NULL, NULL) or
 * PyArray_AssignOne(result, NULL, NULL), because this function raises an
 * exception when there are no elements to reduce (which appropriate iff the
 * reduction operation has no identity).
 *
 * This means it copies the subarray indexed at zero along each reduction axis
 * into 'result', then returns a view into 'operand' excluding those copied
 * elements.
 *
 * result  : The array into which the result is computed. This must have
 *           the same number of dimensions as 'operand', but for each
 *           axis i where 'axis_flags[i]' is True, it has a single element.
 * operand : The array being reduced.
 * axis_flags : An array of boolean flags, one for each axis of 'operand'.
 *              When a flag is True, it indicates to reduce along that axis.
 * out_skip_first_count : This gets populated with the number of first-visit
 *                        elements that should be skipped during the
 *                        iteration loop.
 * funcname : The name of the reduction operation, for the purpose of
 *            better quality error messages. For example, "numpy.max"
 *            would be a good name for NumPy's max function.
 *
 * Returns a view which contains the remaining elements on which to do
 * the reduction.
 */
NPY_NO_EXPORT PyArrayObject *
PyArray_InitializeReduceResult(
                    PyArrayObject *result, PyArrayObject *operand,
                    npy_bool *axis_flags,
                    npy_intp *out_skip_first_count, const char *funcname)
{
    npy_intp *strides, *shape, shape_orig[NPY_MAXDIMS];
    PyArrayObject *op_view = NULL;
    int idim, ndim, nreduce_axes;

    ndim = PyArray_NDIM(operand);

    /* Default to no skipping first-visit elements in the iteration */
    *out_skip_first_count = 0;

    /* Take a view into 'operand' which we can modify. */
    op_view = (PyArrayObject *)PyArray_View(operand, NULL, &PyArray_Type);
    if (op_view == NULL) {
        return NULL;
    }

    /*
     * Now copy the subarray of the first element along each reduction axis,
     * then return a view to the rest.
     *
     * Adjust the shape to only look at the first element along
     * any of the reduction axes. We count the number of reduction axes
     * at the same time.
     */
    shape = PyArray_SHAPE(op_view);
    nreduce_axes = 0;
    memcpy(shape_orig, shape, ndim * sizeof(npy_intp));
    for (idim = 0; idim < ndim; ++idim) {
        if (axis_flags[idim]) {
            if (shape[idim] == 0) {
                PyErr_Format(PyExc_ValueError,
                             "zero-size array to reduction operation %s "
                             "which has no identity",
                             funcname);
                Py_DECREF(op_view);
                return NULL;
            }
            shape[idim] = 1;
            ++nreduce_axes;
        }
    }

    /*
     * Copy the elements into the result to start.
     */
    if (PyArray_CopyInto(result, op_view) < 0) {
        Py_DECREF(op_view);
        return NULL;
    }

    /*
     * If there is one reduction axis, adjust the view's
     * shape to only look at the remaining elements
     */
    if (nreduce_axes == 1) {
        strides = PyArray_STRIDES(op_view);
        for (idim = 0; idim < ndim; ++idim) {
            if (axis_flags[idim]) {
                shape[idim] = shape_orig[idim] - 1;
                ((PyArrayObject_fields *)op_view)->data += strides[idim];
            }
        }
    }
    /* If there are zero reduction axes, make the view empty */
    else if (nreduce_axes == 0) {
        for (idim = 0; idim < ndim; ++idim) {
            shape[idim] = 0;
        }
    }
    /*
     * Otherwise iterate over the whole operand, but tell the inner loop
     * to skip the elements we already copied by setting the skip_first_count.
     */
    else {
        *out_skip_first_count = PyArray_SIZE(result);

        Py_DECREF(op_view);
        Py_INCREF(operand);
        op_view = operand;
    }

    return op_view;
}

/*
 * This function executes all the standard NumPy reduction function
 * boilerplate code, just calling the appropriate inner loop function where
 * necessary.
 *
 * operand     : The array to be reduced.
 * out         : NULL, or the array into which to place the result.
 * wheremask   : NOT YET SUPPORTED, but this parameter is placed here
 *               so that support can be added in the future without breaking
 *               API compatibility. Pass in NULL.
 * operand_dtype : The dtype the inner loop expects for the operand.
 * result_dtype : The dtype the inner loop expects for the result.
 * casting     : The casting rule to apply to the operands.
 * axis_flags  : Flags indicating the reduction axes of 'operand'.
 * reorderable : If True, the reduction being done is reorderable, which
 *               means specifying multiple axes of reduction at once is ok,
 *               and the reduction code may calculate the reduction in an
 *               arbitrary order. The calculation may be reordered because
 *               of cache behavior or multithreading requirements.
 * keepdims    : If true, leaves the reduction dimensions in the result
 *               with size one.
 * subok       : If true, the result uses the subclass of operand, otherwise
 *               it is always a base class ndarray.
 * identity    : If Py_None, PyArray_InitializeReduceResult is used, otherwise
 *               this value is used to initialize the result to
 *               the reduction's unit.
 * loop        : The loop which does the reduction.
 * data        : Data which is passed to the inner loop.
 * buffersize  : Buffer size for the iterator. For the default, pass in 0.
 * funcname    : The name of the reduction function, for error messages.
 * errormask   : forwarded from _get_bufsize_errmask
 *
 * TODO FIXME: if you squint, this is essentially an second independent
 * implementation of generalized ufuncs with signature (i)->(), plus a few
 * extra bells and whistles. (Indeed, as far as I can tell, it was originally
 * split out to support a fancy version of count_nonzero... which is not
 * actually a reduction function at all, it's just a (i)->() function!) So
 * probably these two implementation should be merged into one. (In fact it
 * would be quite nice to support axis= and keepdims etc. for arbitrary
 * generalized ufuncs!)
 */
NPY_NO_EXPORT PyArrayObject *
PyUFunc_ReduceWrapper(PyArrayObject *operand, PyArrayObject *out,
                      PyArrayObject *wheremask,
                      PyArray_Descr *operand_dtype,
                      PyArray_Descr *result_dtype,
                      NPY_CASTING casting,
                      npy_bool *axis_flags, int reorderable,
                      int keepdims,
                      int subok,
                      PyObject *identity,
                      PyArray_ReduceLoopFunc *loop,
                      void *data, npy_intp buffersize, const char *funcname,
                      int errormask)
{
    PyArrayObject *result = NULL, *op_view = NULL;
    npy_intp skip_first_count = 0;

    /* Iterator parameters */
    NpyIter *iter = NULL;
    PyArrayObject *op[2];
    PyArray_Descr *op_dtypes[2];
    npy_uint32 flags, op_flags[2];

    /* More than one axis means multiple orders are possible */
    if (!reorderable && count_axes(PyArray_NDIM(operand), axis_flags) > 1) {
        PyErr_Format(PyExc_ValueError,
                     "reduction operation '%s' is not reorderable, "
                     "so at most one axis may be specified",
                     funcname);
        return NULL;
    }


    /* Validate that the parameters for future expansion are NULL */
    if (wheremask != NULL) {
        PyErr_SetString(PyExc_RuntimeError,
                "Reduce operations in NumPy do not yet support "
                "a where mask");
        return NULL;
    }

    /*
     * This either conforms 'out' to the ndim of 'operand', or allocates
     * a new array appropriate for this reduction.
     *
     * A new array with WRITEBACKIFCOPY is allocated if operand and out have memory
     * overlap.
     */
    Py_INCREF(result_dtype);
    result = PyArray_CreateReduceResult(operand, out,
                            result_dtype, axis_flags,
                            keepdims, subok, funcname);
    if (result == NULL) {
        goto fail;
    }

    /*
     * Initialize the result to the reduction unit if possible,
     * otherwise copy the initial values and get a view to the rest.
     */
    if (identity != Py_None) {
        if (PyArray_FillWithScalar(result, identity) < 0) {
            goto fail;
        }
        op_view = operand;
        Py_INCREF(op_view);
    }
    else {
        op_view = PyArray_InitializeReduceResult(
            result, operand, axis_flags, &skip_first_count, funcname);
        if (op_view == NULL) {
            goto fail;
        }
        /* empty op_view signals no reduction; but 0-d arrays cannot be empty */
        if ((PyArray_SIZE(op_view) == 0) || (PyArray_NDIM(operand) == 0)) {
            Py_DECREF(op_view);
            op_view = NULL;
            goto finish;
        }
    }

    /* Set up the iterator */
    op[0] = result;
    op[1] = op_view;
    op_dtypes[0] = result_dtype;
    op_dtypes[1] = operand_dtype;

    flags = NPY_ITER_BUFFERED |
            NPY_ITER_EXTERNAL_LOOP |
            NPY_ITER_GROWINNER |
            NPY_ITER_DONT_NEGATE_STRIDES |
            NPY_ITER_ZEROSIZE_OK |
            NPY_ITER_REDUCE_OK |
            NPY_ITER_REFS_OK;
    op_flags[0] = NPY_ITER_READWRITE |
                  NPY_ITER_ALIGNED |
                  NPY_ITER_NO_SUBTYPE;
    op_flags[1] = NPY_ITER_READONLY |
                  NPY_ITER_ALIGNED;

    iter = NpyIter_AdvancedNew(2, op, flags,
                               NPY_KEEPORDER, casting,
                               op_flags,
                               op_dtypes,
                               -1, NULL, NULL, buffersize);
    if (iter == NULL) {
        goto fail;
    }

    /* Start with the floating-point exception flags cleared */
    npy_clear_floatstatus_barrier((char*)&iter);

    if (NpyIter_GetIterSize(iter) != 0) {
        NpyIter_IterNextFunc *iternext;
        char **dataptr;
        npy_intp *strideptr;
        npy_intp *countptr;
        int needs_api;

        iternext = NpyIter_GetIterNext(iter, NULL);
        if (iternext == NULL) {
            goto fail;
        }
        dataptr = NpyIter_GetDataPtrArray(iter);
        strideptr = NpyIter_GetInnerStrideArray(iter);
        countptr = NpyIter_GetInnerLoopSizePtr(iter);

        needs_api = NpyIter_IterationNeedsAPI(iter);

        /* Straightforward reduction */
        if (loop == NULL) {
            PyErr_Format(PyExc_RuntimeError,
                    "reduction operation %s did not supply an "
                    "inner loop function", funcname);
            goto fail;
        }

        if (loop(iter, dataptr, strideptr, countptr,
                        iternext, needs_api, skip_first_count, data) < 0) {

            goto fail;
        }
    }
    
    /* Check whether any errors occurred during the loop */
    if (PyErr_Occurred() ||
            _check_ufunc_fperr(errormask, NULL, "reduce") < 0) {
        goto fail;
    }

    NpyIter_Deallocate(iter);
    Py_DECREF(op_view);

finish:
    /* Strip out the extra 'one' dimensions in the result */
    if (out == NULL) {
        if (!keepdims) {
            PyArray_RemoveAxesInPlace(result, axis_flags);
        }
    }
    else {
        PyArray_ResolveWritebackIfCopy(result); /* prevent spurious warnings */
        Py_DECREF(result);
        result = out;
        Py_INCREF(result);
    }

    return result;

fail:
    PyArray_ResolveWritebackIfCopy(result); /* prevent spurious warnings */
    Py_XDECREF(result);
    Py_XDECREF(op_view);
    if (iter != NULL) {
        NpyIter_Deallocate(iter);
    }

    return NULL;
}