NumPy: Arrays and Vectorization
import numpy as np, pandas as pd
import matplotlib.pyplot as plt
plt.rcParams.update({"figure.dpi":110,"axes.grid":True,"grid.alpha":0.25,"font.size":11})
TL, TL2, AM, RD = "#0f766e", "#2dd4bf", "#d97706", "#dc2626"
BASE = "https://raw.githubusercontent.com/johnfisher-ai/Statistics-Data-Science-AI-Visual-Book/main/data/"
fn = "python-for-data-analysis--store-orders.xlsx"
def load(sheet="Orders"):
try: return pd.read_excel("../../data/" + fn, sheet_name=sheet)
except FileNotFoundError: return pd.read_excel(BASE + fn, sheet_name=sheet)
DEMO 1 · Creating arrays¶
An array is a grid of numbers of one type. Build one from a list, or generate it with arange, linspace, zeros, ones, or a random generator.
a = np.array([2, 4, 6, 8])
print("from a list:", a, "| dtype:", a.dtype, "| shape:", a.shape)
print("arange: ", np.arange(0, 10, 2))
print("linspace:", np.linspace(0, 1, 5))
print("zeros: ", np.zeros(4))
rng = np.random.default_rng(0)
print("random: ", rng.normal(100, 15, 4).round(1))
from a list: [2 4 6 8] | dtype: int64 | shape: (4,) arange: [0 2 4 6 8] linspace: [0. 0.25 0.5 0.75 1. ] zeros: [0. 0. 0. 0.] random: [101.9 98. 109.6 101.6]
DEMO 2 · Vectorization: operate on the whole array at once¶
The core move. Instead of looping element by element, you write one expression on the whole array and NumPy runs it in fast compiled code. It is shorter to read and dramatically faster.
import timeit
big = np.arange(1_000_000, dtype=float)
loop = [x * 2 + 1 for x in big] # the same work two ways
vec = big * 2 + 1
# timeit runs each version three times and we keep the fastest, so one slow run doesn't skew the race
t_loop = min(timeit.repeat(lambda: [x * 2 + 1 for x in big], number=1, repeat=3))
t_vec = min(timeit.repeat(lambda: big * 2 + 1, number=1, repeat=3))
print(f"python loop: {t_loop*1000:6.1f} ms")
print(f"numpy vector:{t_vec*1000:6.1f} ms -> about {t_loop/t_vec:.0f}x faster, and one line")
print("same result:", np.allclose(loop, vec))
python loop: 140.0 ms numpy vector: 4.5 ms -> about 31x faster, and one line same result: True
DEMO 3 · Indexing, slicing, and boolean masks¶
Index and slice like a list, but the real power is the boolean mask: write a condition and NumPy returns just the elements that satisfy it. This is exactly how filtering works in pandas.
prices = np.array([39, 12, 119, 6, 45, 22, 34, 9])
print("first three:", prices[:3], "| every other:", prices[::2])
mask = prices >= 30
print("mask:", mask)
print("prices >= 30:", prices[mask])
print("count >= 30:", mask.sum(), "| their mean:", prices[mask].mean().round(2))
prices[prices < 10] = 10 # floor every small price at 10, in place
print("floored:", prices)
first three: [ 39 12 119] | every other: [ 39 119 45 34] mask: [ True False True False True False True False] prices >= 30: [ 39 119 45 34] count >= 30: 4 | their mean: 59.25 floored: [ 39 12 119 10 45 22 34 10]
DEMO 4 · Broadcasting: combining different shapes¶
Broadcasting lets NumPy stretch a smaller array across a bigger one so they line up, without you writing a loop. A scalar broadcasts to every element; a row broadcasts down every row.
cart = np.array([[2, 39.0], [1, 119.0], [3, 6.0]]) # columns: quantity, price
line_totals = cart[:, 0] * cart[:, 1] # elementwise across rows
print("line totals:", line_totals, "| order total:", line_totals.sum())
monthly = np.array([100.0, 120.0, 90.0, 140.0])
print("all up 10%:", (monthly * 1.10).round(1)) # scalar broadcasts to every month
print("as share of total:", (monthly / monthly.sum()).round(3))
line totals: [ 78. 119. 18.] | order total: 215.0 all up 10%: [110. 132. 99. 154.] as share of total: [0.222 0.267 0.2 0.311]
DEMO 5 · Aggregations and the axis argument¶
Reduce many numbers to a summary: sum, mean, std, min, max. On a 2-D array the axis argument chooses the direction: axis=0 collapses rows (a per-column summary), axis=1 collapses columns (a per-row summary).
grid = np.array([[120, 80, 60], [200, 40, 90], [30, 150, 70]]) # rows=regions, cols=months
print("grand total:", grid.sum())
print("per-column total (axis=0):", grid.sum(axis=0)) # down each column
print("per-row total (axis=1):", grid.sum(axis=1)) # across each row
print("column means:", grid.mean(axis=0).round(1), "| overall std:", grid.std().round(1))
grand total: 840 per-column total (axis=0): [350 270 220] per-row total (axis=1): [260 330 250] column means: [116.7 90. 73.3] | overall std: 51.6
DEMO 6 · NumPy on the real order data¶
Pull one column out of the store dataset as a NumPy array and answer questions with pure vectorized code, no loops. Under the hood, this is what pandas is doing for you.
rev = load()["revenue"].to_numpy()
print(f"{rev.size:,} orders | sum {rev.sum():,.2f} | mean {rev.mean():.2f} | median {np.median(rev):.2f} | std {rev.std():.2f}")
big = rev[rev > 100]
print(f"orders over 100 dollars: {big.size} ({big.size/rev.size:.1%}), and they are {big.sum()/rev.sum():.1%} of all revenue")
z = (rev - rev.mean()) / rev.std() # standardize in one vectorized line
print("largest order is", round(z.max(), 1), "standard deviations above the mean")
3,000 orders | sum 236,108.60 | mean 78.70 | median 56.00 | std 77.18 orders over 100 dollars: 820 (27.3%), and they are 59.3% of all revenue largest order is 6.7 standard deviations above the mean
Wrap-up¶
One idea did all the work: put numbers in an array and act on the whole array at once. Vectorization made the code shorter and roughly two orders of magnitude faster, boolean masks did the filtering, broadcasting aligned shapes, and the axis argument aimed the aggregations. pandas is NumPy with labels bolted on, so every reflex here carries straight into the next notebook.