[Numpy-discussion] Subarray with with arbitrary index?

[Huaiyu Zhu]

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]]
"""