NumPy: vectorization and linear algebra

Murad is processing one million sales of an online shop: for every sale he must compute the amount, apply a discount and find the daily change. With a for loop this takes seconds; with NumPy it takes milliseconds. The secret is vectorization: we write the operation once for the whole array instead of once per element. In this lesson we will also work with linear algebra, the mathematical language of data science.

Vectorization: computing without loops

Why is a Python loop slow? On every iteration the interpreter reads an instruction, checks the element's type and creates a new float object for the result. NumPy checks the type once for numbers of the same type lying next to each other in memory and uses the processor's vector (SIMD) instructions. Run the code and see the difference for yourself:

import time
import numpy as np

x = np.random.default_rng(0).random(1_000_000)

t0 = time.perf_counter()
total = 0.0
for v in x:
    total += v * v
t_loop = time.perf_counter() - t0

t0 = time.perf_counter()
total_vec = np.dot(x, x)
t_vec = time.perf_counter() - t0

print(f'loop:       {t_loop:.4f} s')
print(f'vectorized: {t_vec:.4f} s')
print('same result:', bool(np.isclose(total, total_vec)))
print(f'speed-up: about {t_loop / t_vec:.0f} times')

import numpy as np

prices = np.array([12.0, 45.5, 8.9, 120.0, 64.0])
qty = np.array([3, 1, 10, 1, 2])

revenue = prices * qty
discount = np.where(revenue > 100, 0.10, 0.0)
to_pay = (revenue * (1 - discount)).round(2)
print(revenue)
print(to_pay, to_pay.sum())

daily = np.array([520.0, 610.0, 580.0, 700.0, 655.0])
print(np.diff(daily))
print((np.diff(daily) / daily[:-1] * 100).round(1))
print(np.cumsum(daily))

[ 36.   45.5  89.  120.  128. ]
[ 36.   45.5  89.  108.  115.2] 393.7
[ 90. -30. 120. -45.]
[17.3 -4.9 20.7 -6.4]
[ 520. 1130. 1710. 2410. 3065.]

Broadcasting: arrays of different shapes

When you write a * 10, the number 10 is as if “stretched” to an array of the same shape as a. This mechanism is called broadcasting, and it works not only with single numbers but also with arrays of different shapes. NumPy compares the shapes from right to left:

1. If one shape has fewer dimensions, 1s are added on its left: (4,) → (1, 4).
2. Along each axis the two lengths must be equal, or one of them must be 1.
3. An axis of length 1 is “stretched” to the other array's length (without copying memory); the result's shape is the maximum along each axis.
4. If along some axis the lengths differ and neither is 1, you get a ValueError.

A classic use of broadcasting is standardization: from each column we subtract its mean and divide by its standard deviation, so that features on different scales (height in cm, weight in kg) become comparable. A mean of shape (2,) is subtracted from an array of shape (4, 2) — rules 1 and 3 do the work.

import numpy as np

X = np.array([[170.0, 65.0],
              [160.0, 52.0],
              [182.0, 80.0],
              [175.0, 71.0]])
mean = X.mean(axis=0)
std = X.std(axis=0)
Z = (X - mean) / std
print(mean, std.round(3))
print(Z.round(2))
print(np.allclose(Z.mean(axis=0), 0), Z.std(axis=0).round(6))

[171.75  67.  ] [ 8.012 10.173]
[[-0.22 -0.2 ]
 [-1.47 -1.47]
 [ 1.28  1.28]
 [ 0.41  0.39]]
True [1. 1.]

Combining a column vector with a row vector gives an “each with each” table. a[:, np.newaxis] (or a[:, None]) turns an array of shape (3,) into a (3, 1) column:

import numpy as np

a = np.arange(1, 4)
b = np.arange(1, 6)
table = a[:, np.newaxis] * b
print(table.shape)
print(table)

M = np.ones((3, 4))
v = np.array([10.0, 20.0, 30.0])
try:
    M + v
except ValueError as e:
    print('Error:', str(e).strip())
print((M + v[:, None])[:, 0])

(3, 5)
[[ 1  2  3  4  5]
 [ 2  4  6  8 10]
 [ 3  6  9 12 15]]
Error: operands could not be broadcast together with shapes (3,4) (3,)
[11. 21. 31.]

