NumPy Vectorization

Real-World Scenario

A quant analyst calculates daily returns, rolling volatility, and position sizing across 500 stock tickers and 10 years of daily data — roughly 1.25 million data points. A naive Python loop takes 8 seconds per calculation. The vectorized version runs in under 50 milliseconds. Vectorization isn’t a micro-optimization; at data scale it’s the difference between interactive analysis and waiting.

The Core Idea: Eliminate Python Loops

Every Python for loop over array elements pays interpreter overhead: fetch the next object, unpack its type, dispatch the operation, re-box the result. At 1 million iterations, that overhead dominates. NumPy ufuncs (universal functions) are compiled C functions that operate on entire arrays — they pay that overhead exactly once.

import numpy as np
import 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 — slow: interpreter overhead × n
start = time.perf_counter()
result_loop = [a[i] + b[i] for i in range(n)]
loop_time = time.perf_counter() - start

# Vectorized — fast: one C function call
start = time.perf_counter()
result_vec = a + b
vec_time = time.perf_counter() - start

print(f"Loop:       {loop_time:.3f}s")
print(f"Vectorized: {vec_time:.4f}s")
print(f"Speedup:    {loop_time / vec_time:.0f}x")
# Loop:       0.281s
# Vectorized: 0.002s
# Speedup:    140x

Example 1: Vectorizing a Mathematical Formula

Computing compound interest for 10,000 customers — each with a different principal, rate, and term.

import numpy as np

rng = np.random.default_rng(42)
n = 10_000

# Each customer's loan parameters
principal = rng.uniform(1_000, 100_000, n)   # $1k–$100k
annual_rate = rng.uniform(0.03, 0.15, n)      # 3%–15%
years = rng.integers(1, 30, n).astype(float)  # 1–30 year terms

# Compound interest formula: A = P(1 + r)^t
# Applies to all 10,000 customers simultaneously — no loop
final_balance = principal * (1 + annual_rate) ** years

print(f"Min balance: ${final_balance.min():,.2f}")
print(f"Max balance: ${final_balance.max():,.2f}")
print(f"Total portfolio: ${final_balance.sum():,.2f}")

Example 2: Financial Returns Calculation

Vectorized daily and log returns — the bread and butter of quantitative finance.

import numpy as np

# Simulated price series for 5 assets over 252 trading days
rng = np.random.default_rng(42)
prices = np.cumprod(1 + rng.normal(0.0003, 0.015, (252, 5)), axis=0) * 100

# Daily simple returns: (P_t - P_{t-1}) / P_{t-1}
# np.diff computes P[1:] - P[:-1] across axis 0
returns = np.diff(prices, axis=0) / prices[:-1]  # shape (251, 5)

# Log returns: more numerically stable for long holding periods
log_returns = np.log(prices[1:] / prices[:-1])   # shape (251, 5)

# Annualized volatility per asset — std of daily returns × sqrt(252 trading days)
annualized_vol = returns.std(axis=0) * np.sqrt(252)
print("Annualized volatility per asset:")
for i, vol in enumerate(annualized_vol):
    print(f"  Asset {i}: {vol:.1%}")

# Cumulative return — product of (1 + r) for each asset
cumulative_return = np.prod(1 + returns, axis=0) - 1
print("\nCumulative returns:", cumulative_return.round(3))

Example 3: np.where for Conditional Logic

Vectorized if-else without a loop.

import numpy as np

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

# Single threshold
grades = np.where(scores >= 60, "pass", "fail")
print(grades)
# ['pass' 'fail' 'pass' 'pass' 'fail' 'pass' 'pass' 'fail' 'pass' 'pass']

# Multiple thresholds with np.select — vectorized if/elif/else
conditions = [
    scores >= 90,
    scores >= 80,
    scores >= 70,
    scores >= 60,
]
choices = ["A", "B", "C", "D"]
letter_grades = np.select(conditions, choices, default="F")
print(letter_grades)
# ['B' 'F' 'C' 'A' 'F' 'D' 'B' 'F' 'C' 'A']

Example 4: Rolling Calculations with Strides

Computing a rolling mean without a Python loop — using NumPy’s stride tricks to create a sliding window view.

import numpy as np

def rolling_mean(arr: np.ndarray, window: int) -> np.ndarray:
    """Vectorized rolling mean using stride_tricks."""
    shape = (arr.size - window + 1, window)
    strides = (arr.strides[0], arr.strides[0])
    # as_strided creates a 2-D view where each row is one window
    windows = np.lib.stride_tricks.as_strided(arr, shape=shape, strides=strides)
    return windows.mean(axis=1)

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

ma5 = rolling_mean(prices, window=5)
print(ma5.round(2))
# [101.6  103.   104.6  106.6  107.6  108.4]

np.apply_along_axis vs Pure Vectorization

np.apply_along_axis is cleaner than a Python loop over rows/columns, but still calls Python on each slice. Use true ufuncs when possible.

import numpy as np

data = np.random.default_rng(0).standard_normal((1000, 10))

# Slower: applies a Python function to each row
def normalize_row(row):
    return (row - row.mean()) / row.std()

normalized_loop = np.apply_along_axis(normalize_row, axis=1, arr=data)

# Faster: pure NumPy broadcasting — one C call for mean, one for std
mean = data.mean(axis=1, keepdims=True)   # (1000, 1)
std  = data.std(axis=1, keepdims=True)    # (1000, 1)
normalized_vec = (data - mean) / std      # broadcasts (1000, 10)

# Both produce the same result — the vectorized version is ~10x faster at scale
print(np.allclose(normalized_loop, normalized_vec))  # True

Common Mistakes

1. Using Python math functions on arrays:

import math
import numpy as np

arr = np.array([0.0, 1.0, 2.0, 3.0])

# Wrong — math.sqrt operates on scalars, raises TypeError on arrays
# result = math.sqrt(arr)

# Correct — np.sqrt is a ufunc: operates element-wise on the full array
result = np.sqrt(arr)  # [0. 1. 1.414 1.732]

2. Forgetting keepdims when broadcasting after a reduction:

import numpy as np
X = np.ones((5, 3))

# Without keepdims: mean shape is (3,) — broadcasts fine as a row
# but breaks for column normalization
row_mean = X.mean(axis=1)             # shape (5,) — cannot subtract from (5, 3)
row_mean_col = X.mean(axis=1, keepdims=True)  # shape (5, 1) — broadcasts correctly
result = X - row_mean_col             # (5, 3) - (5, 1) → (5, 3)