Add arrays of different dimensions-CodePudding

I am trying to add two arrays of different dimensions. Array A has shape (20,2,2,2,2,3) and array B has shape (20). Currently, I am using np.newaxis 5 times, so B gets the same shape as A and then I add them. In my actual code A is much bigger and this forces me to write np.newaxis many times. Is there a way to avoid repeating the np.newaxis and just tell python to give B the same shape as A?

A = np.zeros([20,2,2,2,2,3])
B = np.arange(1,21)
B = B[:,np.newaxis,np.newaxis,np.newaxis,np.newaxis,np.newaxis]
C = A   B

CodePudding user response：

What you are doing is broadcasting to the number of dimensions of A, but if you look carefully, this operation you did does not make B have the same shape as A. Indeed they are still different:

>>> B[:, None, None, None, None, None].shape
(20, 1, 1, 1, 1, 1)

So this is basically applying np.expand_dims consequently five times. An alternative way is to reshape the array with extra singletons:

>>> B.reshape((-1, *(1,)*(A.ndim-1))).shape
(20, 1, 1, 1, 1, 1)

Which will reshape (*,) to (*, 1, 1, 1, 1, 1). This has the same effect as placing the np.newaxis manually.

CodePudding user response：

If you are broadcasting, this will work:

A= np.zeros([20,2,2,2,2,3])
B = np.arange(1,21)
B = B.reshape([20,1,1,1,1,1]) 
C = A   B

In a more dynamic way:

shape_a = [20,2,2,2,2,3]
A= np.zeros(shape_a)
B = np.arange(1,21)
shape_b = [shape_a[0]]  (len(shape_a)-1)*[1]
B = B.reshape(shape_b) 
C = A   B

With no broadcasting:

A = np.zeros([20,2,2,2,2,3])
B = np.arange(1,21)
C = A.copy()
C[:,0,0,0,0,0]  = B

And if you don't care about A, just the result:

C = np.zeros([20,2,2,2,2,3])
B = np.arange(1,21)
C[:,0,0,0,0,0]  = B