Thinking in arrays with NumPy

Quant code rarely works on one number at a time. It works on a table of returns, days down and assets across, and asks for every asset's volatility, every day's cross-sectional rank, every scenario's price. NumPy lets you write those operations on the whole table at once. This lesson covers the three ideas that make that work, shapes and axes, vectorisation, and broadcasting, and the silent bugs that come with them.

Shapes and axes

Every NumPy array has a shape. A year of daily returns for 50 assets might be (252, 50): axis 0 is time, axis 1 is the asset. Reductions take an axis argument, and the rule to remember is that the axis you name is the one that disappears:

• returns.mean(axis=0) averages down the rows, one value per asset: shape (50,).
• returns.mean(axis=1) averages across the columns, one value per day: shape (252,).

If you keep one convention for every array in a project (time on axis 0, say), most axis arguments become automatic.

Why vectorise

Here is the same calculation twice: standardise every return by its own asset's mean and volatility, first with Python loops and then with whole-array operations.

import time

import numpy as np

rng = np.random.default_rng(0)
returns = rng.normal(0.0005, 0.01, size=(2520, 50))  # 10 years of daily returns, 50 assets

# Loop version: z-score every return by its asset's mean and volatility.
start = time.perf_counter()
z_loop = np.empty_like(returns)
for j in range(returns.shape[1]):
    column = returns[:, j]
    mean, sd = column.mean(), column.std()
    for i in range(returns.shape[0]):
        z_loop[i, j] = (returns[i, j] - mean) / sd
loop_seconds = time.perf_counter() - start

# Vectorised version: the (50,) means and vols broadcast down all 2,520 rows.
start = time.perf_counter()
z = (returns - returns.mean(axis=0)) / returns.std(axis=0)
vector_seconds = time.perf_counter() - start

print(f"same answer: {np.allclose(z, z_loop)}")
print(f"loop {loop_seconds * 1000:.0f} ms, vectorised {vector_seconds * 1000:.2f} ms")

# Demean each day (each row) across assets: keepdims keeps a (2520, 1) column.
by_day = returns - returns.mean(axis=1, keepdims=True)
print(f"row means after demeaning: {np.abs(by_day.mean(axis=1)).max():.1e}")

On the machine this lesson was written on, it prints same answer: True, then about 50 ms for the loops against 0.8 ms vectorised: roughly 60 times faster, for 126,000 numbers. Your timings will differ; the ratio is what matters.

The speed doesn't come from cleverer maths. The loop runs in the Python interpreter, which checks types and creates an object for every single number. The vectorised line runs one compiled loop over a contiguous block of memory. The gap grows with the data, and on a real research dataset it is the difference between seconds and hours.

Key idea. Move loops out of Python and into NumPy. If you are indexing single elements inside a for loop over a large array, there is almost always a whole-array way to write it.

Broadcasting

In the vectorised line above, returns has shape (2520, 50) and returns.mean(axis=0) has shape (50,), yet NumPy subtracts one from the other. The rule that allows it is broadcasting:

1. Line the two shapes up on the right.
2. Compare sizes axis by axis. Each pair must be equal, or one of them must be 1 (a missing axis counts as 1).
3. Size-1 axes are stretched to match. Nothing is copied in memory; NumPy just reuses the same values.

Step through the cases below. The dashed cells are the stretched copies.

Demean each asset: (4, 3) and (3,) line up on the right as 3 and 3; the missing axis is stretched to 4. Demean each day: the day means need to line up with axis 0, so keepdims=True keeps them as a (4, 1) column. Forget it, and the (4,) array is lined up with the last axis, where 3 meets 4 and NumPy raises an error. That error is the good outcome.

Key idea. Broadcasting aligns shapes from the right. To combine something with the rows of a 2-D array, it must have shape (n, 1): use keepdims=True or [:, None].

The silent bug

The last case is the dangerous one. With a square array, 3 days by 3 assets, subtracting prices.mean(axis=1) doesn't fail: the 3 day means line up with the 3 assets, and each column gets a different day's mean subtracted. The code runs and the numbers look plausible. They are wrong.

Real data is rarely square, but research code often is tested on small examples that are, and covariance matrices and many intermediate arrays always are. Three habits catch this class of bug:

• Use keepdims=True whenever a reduction will be combined with the array it came from.
• Check shapes in code that matters: assert z.shape == returns.shape.
• Test with deliberately non-square inputs, like 5 days by 3 assets, so axis mix-ups fail loudly.

Views, copies and masks

Two more behaviours are worth knowing early, because both cause bugs that look like bad data.

Slices are views. prices[1:3] doesn't copy anything. It is a window onto the same memory, so writing to it changes prices. Call .copy() when you want an independent array. (Indexing with a list or a boolean mask, by contrast, always copies.)

Masks and cumulative operations replace most branching loops. A drawdown, the fall from the running peak, is a classic example:

import numpy as np

prices = np.array([100.0, 104.0, 101.0, 107.0, 98.0, 103.0])

# Running peak and drawdown, with no Python loop.
peak = np.maximum.accumulate(prices)
drawdown = prices / peak - 1
print(peak)
print(np.round(drawdown, 3))
print(f"max drawdown {drawdown.min():.1%}")

# Slices are views: changing one changes the original.
window = prices[1:3]
window[0] = 0.0
print(prices[:3])

It prints the running peak [100. 104. 104. 107. 107. 107.], drawdowns of 0, 0, −2.9%, 0, −8.4% and −3.7%, and a maximum drawdown of −8.4%. The last line prints [100. 0. 101.]: setting window[0] overwrote the original price of 104.

Where this shows up in quant work

• Signal research. Cross-sectional ranks, z-scores and demeaning across assets each day are row operations on a days × assets array: keepdims territory.
• Risk. Portfolio variance is w @ cov @ w, and a covariance matrix is square, the exact case where axis mistakes go unnoticed.
• Scenario analysis. Broadcasting a column of spot prices against a row of shocks builds a whole grid of scenarios, priced in one vectorised call.

Exercises

• Given returns of shape (252, 50) and portfolio weights w of shape (50,), compute the daily portfolio returns without a loop. What shape is the result?
• Rank each day's returns across assets (1 = best) using argsort twice along the right axis. Check your answer on a (5, 3) example by hand.
• Write a function that demeans each row, then test it on a (3, 3) input where the wrong-axis version would also run. Which assertion catches the bug?

Key takeaways

• The axis you pass to a reduction is the one that disappears; keep one axis convention per project.
• Vectorised NumPy is typically tens to hundreds of times faster than Python loops.
• Broadcasting aligns shapes from the right; sizes must match or be 1, and size-1 axes stretch.
• Square arrays hide axis mistakes. Use keepdims=True, assert shapes, and test on non-square data.