Optimising hot loops

Do less work, not the same work faster

Once the profiler has pointed at a hot loop, the biggest wins come from changing how much work it does, not from squeezing the same work. A loop that adds up a window of 250 prices at each of a million positions does 250 million additions; a version that slides the window does two million. That difference, the algorithmic complexity, dwarfs anything a clever trick on the inner line can give. So look first at what the loop repeats: a lookup in a list, a sum recomputed from scratch, a value that never changes but is rebuilt each time.

# Membership in a list scans it; in a set it is a hash lookup. Same answer, very different work
import time

ids = list(range(20_000))
as_set = set(ids)
wanted = list(range(19_000, 20_000))

start = time.perf_counter()
found_list = sum(1 for x in wanted if x in ids)             # scans up to 20,000 items for each of 1,000 lookups
slow = time.perf_counter() - start

start = time.perf_counter()
found_set = sum(1 for x in wanted if x in as_set)
fast = time.perf_counter() - start

print(found_list, found_set)                                # 1000 1000
print(slow > 20 * fast)                                     # True  the set is at least twenty times faster here

Slide, look up once, hoist

Three patterns cover most hot loops. Slide instead of recompute: a moving sum adds the value entering the window and subtracts the one leaving, so each position costs two operations whatever the window length. Build the lookup once: turning a list of pairs into a dict costs one pass, after which each lookup is instant, instead of scanning the list for every query. Hoist what does not change: anything computed inside a loop that depends only on values outside it (a constant, len(x), an attribute lookup) belongs before the loop. Test with a counter of reads or operations, as this lab does: a count does not depend on the machine's speed, so it is a stable measure of work.

# A sliding sum does two operations per position; recomputing each window does `window` of them
def moving_sum_slow(values, window):
    operations = 0
    sums = []
    for i in range(len(values) - window + 1):
        total = 0
        for v in values[i : i + window]:
            total += v
            operations += 1
        sums.append(total)
    return sums, operations

def moving_sum_fast(values, window):
    operations = window
    total = sum(values[:window])
    sums = [total]
    for i in range(window, len(values)):
        total += values[i] - values[i - window]
        operations += 2
        sums.append(total)
    return sums, operations

values = list(range(1_000))
slow, slow_ops = moving_sum_slow(values, 250)
fast, fast_ops = moving_sum_fast(values, 250)
print(slow == fast)                                   # True  the same answers
print(slow_ops, fast_ops)                             # 187750 1750  about a hundred times less work

# Hoisting: a value that never changes is computed once, outside the loop
import math

calls = {"count": 0}

def annualisation_factor():                          # imagine something costly that never changes
    calls["count"] += 1
    return math.sqrt(252)

def scale_slow(returns):
    return [r * annualisation_factor() for r in returns]

def scale_fast(returns):
    factor = annualisation_factor()                  # hoisted out of the loop
    return [r * factor for r in returns]

returns = [0.001 * i for i in range(2_000)]
slow = scale_slow(returns)
slow_calls, calls["count"] = calls["count"], 0
fast = scale_fast(returns)
print(slow == fast)                                  # True  the same numbers
print(slow_calls, calls["count"])                    # 2000 1   one call instead of two thousand

Let NumPy loop for you

When the work really is arithmetic on many numbers, move the loop out of Python. A NumPy array stores the numbers in one block of memory and applies an operation to all of them in compiled code, so np.log(prices[1:] / prices[:-1]) replaces a Python loop with a single call that is typically tens of times faster. This is vectorisation, and it is the first optimisation to reach for in numeric code. It works best on whole arrays with no if per element; where a condition is needed use np.where or a boolean mask. Keep an eye on memory too: vectorising makes temporary arrays the size of the data.

# Vectorising: one NumPy call replaces a Python loop and gives the same numbers
import math
import time
import numpy as np

rng = np.random.default_rng(3)
prices = 100 * np.exp(np.cumsum(rng.normal(0, 0.001, size=300_000)))

start = time.perf_counter()
loop = [math.log(prices[i] / prices[i - 1]) for i in range(1, len(prices))]
loop_time = time.perf_counter() - start

start = time.perf_counter()
vector = np.log(prices[1:] / prices[:-1])
vector_time = time.perf_counter() - start

print(np.allclose(loop, vector))                    # True  the same returns
print(loop_time > 5 * vector_time)                  # True  at least five times faster here

• Cut the amount of work first: slide a window instead of re-adding it, and avoid recomputing what you already have.
• Use the right structure: a set or dict for lookups (built once), not a list scanned for every query; hoist constants out of loops.
• Vectorise numeric loops with NumPy, count work rather than seconds in tests, and keep checking the fast version against a slow, obviously correct one.

Next up: Numba, compiling Python numeric loops to machine code with a decorator.

Lab · what's graded

• Slide the window in moving_sum25 pts
• Use a set in has_duplicates25 pts
• Build a dict once in lookup_prices25 pts
• Vectorise log_returns25 pts