PDA·04 · Chapter 4

NumPy Basics

Chapter 4 — arrays and vectorised computation

NumPy is what makes numerical Python fast, and pandas is built on it. Six sessions: the ndarray itself, slicing (and why slices are views), boolean masks, random numbers and ufuncs, replacing loops with array expressions, and a random walk that pulls it all together.

4.1

The ndarray

Concept

NumPy's ndarray is a grid of values that all share one type, stored in one block of memory. That is why it is fast: loops run in C, not in Python.
Operations apply to every element at once — no for loop. This is called vectorisation. Doubling a million numbers is roughly 10–100× faster than a list comprehension.
Every array has a shape (a tuple of sizes per dimension) and a dtype (the element type, such as float64 or int32).
Create arrays with np.array(list), np.zeros, np.ones, np.arange (like range), and change shape with reshape.
astype converts dtype and always makes a new array. Casting floats to ints truncates — it drops the decimals, it does not round.

Key Terms

ndarray: an n-dimensional array of one data type
Vectorisation: applying an operation to whole arrays without writing a loop
Shape: the size of each dimension, e.g. (2, 3)
dtype: the data type of every element, e.g. float64
astype: returns a copy converted to another dtype

Worked example

import numpy as np

data = np.array([[1.5, -0.1, 3], [0, -3, 6.5]])
print(data * 10)
print(data.shape, data.dtype)

arr = np.array([3.7, -1.2, -2.6, 0.5, 12.9, 10.1])
print(arr.astype(np.int32))
print(np.arange(6).reshape((2, 3)))

[[ 15.  -1.  30.]
 [  0. -30.  65.]]
(2, 3) float64
[ 3 -1 -2  0 12 10]
[[0 1 2]
 [3 4 5]]

Exercise

Activity 1

Time [x * 2 for x in range(1_000_000)] against np.arange(1_000_000) * 2 using %timeit in a notebook. How many times faster is NumPy on your machine?

Activity 2

Make a 3×4 array of the numbers 1 to 12 with arange and reshape. Print its shape, dtype, and the array divided by 2.

4.2

Indexing and Slicing

Concept

One-dimensional arrays slice like lists — but a slice is a view, not a copy. Changing the slice changes the original array.
This is deliberate: NumPy works with huge data, and copying on every slice would be slow and use lots of memory. Use .copy() when you really want a copy.
Assigning a single value to a slice broadcasts it to every element: arr[5:8] = 12.
In 2-D, arr2d[0][2] and arr2d[0, 2] are the same element. Slices combine per axis: arr2d[:2, 1:] means "first two rows, columns from 1 on".

Key Terms

View: a new array object that shares the original data
.copy(): an independent copy of the data
Axis: one dimension of an array: axis 0 = rows, axis 1 = columns
arr[i, j]: row i, column j

Worked example

arr = np.arange(10)
arr_slice = arr[5:8]
arr_slice[1] = 12345
print(arr)             # the original changed too

arr2d = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print(arr2d[0, 2])
print(arr2d[:2, 1:])

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

Exercise

Activity 1

Repeat the example with arr_slice = arr[5:8].copy(). Show that the original is now untouched.

Activity 2

From a 4×4 array of 0–15, pull out: the last row, the middle 2×2 block, and the first column as a 1-D array.

4.3

Boolean and Fancy Indexing

Concept

Comparing an array with a value gives a boolean array — a mask: names == "Bob" is True where the name matches.
Index with the mask to keep only the True positions: data[names == "Bob"] keeps the rows whose name is Bob. The mask must be the same length as that axis.
Combine masks with & (and), | (or) and ~ (not) — not the words and/or, which don't work on arrays. Wrap each condition in parentheses.
Masks work for setting values too: data[data < 0] = 0 sets every negative to zero in one line.
Fancy indexing uses an array of integers to pick positions in any order: arr[[4, 3, 0, 6]]. Unlike slicing, boolean and fancy indexing always copy the data.

Key Terms

Boolean mask: an array of True/False used to select elements
& | ~: element-wise and / or / not for masks
Fancy indexing: selecting with an array of integer positions

Worked example

names = np.array(["Bob", "Joe", "Will", "Bob", "Will", "Joe", "Joe"])
data = np.array([[4, 7], [0, 2], [-5, 6], [0, 0], [1, 2], [-12, -4], [3, 4]])

print(names == "Bob")
print(data[names == "Bob", 1])
print(data[data < 0])

arr = np.arange(32).reshape((8, 4))
print(arr[[1, 5, 7, 2], [0, 3, 1, 2]])

[ True False False  True False False False]
[7 0]
[ -5 -12  -4]
[ 4 23 29 10]

