[Numpy-discussion] allbut2ax_iterator

josef.pktd@gmai... josef.pktd@gmai...
Mon Jun 7 08:59:53 CDT 2010


Maybe I missed something, but I didn't find anything like this in numpy

a kind of apply along 2axis, iterates over all 2d arrays defined by
axes argument in a n-dim array with n>2

useful for functions that operate on a 2d array (chisquare_twoway,
np.corrcoef ? ...)

awful function name

Josef

def allbut2ax_iterator(table, axes=(-2,-1)):
    axes = np.arange(table.ndim)[list(axes)]
    othax = [i for i in np.arange(table.ndim) if not i in axes]
    othshape = np.array(table.shape)[othax]
    for othind in np.ndindex(*othshape):
        index = np.array([slice(None)]*table.ndim)
        index[othax] = othind
        yield index.tolist()


>>> for i in allbut2ax_iterator(np.arange(3*2*4*4).reshape(3,2,4,4), axes=(-2,-1)): print i

[0, 0, slice(None, None, None), slice(None, None, None)]
[0, 1, slice(None, None, None), slice(None, None, None)]
[1, 0, slice(None, None, None), slice(None, None, None)]
[1, 1, slice(None, None, None), slice(None, None, None)]
[2, 0, slice(None, None, None), slice(None, None, None)]
[2, 1, slice(None, None, None), slice(None, None, None)]
>>> for i in allbut2ax_iterator(np.arange(3*2*4*4).reshape(3,2,4,4), axes=(0,3)): print i

[slice(None, None, None), 0, 0, slice(None, None, None)]
[slice(None, None, None), 0, 1, slice(None, None, None)]
[slice(None, None, None), 0, 2, slice(None, None, None)]
[slice(None, None, None), 0, 3, slice(None, None, None)]
[slice(None, None, None), 1, 0, slice(None, None, None)]
[slice(None, None, None), 1, 1, slice(None, None, None)]
[slice(None, None, None), 1, 2, slice(None, None, None)]
[slice(None, None, None), 1, 3, slice(None, None, None)]
>>> for i in allbut2ax_iterator(np.arange(3*2*4*4).reshape(3,2,4,4), axes=(1,2)): print i

[0, slice(None, None, None), slice(None, None, None), 0]
[0, slice(None, None, None), slice(None, None, None), 1]
[0, slice(None, None, None), slice(None, None, None), 2]
[0, slice(None, None, None), slice(None, None, None), 3]
[1, slice(None, None, None), slice(None, None, None), 0]
[1, slice(None, None, None), slice(None, None, None), 1]
[1, slice(None, None, None), slice(None, None, None), 2]
[1, slice(None, None, None), slice(None, None, None), 3]
[2, slice(None, None, None), slice(None, None, None), 0]
[2, slice(None, None, None), slice(None, None, None), 1]
[2, slice(None, None, None), slice(None, None, None), 2]
[2, slice(None, None, None), slice(None, None, None), 3]
>>> for i in allbut2ax_iterator(np.arange(2*4*4).reshape(2,4,4), axes=(1,2)): print i

[0, slice(None, None, None), slice(None, None, None)]
[1, slice(None, None, None), slice(None, None, None)]
>>> for i in allbut2ax_iterator(np.arange(2*4*4).reshape(2,4,4), axes=(0,2)): print i

[slice(None, None, None), 0, slice(None, None, None)]
[slice(None, None, None), 1, slice(None, None, None)]
[slice(None, None, None), 2, slice(None, None, None)]
[slice(None, None, None), 3, slice(None, None, None)]

>>> for i in allbut2ax_iterator(np.arange(2*4).reshape(2,4), axes=(0,1)): print i

[slice(None, None, None), slice(None, None, None)]


>>> table = np.arange(2*4*4).reshape(2,4,4)
>>> for i in allbut2ax_iterator(table, axes=(1,2)): print table[i]

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]
[[16 17 18 19]
 [20 21 22 23]
 [24 25 26 27]
 [28 29 30 31]]


More information about the NumPy-Discussion mailing list