note.nkmk.me

NumPy: Add new dimensions to ndarray (np.newaxis, np.expand_dims)

Posted: 2020-09-24 / Tags: Python, NumPy

To add new dimensions (increase dimensions) to the NumPy array ndarray, you can use np.newaxis, np.expand_dims() and np.reshape() (or reshape() method of ndarray).

This article describes the following contents.

  • How to use np.newaxis
    • np.newaxis is None
    • Add new dimensions with np.newaxis
    • Control broadcasting with np.newaxis
  • Add a new dimension with np.expand_dims()
  • np.reshape()

You can use np.reshape() or reshape() method of ndarray to not only add dimensions but also change to any shape.

You can use np.squeeze() to remove dimensions of size 1.

Sponsored Link

How to use np.newaxis

np.newaxis is None

np.newaxis is an alias of None.

import numpy as np

print(np.newaxis is None)
# True

It's just given an alias to make it easier to understand. If you replace np.newaxis in the sample code below with None, it works the same way.

Add new dimensions with np.newaxis

Using np.newaxis inside [] adds a new dimension of size 1 at that position.

a = np.arange(6).reshape(2, 3)
print(a)
# [[0 1 2]
#  [3 4 5]]

print(a.shape)
# (2, 3)

print(a[:, :, np.newaxis])
# [[[0]
#   [1]
#   [2]]
# 
#  [[3]
#   [4]
#   [5]]]

print(a[:, :, np.newaxis].shape)
# (2, 3, 1)
print(a[:, np.newaxis, :])
# [[[0 1 2]]
# 
#  [[3 4 5]]]

print(a[:, np.newaxis, :].shape)
# (2, 1, 3)
print(a[np.newaxis, :, :])
# [[[0 1 2]
#   [3 4 5]]]

print(a[np.newaxis, :, :].shape)
# (1, 2, 3)

The trailing : in [] can be omitted.

print(a[:, np.newaxis])
# [[[0 1 2]]
# 
#  [[3 4 5]]]

print(a[:, np.newaxis].shape)
# (2, 1, 3)
print(a[np.newaxis])
# [[[0 1 2]
#   [3 4 5]]]

print(a[np.newaxis].shape)
# (1, 2, 3)

Consecutive : can be replaced with .... If you want to add a new dimension to the last dimension of ndarray, which has many dimensions, it is easier to use ....

print(a[..., np.newaxis])
# [[[0]
#   [1]
#   [2]]
# 
#  [[3]
#   [4]
#   [5]]]

print(a[..., np.newaxis].shape)
# (2, 3, 1)

You can use multiple np.newaxis at once. Multiple dimensions are added.

print(a[np.newaxis, :, np.newaxis, :, np.newaxis])
# [[[[[0]
#     [1]
#     [2]]]
# 
# 
#   [[[3]
#     [4]
#     [5]]]]]

print(a[np.newaxis, :, np.newaxis, :, np.newaxis].shape)
# (1, 2, 1, 3, 1)

Adding a dimension by np.newaxis returns a view of the original object. Because the original object and the view object share memory, changing one element modifies the other element.

a_newaxis = a[:, :, np.newaxis]

print(np.shares_memory(a, a_newaxis))
# True

Control broadcasting with np.newaxis

In the operation of two NumPy arrays ndarray, they are automatically reshaped into the same shape by broadcasting.

a = np.zeros(27, dtype=np.int).reshape(3, 3, 3)
print(a)
# [[[0 0 0]
#   [0 0 0]
#   [0 0 0]]
# 
#  [[0 0 0]
#   [0 0 0]
#   [0 0 0]]
# 
#  [[0 0 0]
#   [0 0 0]
#   [0 0 0]]]

print(a.shape)
# (3, 3, 3)

b = np.arange(9).reshape(3, 3)
print(b)
# [[0 1 2]
#  [3 4 5]
#  [6 7 8]]

print(b.shape)
# (3, 3)

print(a + b)
# [[[0 1 2]
#   [3 4 5]
#   [6 7 8]]
# 
#  [[0 1 2]
#   [3 4 5]
#   [6 7 8]]
# 
#  [[0 1 2]
#   [3 4 5]
#   [6 7 8]]]

In broadcast, a new dimension is added to the beginning of the array with the smaller number of dimensions.

If you add a new dimension to the beginning with np.newaxis, the result will be the same as if it was automatically converted by broadcasting.

print(b[np.newaxis, :, :].shape)
# (1, 3, 3)

print(a + b[np.newaxis, :, :])
# [[[0 1 2]
#   [3 4 5]
#   [6 7 8]]
# 
#  [[0 1 2]
#   [3 4 5]
#   [6 7 8]]
# 
#  [[0 1 2]
#   [3 4 5]
#   [6 7 8]]]

Changing the position to add will give different results.

print(b[:, np.newaxis, :].shape)
# (3, 1, 3)

