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

[josef.pktd at gmail.com]

On Fri, May 15, 2009 at 5:39 PM, David Huard <david.huard at gmail.com> wrote:
> Josef,
>
> You're right, you can see it as a moving average. For 1D, correlate(a,
> [5,.5]) yields what I expect but does not take an axis keyword. For the 2D
> case, I'm rather looking for
>
>>>> ndimage.filters.correlate(b,0.25*np.ones((2,2)))[1:,1:]
>
> So another one-liner... maybe not worth adding to the numpy namespace.
>

I needed some practice with slice handling. This seems to work, but
only minimally tested. It would be possible to extend it to axis being
a tuple.
ndimage is currently very fast if you give it the correct types and
crashes for wrong function arguments.

Josef

def movmean(a, k=2, axis=None):
    '''moving average along axis for window length k'''
    a = np.asarray(a, dtype=float)  # integers don't work because
return type is also integer
    if axis is None:
        kernshape = [k]*a.ndim
        kern = 1/float(k)**a.ndim * np.ones(kernshape)
        #print kern
        cut = [slice(1,None,None)]*a.ndim
        return ndimage.filters.correlate(a,kern)[cut]
    else:
        kernshape = [1]*a.ndim
        kernshape[axis] = k
        kern = 1/float(k) * np.ones(kernshape)
        #print kern
        cut = [slice(None)]*a.ndim
        cut[axis] = slice(1,None,None)
        return ndimage.filters.correlate(a,kern)[cut]

a = np.arange(5)
b = 1.0*a[:,np.newaxis]*np.arange(1,6,2)
c = b[:,:,np.newaxis]*a

print movmean(a)
print movmean(b)
print "axis=1"
print (b[:,:-1]+b[:,1:])/2
print movmean(b, axis=1)
print "axis=0"
print (b[:-1,:]+b[1:,:])/2
print movmean(b, axis=0)

print (c[:-1,:,:]+c[1:,:,:])/2
print movmean(c, axis=0)