Reshaping

reshape shows the same elements in a different shape; the number of elements must not change, and one dimension may be -1 — NumPy works it out. .T transposes, ravel() flattens to one dimension (returning a view when possible), and np.vstack, np.hstack and np.concatenate glue arrays together.

import numpy as np

a = np.arange(12)
print(a.reshape(3, 4))
print(a.reshape(2, -1).shape)
m = a.reshape(3, 4)
print(m.T.shape, m.ravel()[:5])
print(np.vstack([m[0], m[1]]))
print(np.concatenate([m, m], axis=1).shape)

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
(2, 6)
(4, 3) [0 1 2 3 4]
[[0 1 2 3]
 [4 5 6 7]]
(3, 8)

Linear algebra: np.linalg

The product of matrices is not element-wise multiplication. A * B multiplies matching elements, while A @ B (or np.dot) computes the mathematical matrix product: a row of A is multiplied by a column of B as a dot product. For this, the number of columns of A must equal the number of rows of B.

import numpy as np

A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])
print(A * B)
print(A @ B)
print(np.dot(A, B).shape, (np.ones((2, 3)) @ np.ones((3, 5))).shape)

[[ 5 12]
 [21 32]]
[[19 22]
 [43 50]]
(2, 2) (2, 5)

A system of linear equations is written in matrix form as Ax = b. If det A ≠ 0, the system has a unique solution and A has an inverse: x = A⁻¹b. For a 2 × 2 matrix everything can be computed by hand:

With three unknowns (notebook, pen, ruler) the hand calculation gets long, but NumPy solves it in one line:

import numpy as np

A = np.array([[2.0, 1.0, 3.0],
              [1.0, 3.0, 1.0],
              [4.0, 2.0, 0.0]])
b = np.array([9.5, 5.0, 7.0])

x = np.linalg.solve(A, b)
print(x.round(2))
print(np.allclose(A @ x, b))
print(round(np.linalg.det(A), 2))
x_inv = np.linalg.inv(A) @ b
print(np.allclose(x, x_inv))

[1.5 0.5 2. ]
True
-30.0
True

An eigenvector is a non-zero vector v that the matrix A only stretches or shrinks without changing its direction; the stretch factor λ is the eigenvalue. We find the λ values from the characteristic equation. Applications: principal component analysis (PCA), Google's PageRank algorithm, the natural vibration frequencies of bridges and buildings, and the stability of dynamic systems.

import numpy as np

A = np.array([[2.0, 1.0], [1.0, 2.0]])
values, vectors = np.linalg.eigh(A)
print(values.round(4))
v = vectors[:, 1]
print(np.allclose(A @ v, values[1] * v))
print(round(np.linalg.det(A), 4), round(values.prod(), 4))
print(round(np.trace(A), 4), round(values.sum(), 4))

[1. 3.]
True
3.0 3.0
4.0 4.0

Linear algebra leads straight to machine learning. Fitting the best straight line through points (the least-squares method) leads to the normal equation. If we add a column of ones to the matrix X, the intercept is found automatically as well:

import numpy as np

rng = np.random.default_rng(7)
x = rng.uniform(0, 10, 50)
y = 2.5 * x + 1.0 + rng.normal(0, 1, 50)
X = np.column_stack([np.ones_like(x), x])
w = np.linalg.solve(X.T @ X, X.T @ y)
print(w.round(3))
w2, *_ = np.linalg.lstsq(X, y, rcond=None)
print(w2.round(3))

[1.236 2.458]
[1.236 2.458]

import numpy as np

X = np.array([[2.0, 50.0, 1.0],
              [4.0, 80.0, 3.0],
              [6.0, 65.0, 2.0],
              [8.0, 95.0, 5.0]])
# min-max scale every column

[[0.   0.   0.  ]
 [0.33 0.67 0.5 ]
 [0.67 0.33 0.25]
 [1.   1.   1.  ]]

import numpy as np

# A: kilograms in each purchase, b: amounts paid (manat)
A = np.array([[2.0, 1.0, 1.0],
              [1.0, 2.0, 1.0],
              [1.0, 1.0, 3.0]])
b = np.array([9.5, 10.0, 13.5])
# solve A x = b and print the prices, then the determinant

[2.  2.5 3. ]
7.0