NumPy Intro

Why it exists

A Python list is flexible and slow for numbers: each item is a separate object, and a loop over a million of them is a million interpreted steps. NumPy stores numbers in one block of memory and does the loop in C.

import numpy as np

a = np.array([1, 2, 3, 4])
print(a)
print(a.dtype.kind, a.shape)

[1 2 3 4]
i (4,)

Note the printing: no commas. That is how you can tell an array from a list at a glance. The kind is i for integer; the exact width behind it is int64 on most desktops and int32 on Windows and here in the browser, so do not rely on it unless you set dtype yourself.

Arithmetic applies to everything

This is the whole idea. No loop, no comprehension:

import numpy as np

a = np.array([1, 2, 3, 4])
print(a * 2)
print(a + 10)
print(a ** 2)
print(a + np.array([10, 20, 30, 40]))

[2 4 6 8]
[11 12 13 14]
[ 1  4  9 16]
[11 22 33 44]

With a list, a * 2 would repeat it and a + 10 would be a TypeError.

Making arrays

import numpy as np

print(np.zeros(3))
print(np.ones((2, 3)))
print(np.arange(0, 10, 3))
print(np.linspace(0, 1, 5))

[0. 0. 0.]
[[1. 1. 1.]
 [1. 1. 1.]]
[0 3 6 9]
[0.   0.25 0.5  0.75 1.  ]

arange is like range with a step; linspace gives a fixed number of evenly spaced values, which is usually what you want for plotting.

One type per array

import numpy as np

print(np.array([1, 2, 3]).dtype.kind)
print(np.array([1.0, 2, 3]).dtype.kind)
print(np.array([1, 2, 3], dtype=float))

i
f
[1. 2. 3.]

Everything is promoted to a single type. That fixed layout is where the speed comes from.

Shapes

import numpy as np

grid = np.array([[1, 2, 3], [4, 5, 6]])
print(grid.shape)
print(grid[1, 2])
print(grid[0])
print(grid.reshape(3, 2))

(2, 3)
6
[1 2 3]
[[1 2]
 [3 4]
 [5 6]]

Indexing a grid is grid[row, column], not grid[row][column] — though both work.

Selecting with a condition

A comparison gives an array of booleans, and that can index the array:

import numpy as np

a = np.array([1, 5, 3, 8, 2])
print(a > 3)
print(a[a > 3])
print(a[a > 3].sum())

[False  True False  True False]
[5 8]
13

This is called boolean masking, and it replaces most of the filtering loops you would otherwise write.

Summarising

import numpy as np

a = np.array([[1, 2], [3, 4]])
print(a.sum(), a.mean(), a.max())
print(a.sum(axis=0))
print(a.sum(axis=1))

10 2.5 4
[4 6]
[3 7]

axis=0 collapses down the columns, axis=1 across the rows. Getting those the right way round is most of learning NumPy.

A slice is a view

Unlike a list slice, a NumPy slice shares memory with the original:

import numpy as np

a = np.array([1, 2, 3, 4])
piece = a[0:2]
piece[0] = 99
print(a)

safe = a[0:2].copy()
safe[0] = 0
print(a)

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

That avoids copying large data, and it means an edit through a slice changes the original. Call .copy() when you do not want that.