
Following up to an earlier discussion, I've written up some wrapper functions for subarrays. The 'takes' function only works for 2d arrays, but the 'puts' function works for arbitrary array. Hopefully something like this can be included in Numeric module. Huaiyu Zhu """ Wrapper functions for dealing with arbitrary subarrays. Example: addeqs(x, [[0,2], [1,3,4], y), would add (2,3)array y to the (2,3)subarray of x comprised of rows 0, 2 and columns 1, 3, 4. A more natural notation of addeqs(x, ij, y) would be x[ij] += y, but seems difficult to do with current numpy. (How to let slice handle lists?) """ from Numeric import * def puts(x, irows, icols, v): """ puts(x, i, j, v): Put v in subgrid of 2d array x given by i, j. """ nrow, ncol = x.shape if irows is None: irows = arange(nrow) if icols is None: icols = arange(ncol) if len(shape(icols)) == 1: icols = icols[:ncol] if len(shape(irows)) == 0 or len(shape(icols)) == 0: ii = irows*ncol + icols v1 = v else: ii = (irows*ncol)[:, NewAxis] + icols[NewAxis, :] v1 = reshape(v, shape(ii)) put(x, ii, v1) def takes(x, I): """ takes(x, I): Takes a subgrid from array x. I is a list of list of subindices. """ for i in xrange(len(I)): ii = I[i] if ii is not None: x = take(x, ii, i) return x def addeqs(x, ij, y): """ Simulates x[ij] += y, where ij can be arbitray subarray. """ i, j = ij puts(x, i, j, takes(x, ij) + y) if __name__ == "__main__": a5 = arange(5) a2 = arange(2) a3 = arange(3) d = array([a3, a3+3]) print d; print b = array([a5, a5+5, a5+10, a5+15]); print b; print c = b.copy(); puts(c, None, 3, a5+1000); print c; print c = b.copy(); puts(c, a2*2, 3, a5+1000); print c; print c = b.copy(); puts(c, 2, a2*2, a5+1000); print c; print c = b.copy(); puts(c, a2*2+1, a3*2, d+1000); print c; print c = b.copy(); d1 = takes(c, (a2*2+1, a3*2)) c1 = c print d1; print puts(c, a2*2+1, a3*2, d+1000); print c; print puts(c, a2*2+1, a3*2, d1); print c; print addeqs(c, (a2*2+1, a3*2), d1*0+100); print c; print print c1; print d1 += 20; print c; print # Alas, this does not change c d2 = takes(c, (a2*2+1, None)) print d2; print print shape(c), shape(a2), shape(d2) addeqs(c, (a2*2+1, None), -d2) print c; print """ The expected results are [[0 1 2] [3 4 5]] [[ 0 1 2 3 4] [ 5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19]] [[ 0 1 2 1000 4] [ 5 6 7 1001 9] [ 10 11 12 1002 14] [ 15 16 17 1003 19]] [[ 0 1 2 1000 4] [ 5 6 7 8 9] [ 10 11 12 1001 14] [ 15 16 17 18 19]] [[ 0 1 2 3 4] [ 5 6 7 8 9] [1000 11 1001 13 14] [ 15 16 17 18 19]] [[ 0 1 2 3 4] [1000 6 1001 8 1002] [ 10 11 12 13 14] [1003 16 1004 18 1005]] [[ 5 7 9] [15 17 19]] [[ 0 1 2 3 4] [1000 6 1001 8 1002] [ 10 11 12 13 14] [1003 16 1004 18 1005]] [[ 0 1 2 3 4] [ 5 6 7 8 9] [10 11 12 13 14] [15 16 17 18 19]] [[ 0 1 2 3 4] [105 6 107 8 109] [ 10 11 12 13 14] [115 16 117 18 119]] [[ 0 1 2 3 4] [105 6 107 8 109] [ 10 11 12 13 14] [115 16 117 18 119]] [[ 0 1 2 3 4] [105 6 107 8 109] [ 10 11 12 13 14] [115 16 117 18 119]] [[105 6 107 8 109] [115 16 117 18 119]] (4, 5) (2,) (2, 5) [[ 0 1 2 3 4] [ 0 0 0 0 0] [10 11 12 13 14] [ 0 0 0 0 0]] """