NumPy: The Numerical Engine
1 · The lesson
readPython's built-in lists are arrays of pointers to objects of any type. They are flexible, but very slow for numerical work. NumPy uses a different design. An ndarray is a single block of contiguous memory, where every element has the same fixed-width type, plus a small header that describes the shape and strides. Because the data is stored this way, NumPy can use vectorised C loops, SIMD instructions and linear algebra through BLAS. The result is often 50–200× faster than the same code in plain Python.
You've seen the basics in ds-numpy. This lesson goes further: memory layout, broadcasting rules, advanced indexing, how axis works, and the performance habits that make NumPy code fast enough for production. By the end, you should write array expressions by default, instead of a for loop over rows.
1. The Memory Model: Why NumPy Is Fast
An ndarray is three things bundled together:
1. A data buffer: a flat C array of values that all have the same width.
2. A dtype: the element type (int32, float64, bool, etc.) and its byte width.
3. A shape + strides tuple: how to read the buffer as an n-dimensional grid.
import numpy as np a = np.arange(12, dtype=np.int32).reshape(3, 4) print(a) # [[ 0 1 2 3] # [ 4 5 6 7] # [ 8 9 10 11]] print(a.dtype) # int32 print(a.shape) # (3, 4) print(a.strides) # (16, 4) 16 bytes to the next row, 4 bytes to the next column print(a.nbytes) # 48 12 ints * 4 bytes each
strides=(16, 4) is the important part. Moving down one row jumps 16 bytes (one full row of four int32s). Moving across one column jumps 4 bytes (one int32). Reshaping, slicing and transposing usually only change the strides, and no data is copied. That's why a.T is instantaneous on a 1 GB array.
Compare this with a Python list of lists. Each row is a separate PyListObject, each cell is a separate PyLongObject somewhere else in memory, and every access follows several pointers. NumPy keeps the data close together in memory, and CPUs process data like that much faster.
2. Creating Arrays Efficiently
Avoid np.array([0]*1000), because it builds a Python list first. The right constructors allocate the buffer directly:
import numpy as np np.zeros((3, 4), dtype=np.float32) # filled with 0.0, 12 floats = 48 bytes np.ones((2, 2), dtype=np.int64) # filled with 1 np.full((2, 3), 7.5) # filled with 7.5 np.empty((3, 3)) # memory is not cleared, so overwrite every value before use np.eye(4) # 4x4 identity matrix np.arange(0, 10, 0.5) # like range, but allows a float step and returns an array np.linspace(0, 1, 11) # 11 evenly spaced points from 0 to 1 inclusive
dtype choice matters for memory. A (1000, 1000) float64 array is 8 MB. The same shape as float32 is 4 MB, and as int8 it's 1 MB. Pick the smallest type that can hold your data, but only when you have a reason to. Integer overflow happens without any warning (see Common Mistakes).
For random data, use the newer Generator API, not the legacy np.random.rand / randn functions with their global state:
rng = np.random.default_rng(seed=42) # the modern API (NEP 19) rng.standard_normal((3, 4)) # 3x4 array from the standard normal N(0, 1) rng.uniform(0, 100, size=1000) # 1000 uniform floats in [0, 100) rng.integers(0, 10, size=20) # 20 ints in [0, 10) rng.choice(["A", "B", "C"], size=5, p=[0.5, 0.3, 0.2])
setup added so this can run · re-runs 2 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) np = _AutoMock('np') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import numpy as np except Exception: pass _pm_earlier = [ 'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0 1 2 3]\n# [ 4 5 6 7]\n# [ 8 9 10 11]]\n\nprint(a.dtype) # int32\nprint(a.shape) # (3, 4)\nprint(a.strides) # (16, 4) 16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes) # 48 12 ints * 4 bytes each', 'import numpy as np\n\nnp.zeros((3, 4), dtype=np.float32) # filled with 0.0, 12 floats = 48 bytes\nnp.ones((2, 2), dtype=np.int64) # filled with 1\nnp.full((2, 3), 7.5) # filled with 7.5\nnp.empty((3, 3)) # memory is not cleared, so overwrite every value before use\nnp.eye(4) # 4x4 identity matrix\nnp.arange(0, 10, 0.5) # like range, but allows a float step and returns an array\nnp.linspace(0, 1, 11) # 11 evenly spaced points from 0 to 1 inclusive', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
The legacy functions still work, but they share one global seed, which causes problems when code runs in parallel. Use default_rng in new code.
3. Vectorisation: Whole Arrays, Not Loops
The main rule for writing fast NumPy: write your computation as an operation on whole arrays, never as a Python loop over each element. Element-wise arithmetic, comparisons and unary functions are all vectorised:
import numpy as np a = np.array([1, 2, 3, 4]) b = np.array([10, 20, 30, 40]) a + b # [11, 22, 33, 44] element by element a * b # [10, 40, 90, 160] a ** 2 # [1, 4, 9, 16] np.sqrt(a) # [1., 1.41, 1.73, 2.] np.exp(a) # [2.72, 7.39, 20.09, 54.60] a > 2 # [False, False, True, True] a boolean array # The "I want a Python loop" trap result = np.empty_like(a) for i in range(len(a)): result[i] = a[i] ** 2 + b[i] # SLOW: a Python loop over a C array # The vectorised way result = a ** 2 + b # FAST: one pass in compiled C code
A vectorised expression on a million-element array runs in milliseconds. The same work in a Python loop can take seconds. Every operator and every function in np.* is built to work on whole arrays at once.
4. Broadcasting in Depth
Broadcasting is the rule that lets arrays of different shapes work together without you copying data to make the shapes match. NumPy lines up the shapes from the right. Then, for each axis, one of three things must be true: the sizes are equal, one of them is 1 (and is stretched), or one is missing (and is treated as 1).
The rules
| Operation | Result shape | Reasoning |
|---|---|---|
(3, 4) + (4,) | (3, 4) | RHS becomes (1, 4), stretched to (3, 4) |
(3, 1) * (1, 4) | (3, 4) | Both stretched to (3, 4): outer product |
(3, 4) + (3,) | error | Right-align: 4 != 3 and neither is 1 |
(3, 4) + (3, 1) | (3, 4) | Column vector stretched across columns |
(5,) + 5.0 | (5,) | Scalar broadcasts to anything |
No memory is used for the "stretching". NumPy reads the broadcast axis with a stride of zero, so the same values are reused without being copied.
Common patterns
import numpy as np # Subtract the column mean from each column (centre per feature) X = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) col_mean = X.mean(axis=0) # shape (3,) centred = X - col_mean # (3,3) - (3,) gives (3,3), and each column now has mean 0 # Z-score normalisation per column std = X.std(axis=0) z = (X - col_mean) / std # Outer product via broadcasting (no np.outer needed) u = np.array([1, 2, 3]) # shape (3,) v = np.array([10, 20, 30, 40]) # shape (4,) outer = u[:, None] * v[None, :] # (3,1) * (1,4) gives (3,4) # [[ 10 20 30 40] # [ 20 40 60 80] # [ 30 60 90 120]] # Pairwise distances (Euclidean) between rows of A and rows of B A = np.random.default_rng(0).standard_normal((100, 5)) B = np.random.default_rng(1).standard_normal((50, 5)) # (100,1,5) - (1,50,5) gives (100,50,5). Summing squares on the last axis gives (100,50) dists = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1)) print(dists.shape) # (100, 50)
X[:, None] and X[None, :] are the idioms you will use most. They add a new axis of length 1, so that broadcasting can do the rest.
5. Advanced Indexing
Besides a[0] and a[1:5], there are three more flexible ways to index.
Integer-array indexing: pick any rows or columns
import numpy as np a = np.arange(20).reshape(5, 4) # [[ 0 1 2 3] # [ 4 5 6 7] # [ 8 9 10 11] # [12 13 14 15] # [16 17 18 19]] a[[0, 2, 4]] # rows 0, 2 and 4, shape (3, 4) a[:, [3, 1, 0]] # columns in a new order, shape (5, 3) a[[0, 1, 2], [3, 2, 1]] # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]
Boolean masks: filter by condition
mask = a > 10 a[mask] # 1-D array of every element > 10: [11, 12, 13, 14, 15, 16, 17, 18, 19] a[a % 2 == 0] # all even values # Mask assignment: set every negative entry to zero x = np.array([-3, 1, -1, 5, -2]) x[x < 0] = 0 print(x) # [0 1 0 5 0]
setup added so this can run · re-runs 6 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) a = _AutoMock('a') np = _AutoMock('np') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import numpy as np except Exception: pass _pm_earlier = [ 'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0 1 2 3]\n# [ 4 5 6 7]\n# [ 8 9 10 11]]\n\nprint(a.dtype) # int32\nprint(a.shape) # (3, 4)\nprint(a.strides) # (16, 4) 16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes) # 48 12 ints * 4 bytes each', 'import numpy as np\n\nnp.zeros((3, 4), dtype=np.float32) # filled with 0.0, 12 floats = 48 bytes\nnp.ones((2, 2), dtype=np.int64) # filled with 1\nnp.full((2, 3), 7.5) # filled with 7.5\nnp.empty((3, 3)) # memory is not cleared, so overwrite every value before use\nnp.eye(4) # 4x4 identity matrix\nnp.arange(0, 10, 0.5) # like range, but allows a float step and returns an array\nnp.linspace(0, 1, 11) # 11 evenly spaced points from 0 to 1 inclusive', 'rng = np.random.default_rng(seed=42) # the modern API (NEP 19)\nrng.standard_normal((3, 4)) # 3x4 array from the standard normal N(0, 1)\nrng.uniform(0, 100, size=1000) # 1000 uniform floats in [0, 100)\nrng.integers(0, 10, size=20) # 20 ints in [0, 10)\nrng.choice(["A", "B", "C"], size=5, p=[0.5, 0.3, 0.2])', 'import numpy as np\n\na = np.array([1, 2, 3, 4])\nb = np.array([10, 20, 30, 40])\n\na + b # [11, 22, 33, 44] element by element\na * b # [10, 40, 90, 160]\na ** 2 # [1, 4, 9, 16]\nnp.sqrt(a) # [1., 1.41, 1.73, 2.]\nnp.exp(a) # [2.72, 7.39, 20.09, 54.60]\na > 2 # [False, False, True, True] a boolean array\n\n# The "I want a Python loop" trap\nresult = np.empty_like(a)\nfor i in range(len(a)):\n result[i] = a[i] ** 2 + b[i] # SLOW: a Python loop over a C array\n\n# The vectorised way\nresult = a ** 2 + b # FAST: one pass in compiled C code', 'import numpy as np\n\n# Subtract the column mean from each column (centre per feature)\nX = np.array([[1.0, 2.0, 3.0],\n [4.0, 5.0, 6.0],\n [7.0, 8.0, 9.0]])\n\ncol_mean = X.mean(axis=0) # shape (3,)\ncentred = X - col_mean # (3,3) - (3,) gives (3,3), and each column now has mean 0\n\n# Z-score normalisation per column\nstd = X.std(axis=0)\nz = (X - col_mean) / std\n\n# Outer product via broadcasting (no np.outer needed)\nu = np.array([1, 2, 3]) # shape (3,)\nv = np.array([10, 20, 30, 40]) # shape (4,)\nouter = u[:, None] * v[None, :] # (3,1) * (1,4) gives (3,4)\n# [[ 10 20 30 40]\n# [ 20 40 60 80]\n# [ 30 60 90 120]]\n\n# Pairwise distances (Euclidean) between rows of A and rows of B\nA = np.random.default_rng(0).standard_normal((100, 5))\nB = np.random.default_rng(1).standard_normal((50, 5))\n# (100,1,5) - (1,50,5) gives (100,50,5). Summing squares on the last axis gives (100,50)\ndists = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1))\nprint(dists.shape) # (100, 50)', 'import numpy as np\n\na = np.arange(20).reshape(5, 4)\n# [[ 0 1 2 3]\n# [ 4 5 6 7]\n# [ 8 9 10 11]\n# [12 13 14 15]\n# [16 17 18 19]]\n\na[[0, 2, 4]] # rows 0, 2 and 4, shape (3, 4)\na[:, [3, 1, 0]] # columns in a new order, shape (5, 3)\na[[0, 1, 2], [3, 2, 1]] # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
np.where: a vectorised if/else
x = np.array([-3, 1, -1, 5, -2]) np.where(x < 0, 0, x) # [0, 1, 0, 5, 0] the same result as np.maximum(x, 0) np.where(x < 0, "neg", "pos") # array(['neg', 'pos', 'neg', 'pos', 'neg'], dtype='<U3') # The one-argument form returns the indices where the condition is True idx = np.where(x < 0) print(idx) # (array([0, 2, 4]),) print(x[idx]) # [-3 -1 -2]
setup added so this can run · re-runs 7 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) np = _AutoMock('np') a = _AutoMock('a') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import numpy as np except Exception: pass _pm_earlier = [ 'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0 1 2 3]\n# [ 4 5 6 7]\n# [ 8 9 10 11]]\n\nprint(a.dtype) # int32\nprint(a.shape) # (3, 4)\nprint(a.strides) # (16, 4) 16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes) # 48 12 ints * 4 bytes each', 'import numpy as np\n\nnp.zeros((3, 4), dtype=np.float32) # filled with 0.0, 12 floats = 48 bytes\nnp.ones((2, 2), dtype=np.int64) # filled with 1\nnp.full((2, 3), 7.5) # filled with 7.5\nnp.empty((3, 3)) # memory is not cleared, so overwrite every value before use\nnp.eye(4) # 4x4 identity matrix\nnp.arange(0, 10, 0.5) # like range, but allows a float step and returns an array\nnp.linspace(0, 1, 11) # 11 evenly spaced points from 0 to 1 inclusive', 'rng = np.random.default_rng(seed=42) # the modern API (NEP 19)\nrng.standard_normal((3, 4)) # 3x4 array from the standard normal N(0, 1)\nrng.uniform(0, 100, size=1000) # 1000 uniform floats in [0, 100)\nrng.integers(0, 10, size=20) # 20 ints in [0, 10)\nrng.choice(["A", "B", "C"], size=5, p=[0.5, 0.3, 0.2])', 'import numpy as np\n\na = np.array([1, 2, 3, 4])\nb = np.array([10, 20, 30, 40])\n\na + b # [11, 22, 33, 44] element by element\na * b # [10, 40, 90, 160]\na ** 2 # [1, 4, 9, 16]\nnp.sqrt(a) # [1., 1.41, 1.73, 2.]\nnp.exp(a) # [2.72, 7.39, 20.09, 54.60]\na > 2 # [False, False, True, True] a boolean array\n\n# The "I want a Python loop" trap\nresult = np.empty_like(a)\nfor i in range(len(a)):\n result[i] = a[i] ** 2 + b[i] # SLOW: a Python loop over a C array\n\n# The vectorised way\nresult = a ** 2 + b # FAST: one pass in compiled C code', 'import numpy as np\n\n# Subtract the column mean from each column (centre per feature)\nX = np.array([[1.0, 2.0, 3.0],\n [4.0, 5.0, 6.0],\n [7.0, 8.0, 9.0]])\n\ncol_mean = X.mean(axis=0) # shape (3,)\ncentred = X - col_mean # (3,3) - (3,) gives (3,3), and each column now has mean 0\n\n# Z-score normalisation per column\nstd = X.std(axis=0)\nz = (X - col_mean) / std\n\n# Outer product via broadcasting (no np.outer needed)\nu = np.array([1, 2, 3]) # shape (3,)\nv = np.array([10, 20, 30, 40]) # shape (4,)\nouter = u[:, None] * v[None, :] # (3,1) * (1,4) gives (3,4)\n# [[ 10 20 30 40]\n# [ 20 40 60 80]\n# [ 30 60 90 120]]\n\n# Pairwise distances (Euclidean) between rows of A and rows of B\nA = np.random.default_rng(0).standard_normal((100, 5))\nB = np.random.default_rng(1).standard_normal((50, 5))\n# (100,1,5) - (1,50,5) gives (100,50,5). Summing squares on the last axis gives (100,50)\ndists = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1))\nprint(dists.shape) # (100, 50)', 'import numpy as np\n\na = np.arange(20).reshape(5, 4)\n# [[ 0 1 2 3]\n# [ 4 5 6 7]\n# [ 8 9 10 11]\n# [12 13 14 15]\n# [16 17 18 19]]\n\na[[0, 2, 4]] # rows 0, 2 and 4, shape (3, 4)\na[:, [3, 1, 0]] # columns in a new order, shape (5, 3)\na[[0, 1, 2], [3, 2, 1]] # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]', 'mask = a > 10\na[mask] # 1-D array of every element > 10: [11, 12, 13, 14, 15, 16, 17, 18, 19]\na[a % 2 == 0] # all even values\n\n# Mask assignment: set every negative entry to zero\nx = np.array([-3, 1, -1, 5, -2])\nx[x < 0] = 0\nprint(x) # [0 1 0 5 0]', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
Fancy indexing always creates a copy, never a view (more on that below). Boolean masks also copy. Plain slices (a[1:5]) are views, so changing them changes the original.
6. Aggregations and the axis Argument
Reductions like sum, mean, std, min, max, argmax all accept axis=. The idea: axis=k removes axis k and keeps the others.
import numpy as np a = np.array([[1, 2, 3], [4, 5, 6]]) # shape (2, 3) a.sum() # 21 adds up every element a.sum(axis=0) # [5, 7, 9] collapses the rows, giving one sum per column a.sum(axis=1) # [6, 15] collapses the columns, giving one sum per row a.mean(axis=0) # [2.5, 3.5, 4.5] the column means a.argmax(axis=1) # [2, 2] the index of the max in each row
A way to remember it: axis=0 goes "down through the rows" and gives one result per column. axis=1 goes "across the columns" and gives one result per row. It feels backwards at first, but you get used to it.
keepdims=True keeps the reduced axis with length 1. You need this to broadcast the result back onto the original array:
row_max = a.max(axis=1) # shape (2,) cannot be subtracted from (2,3) row_max = a.max(axis=1, keepdims=True) # shape (2, 1) (this one broadcasts correctly) normalised = a - row_max # subtract each row's max (the trick for a stable softmax)
setup added so this can run · re-runs 4 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) a = _AutoMock('a') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import numpy as np except Exception: pass _pm_earlier = [ 'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0 1 2 3]\n# [ 4 5 6 7]\n# [ 8 9 10 11]]\n\nprint(a.dtype) # int32\nprint(a.shape) # (3, 4)\nprint(a.strides) # (16, 4) 16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes) # 48 12 ints * 4 bytes each', 'import numpy as np\n\na = np.array([1, 2, 3, 4])\nb = np.array([10, 20, 30, 40])\n\na + b # [11, 22, 33, 44] element by element\na * b # [10, 40, 90, 160]\na ** 2 # [1, 4, 9, 16]\nnp.sqrt(a) # [1., 1.41, 1.73, 2.]\nnp.exp(a) # [2.72, 7.39, 20.09, 54.60]\na > 2 # [False, False, True, True] a boolean array\n\n# The "I want a Python loop" trap\nresult = np.empty_like(a)\nfor i in range(len(a)):\n result[i] = a[i] ** 2 + b[i] # SLOW: a Python loop over a C array\n\n# The vectorised way\nresult = a ** 2 + b # FAST: one pass in compiled C code', 'import numpy as np\n\na = np.arange(20).reshape(5, 4)\n# [[ 0 1 2 3]\n# [ 4 5 6 7]\n# [ 8 9 10 11]\n# [12 13 14 15]\n# [16 17 18 19]]\n\na[[0, 2, 4]] # rows 0, 2 and 4, shape (3, 4)\na[:, [3, 1, 0]] # columns in a new order, shape (5, 3)\na[[0, 1, 2], [3, 2, 1]] # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]', 'import numpy as np\n\na = np.array([[1, 2, 3],\n [4, 5, 6]]) # shape (2, 3)\n\na.sum() # 21 adds up every element\na.sum(axis=0) # [5, 7, 9] collapses the rows, giving one sum per column\na.sum(axis=1) # [6, 15] collapses the columns, giving one sum per row\n\na.mean(axis=0) # [2.5, 3.5, 4.5] the column means\na.argmax(axis=1) # [2, 2] the index of the max in each row', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
7. Linear Algebra: @, dot, solve
@ is the matrix-multiplication operator (PEP 465). For arrays with more than two dimensions, it multiplies a batch of matrices at once, using broadcasting.
import numpy as np A = np.array([[1, 2], [3, 4]], dtype=float) b = np.array([5, 6], dtype=float) A @ b # [17., 39.] matrix times vector A @ A # matrix times matrix, a 2x2 result A.T @ A # Gram matrix # Solve Ax = b. Do NOT compute inv(A) @ b x = np.linalg.solve(A, b) # [-4., 4.5] faster and more accurate than the inverse # When you really need the inverse A_inv = np.linalg.inv(A) # Norms: Euclidean, Manhattan and Frobenius np.linalg.norm(b) # 7.81 L2 (Euclidean) by default np.linalg.norm(b, ord=1) # 11.0 L1 (Manhattan) np.linalg.norm(A, ord="fro") # 5.48 Frobenius # Eigendecomposition eigvals, eigvecs = np.linalg.eig(A) # eigvals: [-0.37, 5.37] # eigvecs: columns are the corresponding unit eigenvectors # Singular value decomposition (the method behind PCA) U, s, Vt = np.linalg.svd(A, full_matrices=False)
For Ax = b, always use solve instead of inv @ b. It uses LU decomposition internally, it is faster, and it avoids the rounding problems of computing the inverse.
8. Statistics and Random: Beyond the Basics
import numpy as np data = np.array([2, 4, 4, 4, 5, 5, 7, 9]) np.percentile(data, [25, 50, 75]) # [4., 4.5, 5.5] the quartiles np.median(data) # 4.5 np.cov(data) # 4.57 sample covariance (for 1-D data, the variance) # Correlation matrix x = np.random.default_rng(0).standard_normal((100,)) y = 2 * x + 0.5 * np.random.default_rng(1).standard_normal((100,)) np.corrcoef(x, y) # [[1. 0.97] # [0.97 1. ]] # Statistics on 2-D data: choose the axis carefully X = np.random.default_rng(0).standard_normal((1000, 5)) # 1000 samples, 5 features feature_mean = X.mean(axis=0) # shape (5,) one mean per feature feature_std = X.std(axis=0, ddof=1) # sample std (N-1 denominator)
ddof=1 gives the unbiased sample statistic. NumPy defaults to the population version (ddof=0), but pandas defaults to the sample version. This difference causes many mistakes.
9. Reshape, Transpose, Stack, Split
Reshaping, transposing and splitting usually don't copy data. They only read the same buffer in a different way. Stacking is different: it always creates a new array.
import numpy as np a = np.arange(12) a.reshape(3, 4) # 3x4 view a.reshape(2, -1) # -1 means "work it out", so 2x6 a.reshape(2, 2, 3) # 3-D view # Transpose: swap or reorder the axes M = np.arange(24).reshape(2, 3, 4) M.T.shape # (4, 3, 2) reverses all axes M.transpose(1, 0, 2).shape # (3, 2, 4) axes in the order you give # Stacking and concatenation a = np.array([1, 2, 3]) b = np.array([4, 5, 6]) np.concatenate([a, b]) # [1, 2, 3, 4, 5, 6] np.stack([a, b]) # [[1,2,3],[4,5,6]] adds a new axis at the front np.vstack([a, b]) # [[1,2,3],[4,5,6]] stacks them as rows np.hstack([a, b]) # [1, 2, 3, 4, 5, 6] joins 1-D arrays end to end np.column_stack([a, b]) # [[1,4],[2,5],[3,6]] (each array becomes a column) # Splitting: the opposite of stacking M = np.arange(12).reshape(3, 4) np.split(M, 2, axis=1) # split into 2 along the columns, giving a list of (3,2) arrays
reshape requires the new shape to have the same total size. ravel() flattens to 1-D (a view if the data is contiguous). flatten() always copies.
10. Views vs Copies: The Trap
Slicing returns a view: the same buffer with a different shape or stride. Changing a view changes the original. Fancy indexing returns a copy: separate data.
import numpy as np a = np.arange(10) b = a[2:6] # VIEW: shares memory with a b[0] = 999 print(a) # [ 0 1 999 3 4 5 6 7 8 9] a changed too! c = a[[2, 3, 4]] # COPY: fancy indexing always copies c[0] = -1 print(a) # unchanged # Check explicitly print(b.base is a) # True b is a view of a print(c.base is a) # False c is a separate copy # Force a copy when you want one d = a[2:6].copy() d[0] = 555 print(a) # unchanged
This causes many "why did my data change?" bugs. When you plan to change the data, call .copy() first to be safe. When you want a view (for memory or speed), check it with .base.
11. Performance: Benchmarks, np.einsum and the Cost of Loops
This benchmark answers the question "should I vectorise?":
import numpy as np, time n = 1_000_000 a = np.random.default_rng(0).standard_normal(n) b = np.random.default_rng(1).standard_normal(n) # Python loop t0 = time.perf_counter() result_loop = [a[i] * b[i] + a[i] for i in range(n)] t_loop = time.perf_counter() - t0 # NumPy vectorised t0 = time.perf_counter() result_vec = a * b + a t_vec = time.perf_counter() - t0 print(f"loop: {t_loop:.3f}s vec: {t_vec:.4f}s speedup: {t_loop/t_vec:.0f}x") # loop: ~0.300s vec: ~0.003s speedup: ~100x
For complex multi-axis operations, np.einsum expresses tensor contractions in Einstein notation. It is sometimes faster than equivalent combinations of transpose, reshape, and @:
import numpy as np A = np.random.default_rng(0).standard_normal((100, 50)) B = np.random.default_rng(1).standard_normal((50, 30)) # Standard matmul np.einsum("ij,jk->ik", A, B).shape # (100, 30) the same as A @ B # Batched dot product of corresponding rows X = np.random.default_rng(0).standard_normal((1000, 5)) Y = np.random.default_rng(1).standard_normal((1000, 5)) dots = np.einsum("ij,ij->i", X, Y) # shape (1000,) (one dot product per row)
When NumPy's vectorised operations aren't enough, for example a busy inner loop with branches that can't be written as array operations, use numba (@njit JIT-compiles a Python function) or Cython. Both run that loop as compiled code instead of in the Python interpreter.
12. I/O: Saving and Loading
import numpy as np a = np.arange(12).reshape(3, 4) np.save("a.npy", a) # binary, fastest, and keeps the dtype b = np.load("a.npy") np.savez("bundle.npz", x=a, y=a * 2) # multiple arrays in one file data = np.load("bundle.npz") print(data["x"], data["y"]) np.savetxt("a.csv", a, delimiter=",", fmt="%d") # human-readable, but can lose float precision b = np.loadtxt("a.csv", delimiter=",")
.npy is the default for "save this for later". It loads back exactly the same array, including dtype, shape and byte order. .csv is for sharing data with people or tools that don't use NumPy.
Common Mistakes
1. Silent integer overflow
import numpy as np a = np.array([100, 100], dtype=np.int8) # range is [-128, 127], so 100 fits print(a + a) # [-56 -56] it WRAPPED around with no warning
NumPy integer dtypes wrap around on overflow without any warning. int8, int16 and int32 can all do this. Use int64 (the default on 64-bit platforms) unless you are short of memory. If you use a smaller type, check the range of the values when the data comes in.
2. Mixed dtypes producing object arrays
a = np.array([1, 2, "three", 4]) print(a.dtype) # <U21 Unicode strings, which are slow a = np.array([1, 2, None, 4]) print(a.dtype) # object separate Python objects, the slowest
setup added so this can run · re-runs 17 earlier blocks on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) np = _AutoMock('np') a = _AutoMock('a') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import numpy as np import numpy as np, time except Exception: pass _pm_earlier = [ 'import numpy as np\n\na = np.arange(12, dtype=np.int32).reshape(3, 4)\nprint(a)\n# [[ 0 1 2 3]\n# [ 4 5 6 7]\n# [ 8 9 10 11]]\n\nprint(a.dtype) # int32\nprint(a.shape) # (3, 4)\nprint(a.strides) # (16, 4) 16 bytes to the next row, 4 bytes to the next column\nprint(a.nbytes) # 48 12 ints * 4 bytes each', 'import numpy as np\n\nnp.zeros((3, 4), dtype=np.float32) # filled with 0.0, 12 floats = 48 bytes\nnp.ones((2, 2), dtype=np.int64) # filled with 1\nnp.full((2, 3), 7.5) # filled with 7.5\nnp.empty((3, 3)) # memory is not cleared, so overwrite every value before use\nnp.eye(4) # 4x4 identity matrix\nnp.arange(0, 10, 0.5) # like range, but allows a float step and returns an array\nnp.linspace(0, 1, 11) # 11 evenly spaced points from 0 to 1 inclusive', 'rng = np.random.default_rng(seed=42) # the modern API (NEP 19)\nrng.standard_normal((3, 4)) # 3x4 array from the standard normal N(0, 1)\nrng.uniform(0, 100, size=1000) # 1000 uniform floats in [0, 100)\nrng.integers(0, 10, size=20) # 20 ints in [0, 10)\nrng.choice(["A", "B", "C"], size=5, p=[0.5, 0.3, 0.2])', 'import numpy as np\n\na = np.array([1, 2, 3, 4])\nb = np.array([10, 20, 30, 40])\n\na + b # [11, 22, 33, 44] element by element\na * b # [10, 40, 90, 160]\na ** 2 # [1, 4, 9, 16]\nnp.sqrt(a) # [1., 1.41, 1.73, 2.]\nnp.exp(a) # [2.72, 7.39, 20.09, 54.60]\na > 2 # [False, False, True, True] a boolean array\n\n# The "I want a Python loop" trap\nresult = np.empty_like(a)\nfor i in range(len(a)):\n result[i] = a[i] ** 2 + b[i] # SLOW: a Python loop over a C array\n\n# The vectorised way\nresult = a ** 2 + b # FAST: one pass in compiled C code', 'import numpy as np\n\n# Subtract the column mean from each column (centre per feature)\nX = np.array([[1.0, 2.0, 3.0],\n [4.0, 5.0, 6.0],\n [7.0, 8.0, 9.0]])\n\ncol_mean = X.mean(axis=0) # shape (3,)\ncentred = X - col_mean # (3,3) - (3,) gives (3,3), and each column now has mean 0\n\n# Z-score normalisation per column\nstd = X.std(axis=0)\nz = (X - col_mean) / std\n\n# Outer product via broadcasting (no np.outer needed)\nu = np.array([1, 2, 3]) # shape (3,)\nv = np.array([10, 20, 30, 40]) # shape (4,)\nouter = u[:, None] * v[None, :] # (3,1) * (1,4) gives (3,4)\n# [[ 10 20 30 40]\n# [ 20 40 60 80]\n# [ 30 60 90 120]]\n\n# Pairwise distances (Euclidean) between rows of A and rows of B\nA = np.random.default_rng(0).standard_normal((100, 5))\nB = np.random.default_rng(1).standard_normal((50, 5))\n# (100,1,5) - (1,50,5) gives (100,50,5). Summing squares on the last axis gives (100,50)\ndists = np.sqrt(((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1))\nprint(dists.shape) # (100, 50)', 'import numpy as np\n\na = np.arange(20).reshape(5, 4)\n# [[ 0 1 2 3]\n# [ 4 5 6 7]\n# [ 8 9 10 11]\n# [12 13 14 15]\n# [16 17 18 19]]\n\na[[0, 2, 4]] # rows 0, 2 and 4, shape (3, 4)\na[:, [3, 1, 0]] # columns in a new order, shape (5, 3)\na[[0, 1, 2], [3, 2, 1]] # picks (0,3), (1,2) and (2,1), giving [3, 6, 9]', 'mask = a > 10\na[mask] # 1-D array of every element > 10: [11, 12, 13, 14, 15, 16, 17, 18, 19]\na[a % 2 == 0] # all even values\n\n# Mask assignment: set every negative entry to zero\nx = np.array([-3, 1, -1, 5, -2])\nx[x < 0] = 0\nprint(x) # [0 1 0 5 0]', 'x = np.array([-3, 1, -1, 5, -2])\nnp.where(x < 0, 0, x) # [0, 1, 0, 5, 0] the same result as np.maximum(x, 0)\nnp.where(x < 0, "neg", "pos")\n# array([\'neg\', \'pos\', \'neg\', \'pos\', \'neg\'], dtype=\'<U3\')\n\n# The one-argument form returns the indices where the condition is True\nidx = np.where(x < 0)\nprint(idx) # (array([0, 2, 4]),)\nprint(x[idx]) # [-3 -1 -2]', 'import numpy as np\n\na = np.array([[1, 2, 3],\n [4, 5, 6]]) # shape (2, 3)\n\na.sum() # 21 adds up every element\na.sum(axis=0) # [5, 7, 9] collapses the rows, giving one sum per column\na.sum(axis=1) # [6, 15] collapses the columns, giving one sum per row\n\na.mean(axis=0) # [2.5, 3.5, 4.5] the column means\na.argmax(axis=1) # [2, 2] the index of the max in each row', 'import numpy as np\n\nA = np.array([[1, 2], [3, 4]], dtype=float)\nb = np.array([5, 6], dtype=float)\n\nA @ b # [17., 39.] matrix times vector\nA @ A # matrix times matrix, a 2x2 result\nA.T @ A # Gram matrix\n\n# Solve Ax = b. Do NOT compute inv(A) @ b\nx = np.linalg.solve(A, b) # [-4., 4.5] faster and more accurate than the inverse\n\n# When you really need the inverse\nA_inv = np.linalg.inv(A)\n\n# Norms: Euclidean, Manhattan and Frobenius\nnp.linalg.norm(b) # 7.81 L2 (Euclidean) by default\nnp.linalg.norm(b, ord=1) # 11.0 L1 (Manhattan)\nnp.linalg.norm(A, ord="fro") # 5.48 Frobenius\n\n# Eigendecomposition\neigvals, eigvecs = np.linalg.eig(A)\n# eigvals: [-0.37, 5.37]\n# eigvecs: columns are the corresponding unit eigenvectors\n\n# Singular value decomposition (the method behind PCA)\nU, s, Vt = np.linalg.svd(A, full_matrices=False)', 'import numpy as np\n\ndata = np.array([2, 4, 4, 4, 5, 5, 7, 9])\n\nnp.percentile(data, [25, 50, 75]) # [4., 4.5, 5.5] the quartiles\nnp.median(data) # 4.5\nnp.cov(data) # 4.57 sample covariance (for 1-D data, the variance)\n\n# Correlation matrix\nx = np.random.default_rng(0).standard_normal((100,))\ny = 2 * x + 0.5 * np.random.default_rng(1).standard_normal((100,))\nnp.corrcoef(x, y)\n# [[1. 0.97]\n# [0.97 1. ]]\n\n# Statistics on 2-D data: choose the axis carefully\nX = np.random.default_rng(0).standard_normal((1000, 5)) # 1000 samples, 5 features\nfeature_mean = X.mean(axis=0) # shape (5,) one mean per feature\nfeature_std = X.std(axis=0, ddof=1) # sample std (N-1 denominator)', 'import numpy as np\n\na = np.arange(12)\na.reshape(3, 4) # 3x4 view\na.reshape(2, -1) # -1 means "work it out", so 2x6\na.reshape(2, 2, 3) # 3-D view\n\n# Transpose: swap or reorder the axes\nM = np.arange(24).reshape(2, 3, 4)\nM.T.shape # (4, 3, 2) reverses all axes\nM.transpose(1, 0, 2).shape # (3, 2, 4) axes in the order you give\n\n# Stacking and concatenation\na = np.array([1, 2, 3])\nb = np.array([4, 5, 6])\n\nnp.concatenate([a, b]) # [1, 2, 3, 4, 5, 6]\nnp.stack([a, b]) # [[1,2,3],[4,5,6]] adds a new axis at the front\nnp.vstack([a, b]) # [[1,2,3],[4,5,6]] stacks them as rows\nnp.hstack([a, b]) # [1, 2, 3, 4, 5, 6] joins 1-D arrays end to end\nnp.column_stack([a, b]) # [[1,4],[2,5],[3,6]] (each array becomes a column)\n\n# Splitting: the opposite of stacking\nM = np.arange(12).reshape(3, 4)\nnp.split(M, 2, axis=1) # split into 2 along the columns, giving a list of (3,2) arrays', 'import numpy as np\n\na = np.arange(10)\nb = a[2:6] # VIEW: shares memory with a\nb[0] = 999\nprint(a) # [ 0 1 999 3 4 5 6 7 8 9] a changed too!\n\nc = a[[2, 3, 4]] # COPY: fancy indexing always copies\nc[0] = -1\nprint(a) # unchanged\n\n# Check explicitly\nprint(b.base is a) # True b is a view of a\nprint(c.base is a) # False c is a separate copy\n\n# Force a copy when you want one\nd = a[2:6].copy()\nd[0] = 555\nprint(a) # unchanged', 'import numpy as np, time\n\nn = 1_000_000\na = np.random.default_rng(0).standard_normal(n)\nb = np.random.default_rng(1).standard_normal(n)\n\n# Python loop\nt0 = time.perf_counter()\nresult_loop = [a[i] * b[i] + a[i] for i in range(n)]\nt_loop = time.perf_counter() - t0\n\n# NumPy vectorised\nt0 = time.perf_counter()\nresult_vec = a * b + a\nt_vec = time.perf_counter() - t0\n\nprint(f"loop: {t_loop:.3f}s vec: {t_vec:.4f}s speedup: {t_loop/t_vec:.0f}x")\n# loop: ~0.300s vec: ~0.003s speedup: ~100x', 'import numpy as np\n\nA = np.random.default_rng(0).standard_normal((100, 50))\nB = np.random.default_rng(1).standard_normal((50, 30))\n\n# Standard matmul\nnp.einsum("ij,jk->ik", A, B).shape # (100, 30) the same as A @ B\n\n# Batched dot product of corresponding rows\nX = np.random.default_rng(0).standard_normal((1000, 5))\nY = np.random.default_rng(1).standard_normal((1000, 5))\ndots = np.einsum("ij,ij->i", X, Y) # shape (1000,) (one dot product per row)', 'import numpy as np\n\na = np.arange(12).reshape(3, 4)\n\nnp.save("a.npy", a) # binary, fastest, and keeps the dtype\nb = np.load("a.npy")\n\nnp.savez("bundle.npz", x=a, y=a * 2) # multiple arrays in one file\ndata = np.load("bundle.npz")\nprint(data["x"], data["y"])\n\nnp.savetxt("a.csv", a, delimiter=",", fmt="%d") # human-readable, but can lose float precision\nb = np.loadtxt("a.csv", delimiter=",")', 'import numpy as np\n\na = np.array([100, 100], dtype=np.int8) # range is [-128, 127], so 100 fits\nprint(a + a) # [-56 -56] it WRAPPED around with no warning', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
When NumPy can't find a common type, it uses a string dtype or object (an array of pointers to Python objects). You lose the speed of vectorisation, and an object array is no faster than a list. Replace missing values with np.nan (which is a float) before you create the array, or use pandas for data with mixed types.
3. In-place ops on shared views
import numpy as np a = np.arange(10) b = a[::2] # view of every other element: [0, 2, 4, 6, 8] b *= 100 # IN-PLACE, so this also changes a print(a) # [ 0 1 200 3 400 5 600 7 800 9]
*=, += and similar operators change the underlying buffer. If b is a view, a sees the change too. If you don't want that, use b = b * 100 (which creates a new array), or call .copy() before the in-place operation.
4. Using Python functions when NumPy ops exist
import math import numpy as np x = np.linspace(0, 10, 1_000_000) # Slow: a Python loop hidden inside a list comprehension y = [math.sin(v) for v in x] # ~200ms # Fast: a vectorised loop in C y = np.sin(x) # ~5ms
math.sin, math.exp, math.log are scalar functions. np.sin, np.exp, np.log are vectorised. Using the math versions on arrays loses the benefit of NumPy. The same applies to min/max: np.min(a) is much faster than min(a) on a large array.
5. Forgetting keepdims and breaking broadcast
import numpy as np a = np.random.default_rng(0).standard_normal((100, 5)) # Wrong: the shapes do not match centred = a - a.mean(axis=1) # (100, 5) - (100,) fails because the shapes do not line up # Right: keep the collapsed axis with length 1 centred = a - a.mean(axis=1, keepdims=True) # (100, 5) - (100, 1) broadcasts correctly
Whenever you reduce an array and then combine the result with the original, you almost always want keepdims=True.
🎯 Your Turn: Rolling Mean in O(n)
Write moving_average(arr, window) that returns the rolling mean of arr with a fixed window size. The simple solution is O(n × window). The better one is O(n), using np.cumsum.
Expected behaviour:
moving_average(np.array([1, 2, 3, 4, 5, 6, 7]), window=3) # array([2., 3., 4., 5., 6.]) means of [1,2,3], [2,3,4], ..., [5,6,7] moving_average(np.array([1.0, 2.0, 4.0, 8.0]), window=2) # array([1.5, 3.0, 6.0])
setup added so this can run · defines moving_average, np
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) def moving_average(*_a, **_kw): print('-> moving_average() called') return _AutoMock('moving_average()') np = _AutoMock('np')
Constraints:
- Return a 1-D
ndarrayof lengthlen(arr) - window + 1. - Must run in O(n), in a single pass, using
np.cumsum, not a Python loop. - Raise
ValueErrorifwindow < 1orwindow > len(arr).
Skeleton:
import numpy as np def moving_average(arr, window): # TODO 1: validate window (>= 1, <= len(arr)) # TODO 2: compute cumulative sum, prepended with a 0 # TODO 3: subtract cumsum[:-window] from cumsum[window:] to get window sums # TODO 4: divide by window to get means ...
Hint 1: Cumsum gives you prefix sums
np.cumsum(arr) returns [arr[0], arr[0]+arr[1], arr[0]+arr[1]+arr[2], ...]. The sum of any window arr[i:i+w] equals cumsum[i+w-1] - cumsum[i-1]. Prepend a zero (np.concatenate([[0], cumsum])) so the i=0 case works without a special branch.
Hint 2: Vectorise the subtraction
With the prepended cumsum (lengthn+1), the array of window sums is cumsum[window:] - cumsum[:-window]. Both are length n - window + 1. Divide by window (a scalar, broadcasts automatically) to get means. No loop.
Show full solution
import numpy as np def moving_average(arr, window): """Rolling mean of `arr` over a fixed-size window. Runs in O(n) time using cumsum.""" if window < 1: raise ValueError("window must be >= 1") if window > len(arr): raise ValueError("window cannot exceed array length") cumsum = np.cumsum(arr, dtype=np.float64) # prefix sums cumsum = np.concatenate(([0.0], cumsum)) # put a 0 at the front so cumsum[0] == 0 window_sums = cumsum[window:] - cumsum[:-window] # the sum of each window return window_sums / window # Verify print(moving_average(np.array([1, 2, 3, 4, 5, 6, 7]), 3)) # [2. 3. 4. 5. 6.] print(moving_average(np.array([1.0, 2.0, 4.0, 8.0]), 2)) # [1.5 3. 6. ] # Speed: the naive loop compared with cumsum on a million elements import time arr = np.random.default_rng(0).standard_normal(1_000_000) w = 100 t0 = time.perf_counter() naive = np.array([arr[i:i+w].mean() for i in range(len(arr) - w + 1)]) t_naive = time.perf_counter() - t0 t0 = time.perf_counter() fast = moving_average(arr, w) t_fast = time.perf_counter() - t0 print(f"naive: {t_naive:.3f}s cumsum: {t_fast:.4f}s speedup: {t_naive/t_fast:.0f}x") # naive: ~1.2s cumsum: ~0.01s speedup: ~100x
Why this is O(n): cumsum is a single linear pass. The subtraction and division are also linear, and they run vectorised in C. There is no nested loop anywhere. Doubling the window size doesn't change the runtime.
Why dtype=np.float64 on cumsum: if arr has an integer dtype, the cumulative sums can overflow without warning. Converting to float64 first avoids that. For very long arrays of large floats, even float64 collects rounding error, because np.cumsum adds the numbers in the simplest way. Rolling means in production code use Kahan summation or the algorithm in pandas' rolling().mean(), which subtracts the leaving element rather than recomputing from a global prefix sum (better numerics, same O(n)).
This technique, prefix sums for sliding-window aggregates, works for many problems. Rolling sum, rolling mean, "is there a window of size k with a sum of at least X?" and "longest subarray with sum < S" can all be solved with operations on a cumsum array. It is worth practising until you know it well.
What You Learned
- An
ndarrayis a contiguous buffer + dtype + shape + strides. Reshaping and slicing change the metadata, not the data. - Use the right constructor (
zeros,empty,full,eye,linspace) and the moderndefault_rngfor random data. - Vectorise: write computations as whole-array operations, never per-element Python loops. Expect 50–200× speedups.
- Broadcasting aligns shapes from the right, and size-1 and missing dimensions stretch. Use
[:, None]andkeepdims=Trueto control alignment. - Advanced indexing: integer arrays and boolean masks both copy, but plain slices are views.
np.wherefor vectorised ternaries. axis=kcollapses axis k.axis=0gives per-column results,axis=1per-row.- Linear algebra: prefer
@overnp.dot,np.linalg.solve(A, b)overinv(A) @ b. - Views vs copies: slices are views, and fancy indexing copies. Changes to a view also change the original. Use
.copy()when you are not sure. - Watch for integer overflow without warning, accidental
objectdtype, and the in-place-on-shared-view trap. - For loops that can't be vectorised, use
numbaor Cython. For tensor contractions, usenp.einsum.
Next: Pandas DataFrames: Command Center: the same arrays, with labels, mixed column types and grouping.
Practice this
on practicepython.inShort exercises that run in your browser and tell you what your code did, as well as whether a test passed.
Found this useful? Share it with a friend who’s learning Python.