print(a + b[:, np.newaxis, :])
# [[[0 1 2]
#   [0 1 2]
#   [0 1 2]]
# 
#  [[3 4 5]
#   [3 4 5]
#   [3 4 5]]
# 
#  [[6 7 8]
#   [6 7 8]
#   [6 7 8]]]
print(b[:, :, np.newaxis].shape)
# (3, 3, 1)

print(a + b[:, :, np.newaxis])
# [[[0 0 0]
#   [1 1 1]
#   [2 2 2]]
# 
#  [[3 3 3]
#   [4 4 4]
#   [5 5 5]]
# 
#  [[6 6 6]
#   [7 7 7]
#   [8 8 8]]]

For example, if you want to add or subtract arrays of color image (shape: (height, width, color)) and monochromatic image (shape: (height, width)), it is impossible to broadcast the image as it is, but adding a new dimension at the end of the monochromatic image works well.

Add a new dimension with np.expand_dims()

You can also add a new dimension to ndarray using np.expand_dims().

Specify the original ndarray in the first argument a and the position to add the dimension in the second argument axis.

a = np.arange(6).reshape(2, 3)
print(a)
# [[0 1 2]
#  [3 4 5]]

print(np.expand_dims(a, 0))
# [[[0 1 2]
#   [3 4 5]]]

print(np.expand_dims(a, 0).shape)
# (1, 2, 3)

You can insert a new dimension at any position as follows:

print(np.expand_dims(a, 0).shape)
# (1, 2, 3)

print(np.expand_dims(a, 1).shape)
# (2, 1, 3)

print(np.expand_dims(a, 2).shape)
# (2, 3, 1)

A negative value can be specified for the second argument axis. -1 corresponds to the last dimension and you can specify the position from behind.

print(np.expand_dims(a, -1).shape)
# (2, 3, 1)

print(np.expand_dims(a, -2).shape)
# (2, 1, 3)

print(np.expand_dims(a, -3).shape)
# (1, 2, 3)

In NumPy 1.17, specifying a value such as axis > a.ndim or axis < -a.ndim - 1 in the second argument axis does not cause an error, and the dimension is added at the end or the beginning .

However, as the warning message says, it will cause an error in the future, so you should avoid it.

print(np.expand_dims(a, 3).shape)
# (2, 3, 1)
# 
# /usr/local/lib/python3.7/site-packages/ipykernel_launcher.py:1: DeprecationWarning: Both axis > a.ndim and axis < -a.ndim - 1 are deprecated and will raise an AxisError in the future.
#   """Entry point for launching an IPython kernel.

print(np.expand_dims(a, -4).shape)
# (2, 1, 3)
# 
# /usr/local/lib/python3.7/site-packages/ipykernel_launcher.py:1: DeprecationWarning: Both axis > a.ndim and axis < -a.ndim - 1 are deprecated and will raise an AxisError in the future.
#   """Entry point for launching an IPython kernel.

Only integer values can be specified in the second argument axis. It is not possible to add multiple dimensions at once by specifying multiple positions with list or tuple.

# print(np.expand_dims(a, (0, 1)).shape)
# TypeError: '>' not supported between instances of 'tuple' and 'int'

As with np.newaxis, np.expand_dims() returns a view.

a_expand_dims = np.expand_dims(a, 0)

print(np.shares_memory(a, a_expand_dims))
# True

It is of course possible to control broadcasting by adding a new dimension with np.expand_dims() as in the example of np.newaxis above.

Sponsored Link

np.reshape()

You can reshape ndarray with np.reshape() or reshape() method of ndarray. See the following article for details.

If you specify a shape with a new dimension to reshape(), the result is, of course, the same as when using np.newaxis or np.expand_dims().

a = np.arange(6).reshape(2, 3)
print(a)
# [[0 1 2]
#  [3 4 5]]

print(a.shape)
# (2, 3)

print(a[np.newaxis])
# [[[0 1 2]
#   [3 4 5]]]

print(a[np.newaxis].shape)
# (1, 2, 3)

print(np.expand_dims(a, 0))
# [[[0 1 2]
#   [3 4 5]]]

print(np.expand_dims(a, 0).shape)
# (1, 2, 3)

print(a.reshape(1, 2, 3))
# [[[0 1 2]
#   [3 4 5]]]

print(a.reshape(1, 2, 3).shape)
# (1, 2, 3)

As you can see from the above example, using np.newaxis and np.expand_dims() has the advantage that you don't have to explicitly specify the size of the original dimension.

Even with reshape(), if you want to add a dimension to the beginning or end, you do not have to explicitly specify the size by unpacking the original shape with *.

print(a.reshape(1, *a.shape))
# [[[0 1 2]
#   [3 4 5]]]

print(a.reshape(1, *a.shape).shape)
# (1, 2, 3)
Sponsored Link
Share

Related Categories

Related Articles