lesson 204 NumPy Foundations · advanced

Vectorised rules instead of loops

A loop with an if inside, run over every value, becomes one array expression: a mask for the test, np.where or np.select for the choice, and np.diff or np.cumsum for neighbours and running totals.

You'll be able to

Replace a loop that chooses a value per item with np.where(test, if_true, if_false) or np.select
Combine masks with &, | and ~, with parentheses around each comparison
Use np.diff, np.cumsum, np.abs and np.flatnonzero for changes, running totals and positions

The idea

Many rules have the shape *for each value, if it passes a test, use this, otherwise that*. In plain Python that is a loop with an if. With arrays, the test is a mask and np.where(mask, a, b) picks from a where the mask is True and from b where it is False, for every position at once. The work happens in compiled code, which is why vectorised code is usually many times faster than a Python loop over the same numbers, and shorter to read.

Two tests combine with & (and), | (or) and ~ (not). These operators bind more tightly than <, so each comparison needs its own parentheses.

Run each example to see what it prints. Running examples doesn't count toward progress.

example-1.py · Ice days at Sioux Falls

import numpy as np

tmax = np.array([1.1, 2.8, -1.0, -1.0, 1.1, -0.5, -4.9])    # 1 to 7 January 2024
tmin = np.array([-7.1, -5.5, -6.6, -4.9, -1.0, -6.6, -11.0])
print(np.where(tmax < 0, "ice day", "thaw"))
print((tmax < 0).sum(), "ice days")
hard_freeze = (tmax < 0) & (tmin < -5)
print(np.flatnonzero(hard_freeze))
print(np.diff(tmax))

example-2.py · `and` between two arrays

import numpy as np

tmax = np.array([1.1, -1.0, -4.9])
tmin = np.array([-7.1, -6.6, -11.0])
hard_freeze = tmax < 0 and tmin < -5

With more than two outcomes, np.select(conditions, choices, default=…) checks the conditions in order and uses the first one that is true at each position, like an if/elif chain. np.cumsum builds running totals, and np.abs gives sizes without signs.

example-3.py · Taxi trips: payment type, tips and long card rides

import numpy as np

# Five NYC yellow taxi trips, 15 January 2025 (TLC trip records)
distance = np.array([9.2, 0.9, 3.17, 4.86, 0.5])     # miles
fare = np.array([40.8, 7.2, 17.2, 22.06, 5.1])        # dollars
tip = np.array([10.7, 2.55, 0.0, 0.0, 1.0])
payment = np.array([1, 1, 0, 0, 1])                   # 1 card, 2 cash, 0 flex fare

kind = np.select([payment == 1, payment == 2], ["card", "cash"], default="other")
print(kind)
tip_pct = np.where(payment == 1, tip / fare * 100, 0.0)
print(tip_pct.round(1))
print(np.cumsum(fare).round(2))
long_card = (payment == 1) & (distance > 5)
print(fare[long_card].sum())

When to use it: any per-value rule over many numbers: labels, thresholds, conditional formulas, changes between neighbours. Mistakes to avoid: and/or on arrays (ValueError); missing parentheses around comparisons joined with &; forgetting that np.where evaluates both branches for every value, so a division in the unused branch still runs (guard it or divide only the masked values); and conditions in np.select in the wrong order.

Review

A per-value rule becomes a mask plus np.where (two outcomes) or np.select (an if/elif chain).
Combine masks with &, |, ~, and put each comparison in parentheses; and/or raise on arrays.
np.diff, np.cumsum, np.abs and np.flatnonzero cover changes, running totals, sizes and positions.
Next: statistics along one axis of a table, and readings with gaps.

Practise this