Broadcasting

A ideia: esticar, sem copiar

Pense no array menor sendo esticado para combinar com o maior. O NumPy não cria cópias dos dados: ele reaproveita a memória como se ela estivesse repetida. Uma linha somada a uma grade vale para todas as linhas, e uma coluna, para todas as colunas:

import numpy as np

grade = np.ones((3, 4))
linha = np.array([1, 2, 3, 4])
print(grade + linha)

coluna = np.array([[10], [20], [30]])
print(grade + coluna)
print((np.arange(3).reshape(3, 1) + np.arange(4)).shape)

[[2. 3. 4. 5.]
 [2. 3. 4. 5.]
 [2. 3. 4. 5.]]
[[11. 11. 11. 11.]
 [21. 21. 21. 21.]
 [31. 31. 31. 31.]]
(3, 4)

O último exemplo é o mais poderoso: uma coluna (3, 1) somada a uma linha (4,) produz uma grade (3, 4), com todas as combinações.

A regra

Compare as formas da direita para a esquerda, uma dimensão de cada vez. Duas dimensões são compatíveis se forem iguais ou se uma delas for 1. Se todas as dimensões passam nesse teste, as formas combinam. Faltando dimensões à esquerda, o NumPy as trata como 1. Dá para consultar a regra por código:

def compativel(a, b):
    try:
        return np.broadcast_shapes(a, b)
    except ValueError:
        return None


print(compativel((3, 4), (4,)))
print(compativel((3, 1), (1, 4)))
print(compativel((5, 1, 3), (3,)))
print(compativel((3, 4), (3,)))

(3, 4)
(3, 4)
(5, 1, 3)
None

O último caso falha: (3, 4) e (3,) alinham o 4 com o 3 (da direita para a esquerda), que não são iguais e nenhum é 1. O erro de verdade tem esta mensagem:

try:
    np.ones((3, 4)) + np.ones(3)
except ValueError as erro:
    print(erro)

operands could not be broadcast together with shapes (3,4) (3,) 

Usos reais

Centralizar cada coluna (subtrair a média de cada uma) e normalizar cada linha (dividir pela soma da linha) são as duas operações mais comuns, e as duas dependem do broadcasting. O keepdims=True mantém o eixo reduzido com tamanho 1, que é exatamente o que o broadcasting precisa:

notas = np.array([[7.0, 8.0, 9.0], [4.0, 6.0, 8.0]])
print(notas - notas.mean(axis=0))
print(notas / notas.sum(axis=1, keepdims=True))

[[ 1.5  1.   0.5]
 [-1.5 -1.  -0.5]]
[[0.29166667 0.33333333 0.375     ]
 [0.22222222 0.33333333 0.44444444]]

A subtração usa uma média por coluna (forma (3,)). A divisão precisa de uma soma por linha em forma de coluna, (2, 1), e é por isso que o keepdims=True é necessário.

Outro uso é calcular todas as diferenças entre pares de uma vez, sem laço duplo:

indices = np.arange(4)
print(indices[:, None] - indices)

[[ 0 -1 -2 -3]
 [ 1  0 -1 -2]
 [ 2  1  0 -1]
 [ 3  2  1  0]]

A memória não é copiada

O broadcasting é um esticamento virtual. Dá para ver isso nos strides: um passo de 0 bytes significa "repita o mesmo dado". E o resultado é somente leitura, porque escrever em elementos que, na verdade, são o mesmo seria ambíguo:

esticado = np.broadcast_to(np.array([1, 2, 3]), (1000, 3))
print(esticado.shape, esticado.strides, esticado.flags["WRITEABLE"])

(1000, 3) (0, 8) False

Mil linhas "existem", mas ocupam o espaço de uma. É por isso que o broadcasting é rápido e econômico. O cuidado vem do outro lado: o resultado da conta (a grade (3, 4) do exemplo) é um array real, e em escalas grandes ele pode estourar a memória (capítulo 31).

Não é "qualquer array menor serve"

O erro típico é achar que, se um array é menor, ele sempre combina. O que conta é cada dimensão. Teste a forma com np.broadcast_shapes antes de rodar uma conta grande, ou leia a mensagem de erro: ela mostra as duas formas lado a lado.