[Numpy-discussion] Efficiently defining a multidimensional array

Neil Martinsen-Burrell nmb@wartburg....
Thu Aug 27 19:53:42 CDT 2009


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


More information about the NumPy-Discussion mailing list