Exercise

Activity 1

Use the mask explorer below with each expression. Before pressing Select, predict how many rows or values come back.

Activity 2

Write the mask for "rows that are not Joe" two ways: with != and with ~. Check they select the same rows.

Activity 3

Set every value in data greater than 5 to 5, using one line with a mask.

4.4

Random Numbers and ufuncs

Concept

Modern NumPy makes random numbers from a generator: rng = np.random.default_rng(seed=12345). The same seed gives the same numbers every run.
rng.standard_normal((2, 3)), rng.integers(0, 10, size=5) and rng.permutation(arr) cover most needs. Older code uses np.random.randn — it still works, but prefer a generator.
A ufunc (universal function) runs element-wise: unary ones like np.sqrt and np.exp, binary ones like np.add and np.maximum.
Some ufuncs return more than one array: np.modf splits numbers into fractional and whole parts.

Key Terms

default_rng(seed): creates a reproducible random number generator
Seed: the starting value that makes random output repeatable
ufunc: a fast element-wise function over arrays

Worked example

rng = np.random.default_rng(seed=12345)
print(rng.standard_normal((2, 3)))

print(np.sqrt(np.array([1, 4, 9])))
print(np.maximum(np.array([1, 5, 2]), np.array([4, 3, 6])))

[[-1.42382504  1.26372846 -0.87066174]
 [-0.25917323 -0.07534331 -0.74088465]]
[1. 2. 3.]
[4 5 6]

Exercise

Activity 1

Simulate 20 dice rolls with rng.integers(1, 7, size=20). Run it twice with the same seed, then with no seed. What changes?

Activity 2

Make an array of 10 random normals and use np.abs, np.round and np.maximum(arr, 0) on it. Describe what each did.

4.5

Array-Oriented Programming

Concept

np.where(cond, x, y) is a vectorised if/else: take x where the condition is true, y elsewhere. Either can be a single value or an array.
Statistics are methods: sum, mean, std, min, max, argmin, cumsum. Pass axis=0 to work down columns, axis=1 across rows.
Booleans count as 1 and 0, so (arr > 0).sum() counts the positives. any() and all() test a whole mask.
arr.sort() sorts in place; np.sort(arr) returns a sorted copy.
np.unique gives the sorted distinct values; np.isin(values, [2, 3, 6]) tests membership element by element.

Key Terms

np.where: element-wise choice between two values or arrays
axis=0 / axis=1: aggregate down the columns / across the rows
cumsum: running total
np.unique: sorted distinct values

Worked example

arr = np.array([[2, -1], [-3, 4]])
print(np.where(arr > 0, 2, -2))
print(np.where(arr > 0, 2, arr))   # only replace the positives

arr = np.arange(9).reshape((3, 3))
print(arr.mean(), arr.sum(axis=0), arr.mean(axis=1))
print((np.array([1, -2, 3, -4]) > 0).sum())

[[ 2 -2]
 [-2  2]]
[[ 2 -1]
 [-3  2]]
4.0 [ 9 12 15] [1. 4. 7.]
2

Exercise

Activity 1

Given an array of test scores, use np.where to make a matching array of "pass" or "fail" (pass is 50 or more).

Activity 2

For a 5×3 array of random integers, find the mean of each column, the max of each row, and how many values are above 10.

4.6

Files, Linear Algebra and Random Walks

Concept

np.save("arr.npy", arr) and np.load("arr.npy") store an array in NumPy's binary format; np.savez stores several in one file. For tables, pandas (Chapter 6) is usually better.
@ is matrix multiplication (x.dot(y) does the same). numpy.linalg has inv, det, qr and solve.
The random walk example shows the whole chapter at work: draw 1,000 coin flips, turn them into +1/−1 steps with np.where, and take the cumsum.
Questions about the walk become one-liners: (np.abs(walk) >= 10).argmax() finds the first step where it got 10 away from zero, because argmax returns the first True.

Key Terms

.npy / .npz: NumPy binary files for one / several arrays
@: matrix multiplication operator
argmax: index of the first maximum value (first True in a mask)

Worked example

rng = np.random.default_rng(seed=12345)
draws = rng.integers(0, 2, size=1000)
steps = np.where(draws == 0, 1, -1)
walk = steps.cumsum()

print(walk.min(), walk.max())
print((np.abs(walk) >= 10).argmax())

-8 50
155

Exercise

Activity 1

Simulate 5,000 walks at once with a 2-D array (size=(5000, 1000)) and cumsum(axis=1). How many ever reach 30 or -30?

Activity 2

Solve the equations 2x + y = 5 and x - y = 1 with np.linalg.solve.