Concept units

Arrays in NumPy

Intermediate

Read these first if they are new to you

The idea in one paragraph

A NumPy array is a grid of numbers that all have the same type. Its shape says how many numbers lie along each direction: four points in 3D make an array of shape (4, 3), one row per point. Square brackets pick rows, columns or blocks. Arithmetic on an array acts on every number at once, and functions such as sum and mean can collapse one direction. One line of array code replaces a loop, and it runs much faster.

A table of points is an array with two axes

Take four corners of a box, each a vector with x, y and z. Write them as a table: one row per point, one column per coordinate. NumPy stores that table as an array, a block of numbers of one type laid out in a grid.

Each direction of the grid is an axis. Axis 0 runs down the rows, and axis 1 runs across the columns. The shape lists the length along each axis, so this table has shape (4, 3). The number of axes is ndim. The dtype is the type of every number: float64 means decimal numbers stored in 64 bits.

shape.py

import numpy as np

pts = np.array([[0.0, 0.0, 0.0],
                [2.0, 0.0, 0.0],
                [2.0, 4.0, 0.0],
                [0.0, 4.0, 2.0]])
print(pts.shape, pts.ndim, pts.dtype)
print(pts[2])        # one row: the third point
print(pts[:, 1])     # one column: every y
print(pts[1:3, :2])  # rows 1 and 2, columns 0 and 1

(4, 3) 2 float64
[2. 4. 0.]
[0. 0. 4. 4.]
[[2. 0.]
 [2. 4.]]

Brackets pick rows first, then columns

Inside the brackets, a comma separates the axes: pts[row, column]. Each part can be one index, a slice such as 1:3 (from 1 up to, but not including, 3), or a bare : for "all of them". If you give only one part, NumPy takes all the columns. Reading Python covers indexes and slices for plain lists.

Which numbers does each expression pick?

Predict the highlighted cells before you press a button. Then find the expression that picks the x and y of every point, and only those.

col 0 (x)col 1 (y)col 2 (z)
Result shape: press a button
Rows, then columns. One index keeps fewer axes: pts[2] is a single row of shape (3,), and pts[-1, 2] is one number. A slice keeps its axis, even if it holds one item.

Arithmetic on an array acts on every number

pts * 10 multiplies each of the 12 numbers by 10. pts + pts adds the two arrays number by number, in matching positions. These are elementwise operations. No loop appears in your code, but NumPy runs one inside, in fast compiled code.

That inner loop is why arrays matter. On the computer that ran these snippets, a Python loop that found the length of a million 3D points took about 0.3 seconds. The one-line array version, shown below, took about 0.015 seconds: about 20 times faster.

Different shapes can combine too

pts * 10 combines an array with a single number. NumPy can also combine some arrays of different shapes, such as a table of points and one point. That rule is broadcasting.

axis=0 collapses the rows; axis=1 collapses the columns

Functions such as sum, mean and max can combine numbers along one axis. The axis you name is the one that disappears. pts.sum(axis=0) adds down each column and leaves one number per column. pts.sum(axis=1) adds across each row and leaves one number per point.

The named axis is the one that disappears
pts.sum(axis=0) pts.sum(axis=1)
Shape (4, 3) becomes (3,) or (4,). Summing along axis 0 adds the four rows into one row: the column totals. Summing along axis 1 adds each row's three numbers: one total per point.
Worked example: the center of four points

The centroid of some points is their average position: the mean of the x values, of the y values and of the z values. For the four points above:

  1. Add down each column: x gives 0 + 2 + 2 + 0 = 4, y gives 0 + 0 + 4 + 4 = 8, z gives 0 + 0 + 0 + 2 = 2.
  2. Divide by the 4 points: the centroid is (1, 2, 0.5).
  3. In NumPy that is one call, pts.mean(axis=0), because it averages down the rows.

Each point's distance from the origin uses the other axis. Square every number, add across each row with axis=1, then take the square root. For the third point that is √(2² + 4² + 0²) = √20 ≈ 4.47.

axis.py

import numpy as np

pts = np.array([[0.0, 0.0, 0.0],
                [2.0, 0.0, 0.0],
                [2.0, 4.0, 0.0],
                [0.0, 4.0, 2.0]])
print(pts * 10)               # every number times 10
print(pts.sum(axis=0))        # add down each column
print(pts.mean(axis=0))       # the center of the points
print(np.sqrt((pts ** 2).sum(axis=1)))  # each point's distance from the origin

[[ 0.  0.  0.]
 [20.  0.  0.]
 [20. 40.  0.]
 [ 0. 40. 20.]]
[4. 8. 2.]
[1.  2.  0.5]
[0.         2.         4.47213595 4.47213595]
Two details this page skips

A slice such as pts[:, 1] is usually a view: it shares memory with pts, so changing it changes pts too. Call .copy() when you need a separate array. Also, arrays can have more than two axes. A batch of 8 clouds of 2048 points has shape (8, 2048, 3), and axis=-1 names the last axis.

Check yourself

  1. An array cloud has shape (1000, 3). What are the shapes of cloud[:, 2], cloud[:10] and cloud.max(axis=0)?

    Answer

    (1000,): every z value. (10, 3): the first ten points. (3,): the largest x, the largest y and the largest z, because axis 0 disappears.

  2. Now try a worked example of your own. For the points (1, 1, 0), (3, 1, 0) and (2, 4, 3), find the centroid by hand. Which NumPy call gives it?

    Answer

    Add down the columns: (6, 6, 3). Divide by 3: (2, 2, 1). The call is pts.mean(axis=0).

  3. You want the height of every point above the floor, which is its z value. Someone writes pts[2]. What does that give instead, and what is the fix?

    Answer

    pts[2] is the third point, all three coordinates. The z values are column 2 of every row: pts[:, 2].

Go further

Optional background, with more pictures and more functions.

Where this shows up

Units that build on this one