Overview
As powerful as NumPy is for numerical computations, one common transformation needed in data manipulation and machine learning tasks is altering the dimensions of your arrays. Understanding how to increase the dimensions of a numpy array is essential as it enables seamless compatibility with various machine learning libraries that often require input data to be shaped in a certain way.
In this tutorial, we will explore two primary tools for dimension manipulation in NumPy: np.newaxis and np.expand_dims. Both are intuitive and serve to simplify your data reshaping needs, although its seemingly simple, this knowledge can have important impacts on how you handle and process data. We’ll start with the basics and gradually work our way up to more advanced examples.
Reshaping with np.newaxis
np.newaxis is used to increase the dimensions of the existing array by one more dimension, when used once. This is typically used to convert a one-dimensional array into a row vector or a column vector.
Example #1: Adding a New Axis to Convert a 1D Array into a Column Vector
import numpy as np
# Create a 1D numpy array
a = np.array([1, 2, 3, 4])
# Convert it into a column vector using np.newaxis
a_column = a[:, np.newaxis]
# Output
print(a_column)
Output:
[[1]
[2]
[3]
[4]]
In the above example, we used [:, np.newaxis] to add a new axis along the second dimension, effectively transforming our vector into a column vector.
Similarly, we can convert a 1D array into a row vector:
import numpy as np
# Create a 1D numpy array
a = np.array([1, 2, 3, 4])
# Convert it into a row vector using np.newaxis
a_row = a[np.newaxis, :]
# Output
print(a_row)
Output:
[[1 2 3 4]]
Here, we used [np.newaxis, :] which adds a new axis along the first dimension.
Understanding np.expand_dims
While np.newaxis is quite handy, NumPy also provides a function called np.expand_dims. This function takes in a given array and the position at which you would like to add a new axis.
Example #2: Expanding Dimensions with np.expand_dims
import numpy as np
# Create a 1D numpy array
a = np.array([1, 2, 3, 4])
# Use np.expand_dims to add an axis at index 0
a_expanded = np.expand_dims(a, axis=0)
# Output
print(a_expanded)
# Use np.expand_dims to add an axis at index 1
a_expanded = np.expand_dims(a, axis=1)
# Output
print(a_expanded)
Outputs:
[[1 2 3 4]]
[[1]
[2]
[3]
[4]]
The function np.expand_dims has an advantage that you can specify the axis at which you want to add the new axis directly, making your code quite descriptive.
The Uses of np.newaxis and np.expand_dims
Beyond simple reshaping of 1D vectors, you can apply these functions to 2D arrays or multi-dimensional arrays as needed, for various purposes like preparing data for neural network layers which often expect inputs to have certain dimensions.
Example #3: Increasing Dimensionality of a 2D Array
import numpy as np
# Create a 2D numpy array
a = np.array([[1, 2], [3, 4]])
# Add a new axis using np.newaxis which converts it into a 3D array
a_newaxis = a[:, :, np.newaxis]
# Output
print(a_newaxis.shape)
Output:
(2, 2, 1)
The array a has been expanded from shape (2, 2) to (2, 2, 1) by adding a new axis at the third dimension.
Similarly with np.expand_dims:
import numpy as np
# Add a new axis to a 2D array at the first index
a_expanded = np.expand_dims(a, axis=1)
# Output
print(a_expanded.shape)
Output:
(2, 1, 2)
We’ve now created an array of shape (2, 1, 2) by adding a dimension at the second index.
Advanced Examples and Broadcasting
One advanced use of these dimension modifying tools is in conjunction with NumPy’s powerful broadcasting feature, which allows you to perform arithmetic operations on arrays with different shapes. Let’s explore how adding dimensions can be useful for broadcasting.
Example #4: Combining Arrays with Different Dimensions
import numpy as np
# Define two arrays with different shapes
a = np.array([1, 2, 3])
b = np.array([[0], [10], [20], [30]])
# Use broadcasting to add them. Also make use of np.newaxis
result = a[np.newaxis, :] + b
# Output
print(result)
Output:
[[ 1 2 3]
[11 12 13]
[21 22 23]
[31 32 33]]
In the example above, we added a single-dimension array a to a two-dimensional array b by temporarily increasing the dimensions of a to align with b for broadcasting.
Conclusion
In conclusion, np.newaxis and np.expand_dims are convenient tools in the NumPy library that facilitate the reshaping of arrays by increasing their dimensionality. Mastering their use can greatly simplify a wide range of data preprocessing tasks in scientific computing and machine learning. With the knowledge of these tools, you are now better equipped to structure data for various algorithms that require specific input shapes.
