On 2009-08-27 16:09 , Jonathan T wrote:
> Hi,
>
> I want to define a 3-D array as the sum of two 2-D arrays as follows:
>
> C[x,y,z] := A[x,y] + B[x,z]
>
> My linear algebra is a bit rusty; is there a good way to do this that does not
> require me to loop over x,y,z? Thanks!
Numpy's broadcasting is ideal for this. Using None as an index in a
slice adds a new axis with length 1 to an array. Then, the operation of
addition broadcasts the arrays by repeating them across the singleton
dimensions. So the above operation is
C[x,y,z] = A[x,y,z] + B[x,y,z]
where A is constant across its third dimension and B is constant across
its second dimenision. For a concrete example:
In [3]: A = np.arange(10).reshape((5,2))
In [4]: B = np.arange(15).reshape((5,3))
In [5]: C = A[:,:,None] + B[:,None,:]
In [6]: C
Out[6]:
array([[[ 0, 1, 2],
[ 1, 2, 3]],
[[ 5, 6, 7],
[ 6, 7, 8]],
[[10, 11, 12],
[11, 12, 13]],
[[15, 16, 17],
[16, 17, 18]],
[[20, 21, 22],
[21, 22, 23]]])
In [7]: C.shape
Out[7]: (5, 2, 3)
In [8]: Cprime = np.empty((5,2,3))
In [9]: for x in range(5):
for y in range(2):
for z in range(3):
Cprime[x,y,z] = A[x,y] + B[x,z]
....:
In [13]: (C == Cprime).all()
Out[13]: True
-Neil
_______________________________________________
NumPy-Discussion mailing list
[email protected]
http://mail.scipy.org/mailman/listinfo/numpy-discussion
