NumPy Vectorization

A Loop vs One Array Expression

The same elementwise addition, written both ways:

import timeit
import numpy as np

n = 1_000_000
a = np.random.default_rng(0).standard_normal(n)
b = np.random.default_rng(1).standard_normal(n)

t_loop = min(timeit.repeat(lambda: [a[i] + b[i] for i in range(n)], number=1, repeat=5))
t_vec = min(timeit.repeat(lambda: a + b, number=1, repeat=5))

print(f"loop       {t_loop:.4f} s")
print(f"vectorized {t_vec:.4f} s")
print(f"speedup    {t_loop / t_vec:.0f}x")

$ python loop_vs_vec.py
loop       0.1611 s
vectorized 0.0037 s
speedup    44x

The loop pays interpreter overhead a million times, and each a[i] also boxes a NumPy scalar. a + b makes one call into a compiled loop. Both times are the best of five runs; the exact ratio moves between runs and machines, but the loop stays one to two orders of magnitude behind.

Vectorizing a Formula

Compound interest for 10,000 customers, each with their own principal, rate and term:

import numpy as np

rng = np.random.default_rng(42)
principal = rng.uniform(1_000, 100_000, 10_000)
rate = rng.uniform(0.03, 0.15, 10_000)
years = rng.integers(1, 31, 10_000)

balance = principal * (1 + rate) ** years

print(balance.shape)
print(f"min ${balance.min():,.2f}  max ${balance.max():,.2f}")
print(f"total ${balance.sum():,.0f}")

$ python compound.py
(10000,)
min $1,423.49  max $5,645,724.69
total $3,200,341,757

The formula is written exactly as on paper, and every operator applies to all 10,000 customers at once. The huge maximum comes from customers who drew both a high rate and a 30-year term — compounding is exponential in years.

Returns Along an Axis

import numpy as np

rng = np.random.default_rng(42)
prices = 100 * np.cumprod(1 + rng.normal(0.0003, 0.015, (252, 3)), axis=0)

returns = np.diff(prices, axis=0) / prices[:-1]
log_returns = np.log(prices[1:] / prices[:-1])
vol = returns.std(axis=0) * np.sqrt(252)

print(prices.shape, returns.shape)
print("annual vol:", vol.round(3))
print("total return:", (prices[-1] / prices[0] - 1).round(3))
print("sum of log returns:", log_returns.sum(axis=0).round(3))

$ python returns.py
(252, 3) (251, 3)
annual vol: [0.244 0.233 0.228]
total return: [-0.076 -0.123 -0.097]
sum of log returns: [-0.079 -0.131 -0.101]

np.diff(..., axis=0) subtracts consecutive rows, so 252 prices give 251 returns per asset. The simulated daily std of 0.015 annualizes to about 0.015 × √252 ≈ 0.24. Log returns add up over time: their sum is log(last / first), which is why they are preferred for multi-period math.

Conditional Logic with np.where and np.select

import numpy as np

scores = np.array([88, 45, 72, 91, 55, 67, 83, 39])

print(np.where(scores >= 60, "pass", "fail"))

conditions = [scores >= 90, scores >= 80, scores >= 70, scores >= 60]
print(np.select(conditions, ["A", "B", "C", "D"], default="F"))

$ python where_select.py
['pass' 'fail' 'pass' 'pass' 'fail' 'pass' 'pass' 'fail']
['B' 'F' 'C' 'A' 'F' 'D' 'B' 'F']

np.where is a vectorized if/else. np.select is if/elif/else: for each element it takes the choice for the first condition that is true, so 91 gets A even though it also satisfies the three later conditions.

Rolling Mean with a Sliding Window View

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

prices = np.array([100., 102., 98., 105., 103., 107., 110., 108.])

windows = sliding_window_view(prices, 3)
print(windows)
print(windows.mean(axis=1).round(2))
print(np.shares_memory(windows, prices))

$ python rolling.py
[[100. 102.  98.]
 [102.  98. 105.]
 [ 98. 105. 103.]
 [105. 103. 107.]
 [103. 107. 110.]
 [107. 110. 108.]]
[100.   101.67 102.   105.   106.67 108.33]
True

Each row of windows is three consecutive prices, and mean(axis=1) averages each row, giving six rolling means from eight prices. shares_memory is True: the windows are a view that reuses the original buffer instead of copying it, which keeps this cheap even for long series.

np.vectorize Is Not Compiled

import math
import timeit
import numpy as np

x = np.random.default_rng(0).uniform(0, 10, 500_000)
slow_sigmoid = np.vectorize(lambda v: 1 / (1 + math.exp(-v)))
fast_sigmoid = lambda: 1 / (1 + np.exp(-x))

t_vec = min(timeit.repeat(lambda: slow_sigmoid(x), number=1, repeat=5))
t_ufunc = min(timeit.repeat(fast_sigmoid, number=1, repeat=5))

print(f"np.vectorize  {t_vec:.4f} s")
print(f"ufunc math    {t_ufunc:.4f} s")
print(np.allclose(slow_sigmoid(x), fast_sigmoid()))

$ python vectorize_speed.py
np.vectorize  0.0800 s
ufunc math    0.0045 s
True

Both give the same numbers, but np.vectorize called the lambda half a million times — it is a loop with broadcasting added, not a compiler. Rewriting the body with np.exp is what makes it fast.

math Functions Don’t Take Arrays

import math
import numpy as np

arr = np.array([0.0, 1.0, 4.0])
try:
    print(math.sqrt(arr))
except TypeError as e:
    print("TypeError:", e)
print(np.sqrt(arr))

$ python math_vs_np.py
TypeError: only 0-dimensional arrays can be converted to Python scalars
[0. 1. 2.]

math.sqrt accepts a single number and cannot convert a three-element array to one. The NumPy version is a ufunc and handles every element.

Broadcasting After a Reduction Needs keepdims

import numpy as np

X = np.arange(15, dtype=float).reshape(5, 3)
try:
    X - X.mean(axis=1)
except ValueError as e:
    print("ValueError:", e)

centered = X - X.mean(axis=1, keepdims=True)
print(centered)

$ python keepdims.py
ValueError: operands could not be broadcast together with shapes (5,3) (5,) 
[[-1.  0.  1.]
 [-1.  0.  1.]
 [-1.  0.  1.]
 [-1.  0.  1.]
 [-1.  0.  1.]]

Row means have shape (5,), and broadcasting lines shapes up from the right, so (5, 3) meets (5,) and the 3 and 5 clash. keepdims=True keeps the result as (5, 1), which stretches across the three columns and centers each row.

Practice

1. Vectorize this loop: out = [x * 2 if x > 0 else 0 for x in arr] for arr = np.array([-2, 3, -1, 5]).
2. Compute a 3-point rolling maximum of [3, 8, 1, 2, 4, 7, 2].
3. Standardize each row of np.arange(12.).reshape(3, 4) to mean 0.

import numpy as np

arr = np.array([-2, 3, -1, 5])
print(np.where(arr > 0, arr * 2, 0))

$ python practice_where.py
[ 0  6  0 10]

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

a = np.array([3, 8, 1, 2, 4, 7, 2])
print(sliding_window_view(a, 3).max(axis=1))

$ python practice_rollmax.py
[8 8 4 7 7]

import numpy as np

X = np.arange(12.0).reshape(3, 4)
print(X - X.mean(axis=1, keepdims=True))

$ python practice_rowcenter.py
[[-1.5 -0.5  0.5  1.5]
 [-1.5 -0.5  0.5  1.5]
 [-1.5 -0.5  0.5  1.5]]