Otimização de laços críticos

Fazer menos trabalho, não o mesmo trabalho mais rápido

Quando o profiler aponta um laço crítico, os maiores ganhos vêm de mudar quanto trabalho ele faz, não de espremer o mesmo trabalho. Um laço que soma uma janela de 250 preços em cada uma de um milhão de posições faz 250 milhões de somas; uma versão que desliza a janela faz dois milhões. Essa diferença, a complexidade algorítmica, ofusca qualquer truque na linha interna. Olhe primeiro o que o laço repete: uma busca em lista, uma soma recalculada do zero, um valor que nunca muda mas é reconstruído a cada vez.

# 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

Deslizar, procurar uma vez, tirar do laço

Três padrões cobrem a maioria dos laços críticos. Deslizar em vez de recalcular: uma soma móvel adiciona o valor que entra na janela e subtrai o que sai, então cada posição custa duas operações, seja qual for o tamanho da janela. Construir a busca uma vez: transformar uma lista de pares em dict custa uma passada, e depois cada busca é instantânea, em vez de varrer a lista a cada consulta. Tirar do laço o que não muda: tudo que é calculado dentro de um laço e depende só de valores de fora (uma constante, len(x), uma busca de atributo) vai antes do laço. Teste com um contador de leituras ou operações, como este laboratório: uma contagem não depende da velocidade da máquina, logo é uma medida estável do trabalho.

# 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

Deixe o NumPy fazer o laço por você

Quando o trabalho é de fato aritmética sobre muitos números, tire o laço do Python. Um array do NumPy guarda os números em um único bloco de memória e aplica uma operação a todos em código compilado, então np.log(prices[1:] / prices[:-1]) substitui um laço Python por uma única chamada tipicamente dezenas de vezes mais rápida. Isso é vetorização, e é a primeira otimização a tentar em código numérico. Funciona melhor com arrays inteiros sem um if por elemento; onde há condição, use np.where ou uma máscara booleana. Fique de olho na memória também: vetorizar cria arrays temporários do tamanho dos dados.

# 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

• Reduza primeiro a quantidade de trabalho: deslize a janela em vez de somá-la de novo e evite recalcular o que já tem.
• Use a estrutura certa: um conjunto ou dict para buscas (construído uma vez), não uma lista varrida a cada consulta; tire as constantes dos laços.
• Vetorize laços numéricos com NumPy, conte o trabalho e não os segundos nos testes, e continue comparando a versão rápida com uma lenta e obviamente correta.

A seguir: Numba, compilar laços numéricos de Python em código de máquina com um decorador.

Lab: o que é avaliado

• Deslizar a janela em moving_sum25 pts
• Usar um conjunto em has_duplicates25 pts
• Construir um dict uma só vez em lookup_prices25 pts
• Vetorizar log_returns25 pts