Vetorização

O laço e a operação sobre o array

Em Python puro, para dobrar todos os elementos você percorre um por um. No NumPy, você escreve a operação sobre o array, e ela vale para todos os elementos:

import time

import numpy as np

valores = np.array([1, 2, 3, 4, 5])
resultado = []
for x in valores:
    resultado.append(x * 2)
print(type(resultado).__name__, type(valores * 2).__name__)
print(valores * 2)

list ndarray
[ 2  4  6  8 10]

Repare que o laço devolveu uma lista (de escalares NumPy), e não um array. Quem escreve for x in array: lista.append(...) abandona as vantagens do NumPy: perde o tipo, perde a velocidade e ainda precisa converter a lista de volta com np.array(resultado).

Fórmulas inteiras, sem laço

As fórmulas se escrevem como na matemática, e valem para um elemento ou um milhão:

celsius = np.array([0.0, 25.0, 100.0])
print(celsius * 9 / 5 + 32)

precos = np.array([100.0, 250.0, 80.0])
quantidades = np.array([3, 1, 10])
print((precos * quantidades).sum())

[ 32.  77. 212.]
1350.0

A segunda conta é o total de um pedido: preço vezes quantidade, elemento a elemento, e depois a soma. Duas linhas, nenhum laço.

Por que é mais rápido

Um laço do Python paga um pequeno imposto a cada volta: verificar o tipo, gerenciar o laço, voltar ao interpretador. Para 5 elementos ele é invisível, e para milhões ele domina o tempo. Na versão vetorizada, o laço existe, mas roda em código compilado, sem esse imposto. A diferença é fácil de medir:

def medir(funcao, repeticoes=3):
    melhor = float("inf")
    for _ in range(repeticoes):
        inicio = time.perf_counter()
        funcao()
        melhor = min(melhor, time.perf_counter() - inicio)
    return melhor


dados = np.arange(200_000, dtype=np.float64)


def com_laco(a):
    saida = np.empty_like(a)
    for i in range(a.size):
        saida[i] = a[i] ** 2 + 1
    return saida


def vetorizado(a):
    return a ** 2 + 1


print(np.array_equal(com_laco(dados), vetorizado(dados)))
mais_rapido = medir(lambda: com_laco(dados)) > 50 * medir(lambda: vetorizado(dados))
print("vetorizado pelo menos 50 vezes mais rápido:", mais_rapido)

True
vetorizado pelo menos 50 vezes mais rápido: True

O resultado é idêntico, e a versão vetorizada é, nesse caso, muito mais rápida (com laço sobre um array, é pior do que sobre uma lista, porque cada a[i] cria um objeto escalar do NumPy).

O erro de usar funções que não são do NumPy

As funções do módulo math esperam um número, e não um array. O NumPy tem uma versão vetorizada de cada uma:

import math

try:
    math.sqrt(np.array([4.0, 9.0]))
except TypeError as erro:
    print("math.sqrt não aceita arrays:", type(erro).__name__)
print(np.sqrt(np.array([4.0, 9.0])))

math.sqrt não aceita arrays: TypeError
[2. 3.]

Quando o laço é a resposta certa

Nem tudo se vetoriza. Quando cada passo depende do resultado do anterior (um saldo que rende juros e recebe aportes, uma simulação, um algoritmo iterativo), o laço continua sendo a forma natural. O que vale é a pergunta: "existe uma versão vetorizada?":

juros = 0.01
saldo = 1000.0
saldos = []
for _ in range(3):
    saldo = saldo * (1 + juros) + 100
    saldos.append(round(saldo, 2))
print(saldos)

[1110.0, 1221.1, 1333.31]

A regra prática

Ao escrever um for sobre elementos de um array, eu paro e procuro a forma vetorizada primeiro: operadores aritméticos, comparações, np.where, agregações com axis, funções universais. Quando elas não resolvem, é que o problema tem dependência entre os passos, e aí o laço é legítimo (ou é hora de olhar cumsum, accumulate e o capítulo 34).