Introduction to NumPy and Matplotlib
Chapter 4: Array reshaping
Adding a new axis
Ah, I feel better squeezing that singleton dimension out of my system!
But wait… apparently the singleton dimension could be useful after all!
Mission 6: Put that lonely singleton back into the array
So now I have an (M \times K \times N) array called x.
And an (M \times N) array called y.
I need to add these two arrays, but unfortunately I cannot do so because their dimensions do not match!
>>> import numpy as np
>>> x = np.random.random_sample((10, 8, 5))
>>> y = np.random.random_sample((10, 5))
>>> print(x.shape)
(10, 8, 5)
>>> print(y.shape)
(10, 5)
>>> z = x + y
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
ValueError: operands could not be broadcast together with shapes (10,8,5) (10,5)
I think y.reshape(10, 1, 5) can solve this. NumPy can then perform its array broadcasting magic to add x and y. But I find it a hassle having to explicitly specify the sizes of M (10) and N (5) each time. And they can change any time too! y.reshape(-1, 1, -1) will not work because there can at most be only one -1.
Is there an easier way to add a new singleton dimension to y?
This is your final mission to pass your probation. Help me find a better way to solve this issue!
Hint: A little bird told me that the answer can be found on this page in the official documentation
The first possible solution is to use array indexing, specifying the np.newaxis object where you want to add the dimension. The array will be ‘expanded’ at that axis.
Apparently np.newaxis is just an alias for None, which you can also use to achieve the same effect.
>>> expanded_y = y[:, np.newaxis, :]
>>> print(expanded_y.shape)
(10, 1, 5)
>>> z = x + expanded_y
>>> expanded_y = y[:, None, :]
>>> print(expanded_y.shape)
(10, 1, 5)
>>> z = x + expanded_y
Another solution is to use the np.expand_dims() function to achieve the same thing.
>>> expanded_y = np.expand_dims(y, axis=1)
>>> print(expanded_y.shape)
(10, 1, 5)
The np.newaxis approach has an advantage though. You can add expand multiple dimensions simultaneously! You will have to do it multiple times for np.expand_dims().
>>> expanded_y = y[np.newaxis, :, np.newaxis, np.newaxis, :]
>>> print(expanded_y.shape)
(1, 10, 1, 1, 5)
Another practical use is to easily convert a vector (1D array) to a 2D array. For example, you may need to perform matrix multiplication, which you cannot do with a 1D array.
>>> x = np.array([1, 2, 3])
>>> y = np.array([[4, 5, 6]])
>>> print(x.shape)
(3, )
>>> print(y.shape)
(1, 3)
>>> print(x @ y)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 1 is different from 3)
>>>
>>> x_new = x[:, np.newaxis]
>>> print(x_new.shape)
(3, 1)
>>> print(x_new @ y)
[[ 4 5 6]
[ 8 10 12]
[12 15 18]