M2.2 — NumPy and Data Manipulation

Learning objectives

After this lesson, you should be able to:

• Create arrays and read their shape, number of dimensions and dtype.
• Index and slice arrays, and select with a boolean mask.
• Use whole-array arithmetic instead of loops.
• Predict the result of broadcasting two shapes.
• Reduce an array along an axis and normalise columns.
• Explain the difference between a view and a copy.

The one-sentence idea

An array is a grid of numbers with a shape, and almost everything you do with data is selecting from it, combining it, or summarising it along an axis.

When something goes wrong, print the shape.

Why NumPy

Python lists are flexible but slow for numbers, and they know nothing about rows, columns or images. NumPy adds an array type: a grid of numbers of one type, stored compactly, with fast operations that act on the whole grid at once. Images, robot trajectories, datasets and the inputs to neural networks are all handled as arrays.

NumPy is also the model for PyTorch tensors, which you meet in the next lesson. The two libraries share most of their vocabulary: shapes, indexing, slicing, broadcasting, and reductions along an axis. Learning it here pays off twice.

Arrays, shape and dtype

Every array has a shape (its size along each dimension), a number of dimensions ndim, and a dtype (the type of its numbers, such as float64). These three facts are the first things to check when code misbehaves, in the same way the shape rule in M1.1 told you whether a multiplication was valid.

import numpy as np

a = np.array([[1.0, 2.0, 3.0],
              [4.0, 5.0, 6.0]])
print(a.shape, a.ndim, a.dtype)
print(np.zeros((2, 3)))
print(np.arange(0, 1, 0.25))
print(np.linspace(0, 1, 5))

(2, 3) 2 float64
[[0. 0. 0.]
 [0. 0. 0.]]
[0.   0.25 0.5  0.75]
[0.   0.25 0.5  0.75 1.  ]

np.zeros and np.ones make arrays of a given shape, np.arange counts with a step, and np.linspace makes a given number of evenly spaced values including both ends. Passing np.array a list of lists makes a 2-dimensional array, with one inner list per row.

Indexing and slicing

Indexing works as it does for lists, with one index per dimension, separated by commas. The rules from Python still hold: counting starts at 0, and a slice a:b stops before b. A lone colon : means “everything along this dimension”.

a = np.arange(12).reshape(3, 4)     # numbers 0..11 arranged as 3 rows, 4 columns
print(a)
print(a[1, 2])        # row 1, column 2
print(a[0])           # the whole first row
print(a[:, 1])        # the whole second column
print(a[1:, 2:])      # rows 1 onward, columns 2 onward

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
6
[0 1 2 3]
[1 5 9]
[[ 6  7]
 [10 11]]

You can also select with a boolean mask: a comparison produces an array of true/false values, and using it as an index keeps only the true positions. It can be used to read or to change values.

v = np.array([0.1, -0.5, 1.2, 0.3])
print(v > 0.25)
print(v[v > 0.25])
v[v < 0] = 0.0         # replace negatives with zero
print(v)

[False False  True  True]
[1.2 0.3]
[0.1 0.  1.2 0.3]

Whole-array operations and broadcasting

Arithmetic on arrays works element by element, without a loop. a + b, a * b and np.sqrt(a) all act on every entry. This is both shorter and much faster than looping in Python.

Broadcasting lets arrays of different shapes combine. NumPy lines the shapes up from the right, and two dimensions are compatible if they are equal or one of them is 1; the size-1 dimension is then stretched to match. A common use is adding one vector to every row of a table:

traj = np.array([[0.0, 1.0],
                 [0.5, 1.5],
                 [1.0, 2.5]])        # shape (3, 2): 3 time steps, 2 joints
offset = np.array([1.0, -1.0])       # shape (2,)
print(traj + offset)                 # offset is added to each row

[[1.  0. ]
 [1.5 0.5]
 [2.  1.5]]

Shapes (3, 2) and (2,) line up from the right as 2 with 2, so the vector is applied to each of the 3 rows. When shapes do not fit, NumPy raises an error instead of guessing:

try:
    np.ones((2, 3)) + np.ones((2, 2))
except ValueError as e:
    print(type(e).__name__)

ValueError

Reducing along an axis

Functions such as sum, mean, max and std collapse an array. By default they use every entry; with axis= they collapse one dimension and keep the others. axis=0 goes down the rows, giving one result per column. axis=1 goes across the columns, giving one result per row.

print(traj.mean(axis=0))     # per joint, averaged over time
print(traj.max(axis=1))      # per time step, the largest joint value
print(traj.sum())            # everything

[0.5        1.66666667]
[1.  1.5 2.5]
6.5

A very common data step is normalisation: subtracting the mean and dividing by the standard deviation, per column, so that every feature has a comparable scale. Broadcasting does it in one line:

mean = traj.mean(axis=0)
std = traj.std(axis=0)
norm = (traj - mean) / std
print(norm.round(2))
print(np.allclose(norm.mean(axis=0), 0), np.allclose(norm.std(axis=0), 1))

[[-1.22 -1.07]
 [ 0.   -0.27]
 [ 1.22  1.34]]
True True

The same step, with statistics computed from a training set, is how action and observation values are scaled in many robot-learning pipelines. M4 and later lessons return to normalising and decoding actions.

Reshaping, transposing and matrix multiplication

reshape changes the shape without changing the data, as long as the total count of entries is the same. A -1 in one position means “work this size out for me”. .T transposes, swapping rows and columns. The @ operator is matrix multiplication, exactly the operation from M1.1:

W = np.array([[1, 2, 0],
              [0, 1, 3]])           # shape (2, 3)
x = np.array([2, 1, 1])             # shape (3,)
print(W @ x)                        # the example from M1.1
print(W.T.shape)
print(np.arange(6).reshape(2, -1).shape)

[4 4]
(3, 2)
(2, 3)

Arrays can have more than two dimensions. A colour image is usually stored with shape (height, width, 3), where the last axis holds red, green and blue. A batch of such images adds a leading axis. Selecting one colour channel is ordinary slicing:

img = np.zeros((4, 6, 3))           # height 4, width 6, three channels
print(img.shape)
print(img[:, :, 0].shape)           # the red channel as a 4 x 6 grid

(4, 6, 3)
(4, 6)

Views, copies and random numbers

A slice of an array is usually a view: it shares memory with the original, so changing the slice changes the array. Use .copy() when you want an independent array.

a = np.arange(5)
b = a[1:3]
b[0] = 99
print(a)               # the original changed too
c = a.copy()
c[0] = -1
print(a)               # the copy did not affect it

[ 0 99  2  3  4]
[ 0 99  2  3  4]

For random numbers, create a generator with a seed so that results can be repeated. Fixing the seed is how experiments and tests are made reproducible.

rng = np.random.default_rng(0)
print(rng.random(3).round(3))

[0.637 0.27  0.041]

Common misconceptions

Engineering takeaways

1. An array has a shape, a number of dimensions and a dtype. Check them first.
2. Indexing starts at 0, a slice stops before its end, and a colon means everything.
3. Boolean masks select or change the entries that satisfy a condition.
4. Whole-array arithmetic replaces loops and is much faster.
5. Broadcasting aligns shapes from the right, and dimensions must be equal or 1.
6. axis says which dimension a reduction collapses.
7. Subtracting the mean and dividing by the standard deviation per column is one line with broadcasting.
8. Slices are views. Copy when you need independence, and seed random generators for repeatability.

Know the shape of every array you touch, and most data bugs announce themselves.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

You can now handle arrays in NumPy. Next is the tensor library that most robot-learning code is written in, which uses the same ideas and adds learning.

Back to the VLA roadmap prerequisites