reduce.py 文件源码

def _ireduce_ufunc_existing_axis(arrays, ufunc, **kwargs):
    """
    Reduction operation for arrays, in the direction of an existing axis.

    Parameters
    ----------
    arrays : iterable
        Arrays to be reduced.
    ufunc : numpy.ufunc
        Binary universal function. Must have a signature of the form ufunc(x1, x2, ...)
    kwargs
        Keyword arguments are passed to ``ufunc``. The ``out`` parameter is ignored.

    Yields 
    ------
    reduced : ndarray
    """
    arrays = iter(arrays)
    first = next(arrays)

    if kwargs['axis'] not in range(first.ndim):
        raise ValueError('Axis {} not supported on arrays of shape {}.'.format(kwargs['axis'], first.shape))

    # Remove the out-parameter if provided.
    kwargs.pop('out', None)

    dtype = kwargs.get('dtype')
    if dtype is None:
        dtype = first.dtype

    axis_reduce = partial(ufunc.reduce, **kwargs)

    accumulator = np.atleast_1d(axis_reduce(first))
    yield accumulator

    # On the first pass of the following loop, accumulator is missing a dimensions
    # therefore, the stacking function cannot be 'concatenate'
    second = next(arrays)
    accumulator = np.stack([accumulator, np.atleast_1d(axis_reduce(second))], axis = -1)
    yield accumulator

    # On the second pass, the new dimensions exists, and thus we switch to
    # using concatenate.
    for array in arrays:
        reduced = np.expand_dims(np.atleast_1d(axis_reduce(array)), axis = accumulator.ndim - 1)
        accumulator = np.concatenate([accumulator, reduced], axis = accumulator.ndim - 1)
        yield accumulator