Faster implementation of argsreduce in scipy.stats.distributions

[pbrod]

The implementation of argsreduce below is about 10% faster than the original implementation in scipy.stats.distributions.py

def argsreduce(cond, *args):
    """ Return the sequence of ravel(args[i]) where ravel(condition) is True in
1D

    Example
      >>> import numpy as np
      >>> rand = np.random.random_sample
      >>> A = rand((4,5))
      >>> B = 2
      >>> C = rand((1,5))
      >>> cond = np.ones(A.shape)
      >>> [A1,B1,C1] = argsreduce(cond,A,B,C)
      >>> B1.shape
      (20,)
      >>> cond[2,:] = 0
      >>> [A2,B2,C2] = argsreduce(cond,A,B,C)
      >>> B2.shape
      (15,)

    """

    newargs = atleast_1d(*args)
    if not isinstance(newargs,list):
        newargs = [newargs,]
    expand_arr = (cond==cond)
    return [extract(cond,arr1*expand_arr) for arr1 in newargs]}}}

[stefan]

Thank you for optimising argsreduce. Would you please add some tests, and reformat the docstring according to the NumPy? docstring format? As soon as that is done, we can merge the patch. Thanks!

[pbrod]

Tests for argsreduce

[pbrod]

Replying to stefan:

Thank you for optimising argsreduce. Would you please add some tests, and reformat the docstring according to the NumPy? docstring format? As soon as that is done, we can merge the patch. Thanks!

Ok. Some test_runs of old implementation:

In [175]: A = np.ones((4,5))
In [176]: timeit out = argsreduce(cond,A,A,A,A,A,A,A,A,A)
1000 loops, best of 3: 214 µs per loop
In [177]: timeit out = argsreduce(cond,A,A,A)
10000 loops, best of 3: 77.9 µs per loop

compared to the new implementation

In [180]: timeit out = dst.argsreduce(cond,A,A,A,A,A,A,A,A,A)
10000 loops, best of 3: 194 µs per loop
In [181]: timeit out =dst.argsreduce(cond,A,A,A)
10000 loops, best of 3: 70.7 µs per loop

Some tests of argsreduce is attached in the test_argsreduce.py file. My suggestion for a new doc-string for argsreduce is:

    """ 
    Return the elements of each input array that satisfy some condition.

    Parameters
    ----------
    condition : array_like
        An array whose nonzero or True entries indicate the elements of each
        input array to extract. The shape of 'condition' must match the common
        shape of the input arrays according to the broadcasting rules in numpy.
    arg1, arg2, arg3, ... : array_like
        one or more input arrays.

    Returns
    -------
    narg1, narg2, narg3, ... : ndarray
        sequence of copies of the input arrays converted to the same size as the
        nonzero values of condition.

    Example
    -------
    >>> import numpy as np
    >>> rand = np.random.random_sample
    >>> A = rand((4,5))
    >>> B = 2
    >>> C = rand((1,5))
    >>> cond = np.ones(A.shape)
    >>> [A1,B1,C1] = argsreduce(cond,A,B,C)
    >>> B1.shape
    (20,)
    >>> cond[2,:] = 0
    >>> [A2,B2,C2] = argsreduce(cond,A,B,C)
    >>> B2.shape
    (15,)

    See also
    --------
    numpy.extract
    """

[stefan]

Perfect, thank you! We can merge this as soon as 0.7 is out the door.

[stefan]

Applied in r5622. Thanks!