multidimensional take

[Johannes Nix]

Johannes Nix <jnix at medi.physik.uni-oldenburg.de> writes:

> Now I want to have a function defined as following:
> 
> >>> a
> array([[140, 455, 325, 360, 498],
>        [372, 647, 636, 365, 462],
>        [893, 141, 776, 238, 259]])
> 
> >>> mtake(a, [[1, 0], [ 1, 1], [2, 4]])
> array([372, 647, 259])
> 
> 

Below is the solution I got after some hours - i does not has any
interpreted loop.  It differs a bit from the list comprehension
definition in that if a e.g. has shape (1, 2, 3, 5, 7) and b has shape
(9, 11, 13), mtake(a, b) will have shape (1, 2, 3, 9) and not (9, 1,
2, 3) , ergo 

>>> a2 = array((a, 10 * a))
>>> mtake(a2, [[1, 0], [ 1, 1], [2, 4]])
array([[ 372,  647,  259],
       [3720, 6470, 2590]])
>>> a2.shape
(2, 3, 5)


I've omitted the axis parameter because it seems confusing to
me, if necessary the axes can be changed by Numeric.transpose . In my
application, shapes of a and b will be (10000, 3, 3) and 
(2, 43, 24, 50) so this will save a large amount of computing time.
Perhaps this would be a useful addition to Numeric ?

Johannes

-------------------------------------------------------------------------
def mtake(ar, indices):
    """similar to take but uses last dimension of indices for
       coordinate lookup along last axes of ar.
    """
    indices = array(indices)
    ar = array(ar)
    if len(indices.shape) == 1:
        return a[indices]
    elif (len(indices.shape) == 2) and (indices.shape[-1] == 1):
        return take(a, indices)

    numrdims = len(ar.shape) - indices.shape[-1]
    strides = multiply.accumulate(array(ar.shape[numrdims+1:] + 
        (1,))[::-1])[::-1]
    newshape = ar.shape[:numrdims] + (strides[0] * ar.shape[numrdims],)
    ar = reshape(ar, newshape)
    return take(ar, add.reduce(strides * indices, -1), -1)

--------------------------------------------------------------------------