[Numpy-discussion] Indexing with callables (was: Yorick-like functionality)

[Pauli Virtanen]

Wed, 13 May 2009 13:18:45 -0700, David J Strozzi kirjoitti:
[clip]
> Many of you probably know of the interpreter yorick by Dave Munro. As a
> Livermoron, I use it all the time.  There are some built-in functions
> there, analogous to but above and beyond numpy's sum() and diff(), which
> are quite useful for common operations on gridded data. Of course one
> can write their own, but maybe they should be cleanly canonized?

+0 from me for zcen and other, having small functions probably won't hurt
   much

[clip]
> Besides zcen, yorick has builtins for "point centering", "un-zone
> centering," etc.  Also, due to its slick syntax you can give these
> things as array "indexes":
> 
> x(zcen), y(dif), z(:,sum,:)

I think you can easily subclass numpy.ndarray to offer the same feature,
see below. I don't know if we want to add this feature (indexing with
callables) to the Numpy's fancy indexing itself. Thoughts?

-----

import numpy as np
import inspect

class YNdarray(np.ndarray):
    """
    A subclass of ndarray that implements Yorick-like indexing with functions.

    Beware: not adequately tested...
    """

    def __getitem__(self, key_):
        if not isinstance(key_, tuple):
            key = (key_,)
            scalar_key = True
        else:
            key = key_
            scalar_key = False

        key = list(key)

        # expand ellipsis manually
        while Ellipsis in key:
            j = key.index(Ellipsis)
            key[j:j+1] = [slice(None)] * (self.ndim - len(key))

        # handle reducing or mutating callables
        arr = self
        new_key = []
        real_axis = 0
        for j, v in enumerate(key):
            if callable(v):
                arr2 = self._reduce_axis(arr, v, real_axis)
                new_key.extend([slice(None)] * (arr2.ndim - arr.ndim + 1))
                arr = arr2
            elif v is not None:
                real_axis += 1
                new_key.append(v)
            else:
                new_key.append(v)

        # final get
        if scalar_key:
            return np.ndarray.__getitem__(arr, new_key[0])
        else:
            return np.ndarray.__getitem__(arr, tuple(new_key))
                

    def _reduce_axis(self, arr, func, axis):
        return func(arr, axis=axis)

x = np.arange(2*3*4).reshape(2,3,4).view(YNdarray)

# Now, 

assert np.allclose(x[np.sum,...], np.sum(x, axis=0))
assert np.allclose(x[:,np.sum,:], np.sum(x, axis=1))
assert np.allclose(x[:,:,np.sum], np.sum(x, axis=2))
assert np.allclose(x[:,np.sum,None,np.sum], x.sum(axis=1).sum(axis=1)[:,None])

def get(v, s, axis=0):
    """Index `v` with slice `s` along given axis"""
    ix = [slice(None)] * v.ndim
    ix[axis] = s
    return v[ix]

def drop_last(v, axis=0):
    """Remove one element from given array in given dimension"""
    return get(v, slice(None, -1), axis)

assert np.allclose(x[:,drop_last,:], x[:,:-1,:])

def zcen(v, axis=0):
    return .5*(get(v, slice(None,-1), axis) + get(v, slice(1,None), axis))

assert np.allclose(x[0,1,zcen], .5*(x[0,1,1:] + x[0,1,:-1]))

def append_one(v, axis=0):
    """Append one element to the given array in given dimension,
    fill with ones"""
    new_shape = list(v.shape)
    new_shape[axis] += 1
    v2 = np.empty(new_shape, dtype=v.dtype)
    get(v2, slice(None, -1), axis)[:] = v
    get(v2, -1, axis)[:] = 1
    return v2

assert np.allclose(x[:,np.diff,0], np.diff(x.view(np.ndarray)[:,:,0], axis=1))
assert np.allclose(x[0,append_one,:], [[0,1,2,3],
                                       [4,5,6,7],
                                       [8,9,10,11],
                                       [1,1,1,1]])
assert np.allclose(x[:,append_one,0], [[0,4,8,1],
                                       [12,16,20,1]])