Python é pequeno, e o ecossistema é grande

O Python em si não sabe treinar uma rede neural, desenhar um gráfico nem ler uma planilha. Nada disso vem embutido. O que o Python tem é uma sintaxe fácil de ler e uma facilidade enorme de chamar código escrito em outras linguagens. Foi isso que atraiu quem trabalhava com cálculo numérico: dava para escrever o miolo pesado em C ou Fortran e usá-lo por uma interface simples em Python.

O NumPy nasceu em meados dos anos 2000, da unificação de dois projetos anteriores (o Numeric e o Numarray). Todo o resto da pilha de dados foi construído em cima dele:

Nenhuma dessas bibliotecas faz parte do Python. Cada uma foi escrita por pessoas, publicada de graça e mantida por uma comunidade. Dá para conferir isso no próprio interpretador:

import importlib.util
import sys

import numpy as np

print("numpy faz parte da biblioteca padrão?", "numpy" in sys.stdlib_module_names)
print("numpy está instalado aqui?", importlib.util.find_spec("numpy") is not None)
print(type(np.add).__name__, type(np.ndarray).__name__)

numpy faz parte da biblioteca padrão? False
numpy está instalado aqui? True
ufunc type

O np.add é um objeto do tipo ufunc, uma função compilada que opera sobre arrays inteiros. O trabalho pesado acontece nesse código compilado, e o Python só orquestra. Essa é a forma saudável de pensar em toda a pilha: Python é a cola, e a velocidade vem das camadas de baixo.

modulos = ["math", "numpy", "json", "pandas"]
da_biblioteca_padrao = [m for m in modulos if m in sys.stdlib_module_names]
assert da_biblioteca_padrao == ["math", "json"]
print("ok")

ok

Instalar

O NumPy se instala como qualquer pacote Python, dentro de um ambiente virtual, como o livro curso de Python mostrou. Escolha a ferramenta que você usa:

uv init meu-projeto
cd meu-projeto
uv add numpy

python3 -m venv .venv
source .venv/bin/activate
python -m pip install numpy

py -m venv .venv
.venv\Scripts\Activate.ps1
python -m pip install numpy

O NumPy distribui versões já compiladas (wheels) para os sistemas e versões do Python mais comuns. Por isso a instalação leva segundos e não exige compilador.

Importar, e a convenção `np`

import sys
from pathlib import Path

import numpy as np

print("versão 2 ou superior:", int(np.__version__.split(".")[0]) >= 2)
print("onde está o NumPy:", Path(np.__file__).parent.name)
print("este Python existe no disco:", Path(sys.executable).exists())

versão 2 ou superior: True
onde está o NumPy: numpy
este Python existe no disco: True

O apelido np não é uma regra da linguagem, mas é uma convenção tão forte que toda documentação, todo tutorial e toda resposta na internet a assumem. Usar outro apelido só faz você traduzir o que lê.

Para ver como o seu NumPy foi construído (quais bibliotecas de álgebra linear ele usa, por exemplo), existe o show_config:

np.show_config()

Quando o import falha

O erro mais comum é instalar com o pip em um terminal e receber No module named 'numpy' no editor. Quase sempre isso significa que são dois Pythons diferentes: o do terminal recebeu o pacote, e o do editor (ou do notebook) não. Para descobrir qual Python está rodando e se ele enxerga o NumPy:

python -c "import sys; print(sys.executable)"
python -m pip show numpy
python -c "import numpy; print(numpy.__version__, numpy.__file__)"

def diagnostico():
    return {
        "numpy": np.__version__,
        "python": sys.version_info[:2],
        "executavel_existe": Path(sys.executable).exists(),
    }


resultado = diagnostico()
assert resultado["executavel_existe"]
assert resultado["python"] >= (3, 12)
assert resultado["numpy"].split(".")[0].isdigit()
print("ok")

ok

Script ou notebook

Um script roda de cima para baixo, do começo ao fim, toda vez. Isso é ótimo para um programa que você vai executar no servidor, mas ruim para explorar dados, porque reprocessar tudo para olhar uma coluna é lento. Um notebook é dividido em células que você executa uma de cada vez, e o resultado fica visível logo abaixo, como em um caderno de laboratório.

Para usar o JupyterLab em um projeto com uv, instale-o como dependência de desenvolvimento e abra:

uv add --dev jupyterlab ipykernel
uv run jupyter lab

O kernel e a armadilha do estado escondido

Cada notebook é ligado a um kernel, o processo Python que lembra de todas as variáveis que você criou. As células não precisam rodar em ordem: você pode executar a quinta, voltar e reexecutar a segunda. É exatamente isso que causa o problema. O notebook pode mostrar um resultado que só existe por causa de uma célula que você editou depois, e o texto que você lê na tela não reproduz mais o resultado:

celulas = ["total = 10", "dobro = total * 2"]

kernel = {}
for celula in celulas:
    exec(celula, kernel)

celulas[0] = "total = 99"
exec(celulas[0], kernel)
print("o que a tela mostra:", kernel["dobro"])

novo_kernel = {}
for celula in celulas:
    exec(celula, novo_kernel)
print("depois de reiniciar e rodar tudo:", novo_kernel["dobro"])

o que a tela mostra: 20
depois de reiniciar e rodar tudo: 198

Eu editei a primeira célula para 99 e reexecutei só ela. O dobro na tela continuou 20, um valor que não corresponde mais ao código visível. Só ao reiniciar o kernel e rodar tudo em ordem aparece o valor verdadeiro, 198.

O melhor dos dois mundos: células em um arquivo `.py`

O VS Code e o JupyterLab (com a extensão jupytext) entendem um arquivo Python comum com marcadores # %%: cada marcador abre uma célula. O arquivo continua sendo um script que roda inteiro, com a vantagem de diferenças legíveis no Git, o que o formato .ipynb (um JSON gigante com saídas embutidas) não dá:

# %% [markdown]
# # Notas de vendas
# Cada bloco `# %%` é uma célula no VS Code e no JupyterLab (com jupytext).

# %%
import numpy as np

vendas = np.array([120.0, 80.5, 99.9, 150.0])

# %%
print(vendas.mean().round(2))

# %%
print(vendas.max())

Cada elemento de uma lista é um objeto

Em Python, 7 não é só o número sete. É um objeto inteiro, com cabeçalho, contagem de referências e tipo. A lista guarda um ponteiro para cada um desses objetos, e cada objeto mora em um lugar diferente da memória. Isso permite misturar tipos em uma mesma lista, ao custo de espaço e de velocidade. O array renuncia a essa flexibilidade de propósito:

import sys

import numpy as np

lista = list(range(1_000))
array = np.arange(1_000)

print(type(lista[0]).__name__, "ocupa", sys.getsizeof(lista[0]), "bytes cada")
print(array.dtype, "ocupa", array.itemsize, "bytes por elemento")

bytes_lista = sys.getsizeof(lista) + sum(sys.getsizeof(x) for x in lista)
print("a lista gasta mais de 3 vezes a memória do array:", bytes_lista > 3 * array.nbytes)

int ocupa 28 bytes cada
int64 ocupa 8 bytes por elemento
a lista gasta mais de 3 vezes a memória do array: True

Cada inteiro Python ocupa 28 bytes, enquanto cada elemento do array ocupa 8. E o array não tem ponteiros: os mil números estão em um bloco contínuo de 8 mil bytes.

Um tipo só, um bloco só

Como todos os elementos têm o mesmo tipo e o mesmo tamanho, o NumPy não precisa verificar o tipo de cada um nem seguir ponteiros. O processador lê a memória em sequência, o que aproveita o cache, e consegue operar em vários números por instrução. É a combinação "mesmo tipo, bloco contínuo" que permite entregar a operação inteira a código compilado.

O preço é que o array não aceita misturar. Se você tentar, o NumPy não reclama: converte tudo para um tipo comum, silenciosamente:

misturado = np.array([7, "a", 3.2, 9])
print(misturado.dtype, misturado)

<U32 ['7' 'a' '3.2' '9']

O <U32 quer dizer "texto Unicode de até 32 caracteres". Os números viraram texto, e você não poderia mais somá-los. Esse é um erro clássico: ele não faz barulho, e só aparece mais tarde, quando uma conta dá errado.

O bloco contínuo também é visível nas propriedades do array:

print(array.flags["C_CONTIGUOUS"], array.strides)

True (8,)

O strides diz quantos bytes o NumPy avança para chegar ao próximo elemento: 8, exatamente o tamanho de um int64. O capítulo 29 explica esse detalhe, que é a base de como o NumPy fatia arrays sem copiar nada.

Medir, em vez de acreditar

A diferença de velocidade é grande o bastante para uma medição simples a mostrar. Eu repito cada medida algumas vezes e fico com a melhor, que é a menos afetada pelo ruído da máquina. Como os tempos absolutos variam de computador para computador, o que eu mostro é a comparação:

import time


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


grande_lista = list(range(1_000_000))
grande_array = np.arange(1_000_000)
t_lista = medir(lambda: [x * 2 for x in grande_lista])
t_array = medir(lambda: grande_array * 2)
print("o array foi mais rápido:", t_array < t_lista)
print("pelo menos 10 vezes mais rápido:", t_lista > 10 * t_array)

o array foi mais rápido: True
pelo menos 10 vezes mais rápido: True

Na prática, a diferença costuma ficar em dezenas de vezes. E ela cresce com o tamanho: com poucos elementos, o custo fixo de chamar o NumPy domina e o ganho some. O capítulo 18 volta a esse ponto.

def bytes_do_array(n, dtype):
    return n * np.dtype(dtype).itemsize


assert bytes_do_array(1_000_000, "float64") == 8_000_000
assert bytes_do_array(1_000_000, "float32") == 4_000_000
assert bytes_do_array(1_000_000, "int8") == 1_000_000
print("ok")

ok

A partir dos seus dados

O np.array converte listas (e listas de listas) em um array. Listas aninhadas viram dimensões:

import numpy as np

print(np.array([1, 2, 3]))
print(np.array([[1, 2], [3, 4]]))

[1 2 3]
[[1 2]
 [3 4]]

Todas as linhas precisam ter o mesmo tamanho. Um array é uma grade retangular, e o NumPy recusa uma estrutura "irregular" em vez de adivinhar o que você quis dizer:

try:
    np.array([[1, 2], [3]])
except ValueError as erro:
    print("erro de forma irregular:", "inhomogeneous" in str(erro))

erro de forma irregular: True

Arrays de preenchimento: você dá a forma

Quando você ainda não tem os valores, mas sabe a forma, os arrays de preenchimento reservam o espaço. A forma é uma tupla, e é aí que mora o erro mais comum:

print(np.zeros((2, 3)))
print(np.ones((2, 3), dtype=int))
print(np.full((2, 3), 7))

[[0. 0. 0.]
 [0. 0. 0.]]
[[1 1 1]
 [1 1 1]]
[[7 7 7]
 [7 7 7]]

try:
    np.zeros(2, 3)
except TypeError as erro:
    print(erro)

Cannot interpret '3' as a data type

Escrevendo np.zeros(2, 3), o NumPy entende o 3 como o tipo de dado (o segundo parâmetro), e por isso a mensagem fala em "data type". A forma é um único argumento: np.zeros((2, 3)), com parênteses duplos.

O np.empty também reserva a forma, mas não inicializa a memória: o conteúdo é o que já estava lá, e não zero. Ele é um pouco mais rápido que o zeros, e só serve quando você vai sobrescrever todos os valores logo em seguida. Eu raramente o uso, porque o ganho é pequeno e o risco de ler lixo é real.

Matrizes especiais

print(np.eye(3))
print(np.diag([1, 2, 3]))

[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]
[[1 0 0]
 [0 2 0]
 [0 0 3]]

O np.eye(3) é a matriz identidade (uns na diagonal). O np.diag monta uma matriz diagonal a partir de uma lista, e, se receber uma matriz, faz o caminho inverso e extrai a diagonal.

Sequências

print(np.arange(0, 10, 2))
print(np.linspace(0, 1, 5))
print(np.logspace(0, 2, 3))

[0 2 4 6 8]
[0.   0.25 0.5  0.75 1.  ]
[  1.  10. 100.]

O arange funciona como o range do Python: início, fim (excluído) e passo. O linspace recebe o número de pontos e inclui os dois extremos. O logspace espaça os pontos em escala logarítmica, de 10⁰ a 10², três pontos.

print(len(np.arange(0, 1, 0.1)), len(np.linspace(0, 1, 11)))

10 11

Arrays com a forma de outro

As funções *_like criam um array com a mesma forma (e o mesmo tipo, a menos que você diga o contrário) de um array existente:

base = np.array([[1, 2, 3], [4, 5, 6]])
print(np.zeros_like(base))
print(np.ones_like(base, dtype=float).shape)

[[0 0 0]
 [0 0 0]]
(2, 3)

O capítulo 16 apresenta os arrays aleatórios, que também são uma forma de criação.

def multiplos_de(passo, limite):
    return np.arange(0, limite + 1, passo)


def matriz_constante(linhas, colunas, valor):
    return np.full((linhas, colunas), valor)


assert multiplos_de(5, 50).tolist() == [0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50]
assert matriz_constante(2, 3, 7).shape == (2, 3)
assert (matriz_constante(2, 3, 7) == 7).all()
print("ok")

ok

Os atributos

import numpy as np

arr = np.array([[1, 2, 3], [4, 5, 6]])
print(arr.shape, arr.ndim, arr.size)
print(arr.dtype, arr.itemsize, arr.nbytes)
print(arr.T)
print(arr.strides)

(2, 3) 2 6
int64 8 48
[[1 4]
 [2 5]
 [3 6]]
(24, 8)

O strides é o detalhe interno que torna o fatiamento tão rápido: para andar uma linha, o NumPy pula 24 bytes (3 elementos de 8 bytes), e para andar uma coluna, pula 8. Você não precisa calculá-lo, mas o capítulo 29 mostra como ele permite criar "visões" sem copiar dados.

Forma não é tamanho

O erro mais frequente de quem começa é trocar shape por size. Eles respondem a perguntas diferentes: o shape diz como os elementos estão organizados, e o size diz quantos são. Um array de forma (4, 5) tem tamanho 20.

Um array de uma dimensão tem forma de um elemento, e é preciso a vírgula na tupla. E há arrays de zero dimensões (escalares), cuja forma é a tupla vazia:

print(len(arr), np.array([1, 2, 3]).shape, np.array(5).shape, np.array(5).ndim)

2 (3,) () 0

O len(arr) devolve apenas o tamanho da primeira dimensão (2 linhas). Para o total de elementos, use size.

def resumo(a):
    return {"forma": a.shape, "dimensoes": a.ndim, "tamanho": a.size, "bytes": a.nbytes}


r = resumo(np.zeros((4, 5)))
assert r == {"forma": (4, 5), "dimensoes": 2, "tamanho": 20, "bytes": 160}
print("ok")

ok

Os tipos mais comuns

O NumPy descobre o tipo a partir dos valores, mas você pode (e muitas vezes deve) escolher:

import numpy as np

print(np.array([1, 2, 3]).dtype)
print(np.array([1.0, 2, 3]).dtype)
print(np.array([True, False]).dtype)
print(np.array([1 + 2j]).dtype)

int64
float64
bool
complex128

Converter com astype

O astype cria um array novo com outro tipo. A conversão de decimal para inteiro trunca (descarta a parte decimal), em vez de arredondar:

notas = np.array([7.9, 8.5, 9.99])
print(notas.astype(int))
print(np.round(notas).astype(int))

[7 8 9]
[ 8  8 10]

O np.round usa o arredondamento "para o par mais próximo" quando o valor está exatamente no meio: o 8.5 virou 8. É o mesmo comportamento do round do Python.

Quando o número não cabe

Cada tipo tem uma faixa fixa. Dentro de uma conta, passar do limite dá a volta (wraparound), sem erro e sem aviso. Um uint8 que chega a 255 e recebe mais 10 volta ao começo:

pequeno = np.array([250, 5], dtype=np.uint8)
print(pequeno + 10)

[ 4 15]

O 250 + 10 deu 4, e não 260. É um erro silencioso, e por isso um dos mais perigosos: nenhum aviso, e o resultado parece plausível. Já na criação, o NumPy 2 protege você quando o valor é impossível:

try:
    np.array([300], dtype=np.uint8)
except OverflowError as erro:
    print("OverflowError:", erro)

OverflowError: Python integer 300 out of bounds for uint8

Precisão de decimais

Os decimais binários não representam a maioria das frações exatamente, e o float32 tem ainda menos precisão do que o float64. A comparação por igualdade é uma armadilha, e a comparação com tolerância é o remédio:

print(0.1 + 0.2 == 0.3, np.isclose(0.1 + 0.2, 0.3))
print(np.finfo(np.float64).eps)
print(np.float32(16_777_216) + np.float32(1))

False True
2.220446049250313e-16
1.6777216e+07

O eps é a menor diferença que o float64 distingue perto de 1. E o último resultado (1.6777216e+07, ou seja, 16.777.216) mostra o limite do float32: a partir de 2²⁴ ele não consegue mais representar todos os inteiros, e somar 1 a 16.777.216 devolve o mesmo número.

O que acontece quando tipos diferentes se encontram

Quando você combina tipos em uma conta, o NumPy decide o tipo do resultado por regras de promoção. O NumPy 2 mudou uma regra importante (a NEP 50): um número Python puro ("escalar fraco") não promove o tipo do array, enquanto um escalar NumPy promove:

a = np.array([1, 2, 3], dtype=np.int8)
print((a + 1).dtype)
print((a + np.int64(1)).dtype)
print((np.float32(3) + 3.0).dtype)
print(np.result_type(np.int8, np.float32), np.result_type(np.int64, np.float32))

int8
int64
float32
float32 float64

O a + 1 continua int8, porque o 1 é um inteiro Python e "se adapta" ao array. Mas o a + np.int64(1) virou int64. E o np.result_type responde, sem fazer conta, qual seria o tipo do resultado.

def economizar(idades):
    return idades.astype(np.min_scalar_type(idades.max()))


idades = np.array([34, 71, 12, 98, 45, 120], dtype=np.int64)
menor = economizar(idades)
assert menor.dtype == np.uint8
assert menor.nbytes < idades.nbytes
assert menor.tolist() == idades.tolist()
print("ok")

ok

Um elemento, uma linha, uma coluna

Em arrays de várias dimensões, você separa os índices de cada eixo por vírgula, dentro de um único colchete: grade[linha, coluna]. Os dois pontos (:) significam "todos os elementos desse eixo":

import numpy as np

grade = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print(grade[1, 2])
print(grade[:, 1])
print(grade[0, ::-1])
print(grade[1:, :2])
print(grade[-1])

6
[2 5 8]
[3 2 1]
[[4 5]
 [7 8]]
[7 8 9]

O grade[1][2] também funciona, mas cria um array intermediário (a linha 1) antes de pegar o elemento, o que é mais lento e menos claro. Prefira sempre grade[1, 2].

Fatias: `inicio:fim:passo`

Como nas listas, o fim é excluído, os índices negativos contam do fim e o passo pode ser negativo. A diferença é que, no array, a fatia pode combinar um eixo por vez:

print(np.arange(10)[::3])

[0 3 6 9]

Atribuir a uma fatia

Uma fatia pode receber um valor (ou outro array), e a mudança acontece no array original. Para o exemplo não alterar a grade, eu trabalho em uma cópia:

copia = grade.copy()
copia[0, :] = 0
copia[:, 2] += 100
print(copia)

[[  0   0 100]
 [  4   5 106]
 [  7   8 109]]

Três dimensões e a reticência

Em arrays de 3 ou mais dimensões, a reticência (...) significa "todos os eixos que eu não mencionei". É muito útil para escrever código que funciona em qualquer número de dimensões:

cubo = np.arange(24).reshape(2, 3, 4)
print(cubo[1, 2, 3], cubo[:, 0, :].shape, cubo[..., 0].shape)

23 (2, 4) (2, 3)

O cubo[..., 0] pega o primeiro elemento do último eixo, de todos os outros, e por isso o resultado tem forma (2, 3).

Índice fora dos limites

Diferente das fatias, que são tolerantes, um índice isolado fora do limite levanta erro, e a mensagem diz exatamente qual eixo e qual tamanho:

try:
    grade[5]
except IndexError as erro:
    print(erro)

index 5 is out of bounds for axis 0 with size 3

def centro(a):
    return a[1:-1, 1:-1]


quadro = np.arange(16).reshape(4, 4)
assert np.array_equal(centro(quadro), np.array([[5, 6], [9, 10]]))
print("ok")

ok

Uma condição vira uma máscara

Comparar um array com um valor produz outro array, de booleanos, chamado de máscara. Usá-la como índice filtra os elementos verdadeiros:

import numpy as np

notas = np.array([78, 92, 45, 88, 91, 30, 85])
mascara = notas > 75
print(mascara)
print(notas[mascara])
print(notas[notas > 75])
print((notas > 75).sum(), (notas > 75).mean().round(3))

[ True  True False  True  True False  True]
[78 92 88 91 85]
[78 92 88 91 85]
5 0.714

Como True vale 1 e False vale 0, o sum() de uma máscara conta quantos elementos passam, e o mean() dá a proporção. Aqui, 5 de 7 notas (71,4%) passaram de 75. É um dos truques mais usados em análise de dados.

Várias condições: `&`, `|` e `~`

Para combinar condições, use os operadores & (e), | (ou) e ~ (não), e sempre coloque cada condição entre parênteses, porque esses operadores têm precedência maior do que a comparação:

presenca = np.array([95, 60, 80, 92, 70, 88, 97])
print(notas[(notas > 80) & (presenca > 90)])
print(notas[(notas < 50) | (presenca < 65)])
print(notas[~(notas > 80)])

[88 85]
[92 45 30]
[78 45 30]

Usar and e or do Python é o erro clássico, e o NumPy responde com uma mensagem que parece enigmática:

try:
    notas[(notas > 80) and (presenca > 90)]
except ValueError as erro:
    print(erro)

The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

O and do Python precisa decidir se um array inteiro é verdadeiro ou falso, e um array com vários elementos não tem resposta única. É por isso que existem o & e o |, que operam elemento a elemento.

Indexação fancy: uma lista de posições

Passando uma lista de índices, você escolhe as posições que quiser, na ordem que quiser, inclusive repetindo:

nums = np.array([10, 20, 30, 40, 50])
print(nums[[0, 2, 4]])
print(nums[[4, 4, 0]])

tabela = np.arange(12).reshape(3, 4)
print(tabela[[0, 2], :])
print(tabela[[0, 1], [1, 3]])

[10 30 50]
[50 50 10]
[[ 0  1  2  3]
 [ 8  9 10 11]]
[1 7]

Com duas listas, o NumPy as combina par a par: o último resultado é tabela[0, 1] e tabela[1, 3], e não um bloco 2×2.

Máscara e fancy devolvem cópias

Diferente da fatia, a indexação booleana e a fancy sempre devolvem uma cópia. Alterar o resultado não toca o original:

recorte = nums[[0, 1]]
recorte[0] = 999
print(nums[0], np.shares_memory(nums, recorte))

10 False

def aprovados(notas, minimo=60):
    return notas[notas >= minimo]


def acima_da_media(notas):
    return np.nonzero(notas > notas.mean())[0]


turma = np.array([50, 90, 60, 40, 80])
assert aprovados(turma).tolist() == [90, 60, 80]
assert acima_da_media(turma).tolist() == [1, 4]
print("ok")

ok

Uma view compartilha a memória

Quando você fatia um array, o NumPy não copia os números: cria um novo objeto que aponta para a mesma memória, com outro começo e outro tamanho. É por isso que o fatiamento custa quase nada, mesmo em arrays enormes. A consequência é que escrever na fatia escreve no original:

import numpy as np

original = np.arange(6)
fatia = original[1:4]
fatia[0] = 99
print(original)
print(np.shares_memory(original, fatia), fatia.base is original)

[ 0 99  2  3  4  5]
True True

O atributo base aponta para o array que é o dono da memória. Uma view tem base, e um array que possui os próprios dados tem base igual a None.

Quando você quer um array independente

Para uma cópia de verdade, peça explicitamente com .copy():

independente = original[1:4].copy()
independente[0] = -1
print(original[1], np.shares_memory(original, independente), independente.base is None)

99 False True

O mesmo objeto, dois nomes

A atribuição b = a não copia nada: cria um segundo nome para o mesmo array (o mesmo modelo de etiquetas do livro de Python). E os operadores de atribuição composta (+=) alteram o array no lugar, enquanto b = b + 10 cria um novo:

a = np.array([1, 2, 3])
b = a
b += 10
print(a)
b = b + 10
print(a, b)

[11 12 13]
[11 12 13] [21 22 23]

O b += 10 mudou também o a, porque os dois são o mesmo objeto. Já o b = b + 10 criou um array novo e o ligou ao nome b, e o a ficou como estava.

Funções que alteram o que recebem

Um array passado para uma função é o mesmo objeto, e não uma cópia. Se a função altera o array, quem a chamou vê a mudança. Isso é útil para economizar memória e perigoso quando acontece sem querer:

def zerar_negativos(a):
    a[a < 0] = 0
    return a


dados = np.array([3, -1, 4, -5])
resultado = zerar_negativos(dados)
print(dados, resultado is dados)

[3 0 4 0] True

Os dados originais foram modificados, e o resultado é o mesmo objeto. Uma versão que não altera a entrada usa uma operação que cria resultado novo, como np.maximum(a, 0).

`asarray` e `array`

O np.asarray(x) não copia se x já for um array, e o np.array(x) copia por padrão. Funções que aceitam "qualquer coisa parecida com um array" costumam começar com np.asarray(entrada), o que converte listas sem copiar o que já é array:

x = np.arange(3)
print(np.asarray(x) is x, np.array(x) is x)

True False

def normalizar(a):
    return (a - a.min()) / (a.max() - a.min())


entrada = np.array([10.0, 20.0, 30.0])
copia_da_entrada = entrada.copy()
saida = normalizar(entrada)
assert np.array_equal(entrada, copia_da_entrada)
assert saida.tolist() == [0.0, 0.5, 1.0]
assert not np.shares_memory(entrada, saida)
print("ok")

ok

Remodelar sem mexer nos dados

O reshape reorganiza os mesmos elementos em outra forma. O total precisa continuar o mesmo: ele nunca cria nem remove elementos. E o -1 em uma das dimensões significa "calcule você":

import numpy as np

arr = np.arange(12)
print(arr.reshape(3, 4))
print(arr.reshape(2, -1).shape, arr.reshape(-1, 3).shape)
try:
    arr.reshape(5, 3)
except ValueError as erro:
    print(erro)

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
(2, 6) (4, 3)
cannot reshape array of size 12 into shape (5,3)

Cinco vezes três dá 15, e o array tem 12 elementos, então o NumPy recusa. O reshape(2, -1) descobriu sozinho que a segunda dimensão é 6.

Transpor, achatar e inserir eixos

grade = np.arange(6).reshape(2, 3)
print(grade.T.shape, grade.transpose().shape)
print(grade.flatten(), grade.ravel())
print(grade[:, np.newaxis, :].shape, grade[..., None].shape)
print(np.squeeze(np.zeros((1, 3, 1))).shape)

(3, 2) (3, 2)
[0 1 2 3 4 5] [0 1 2 3 4 5]
(2, 1, 3) (2, 3, 1)
(3,)

A diferença entre flatten e ravel é a mesma do capítulo anterior: um copia, o outro compartilha memória, e alterar o resultado do ravel muda o original:

r = grade.ravel()
r[0] = 99
print(grade[0, 0])
f = grade.flatten()
f[0] = -1
print(grade[0, 0])

99
99

O truque do eixo extra

Inserir um eixo de tamanho 1 é a forma de preparar um vetor para o broadcasting do próximo capítulo. Um vetor de forma (3,) vira uma coluna (3, 1) com reshape(-1, 1) ou com [:, None]:

vetor = np.array([1, 2, 3])
print(vetor.reshape(-1, 1))

[[1]
 [2]
 [3]]

Um caso real: lote de imagens

Um lote de 10 imagens de 28 por 28 pixels tem forma (10, 28, 28). Muitos modelos esperam cada imagem como um vetor de 784 números, ou seja, forma (10, 784). O -1 calcula o 784 sozinho:

lote = np.zeros((10, 28, 28))
print(lote.reshape(10, -1).shape)

(10, 784)

def para_coluna(v):
    return v.reshape(-1, 1)


def lote_para_planilha(imagens):
    return imagens.reshape(imagens.shape[0], -1)


assert para_coluna(np.array([1, 2, 3])).shape == (3, 1)
assert lote_para_planilha(np.zeros((10, 28, 28))).shape == (10, 784)
print("ok")

ok

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]

def imc(pesos, alturas):
    return pesos / alturas ** 2


def soma_de_quadrados_pares(a):
    return (a[a % 2 == 0] ** 2).sum()


resultado = imc(np.array([70.0, 90.0]), np.array([1.75, 1.80]))
assert np.allclose(resultado, [22.857, 27.778], atol=0.001)
assert soma_de_quadrados_pares(np.arange(10)) == 120
print("ok")

ok

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).

def padronizar_colunas(x):
    return (x - x.mean(axis=0)) / x.std(axis=0)


dados = np.array([[1.0, 100.0], [2.0, 200.0], [3.0, 300.0], [4.0, 400.0]])
z = padronizar_colunas(dados)
assert np.allclose(z.mean(axis=0), 0)
assert np.allclose(z.std(axis=0), 1)
print("ok")

ok

As agregações mais usadas

import numpy as np

dados = np.array([[4, 9, 2], [11, 6, 15]])
print(dados.sum(), dados.mean(), np.median(dados))
print(dados.sum(axis=0), dados.sum(axis=1))
print(dados.max(axis=1), dados.argmax(), dados.argmax(axis=1))

47 7.833333333333333 7.5
[15 15 17] [15 32]
[ 9 15] 5 [1 2]

Sem axis, a agregação achata o array e resume tudo em um número só. Com axis, o NumPy resume ao longo daquele eixo, e a dimensão some do resultado.

O que o axis faz

Eu uso este modelo mental: axis=k é o eixo que desaparece do resultado. Em axis=0, o eixo das linhas desaparece e sobram as colunas.

O argmax e o argmin devolvem a posição do maior e do menor valor (sem axis, a posição no array achatado), em vez do valor.

Dispersão e posição

O desvio padrão e a variância têm uma sutileza: por padrão, o NumPy divide por N (a variância da população). Para tratar os dados como uma amostra, divide-se por N - 1, com ddof=1:

print(dados.std().round(3), dados.var().round(3), dados.std(ddof=1).round(3))
print(np.percentile(dados, [25, 50, 75]))

4.375 19.139 4.792
[ 4.5  7.5 10.5]

O np.percentile devolve os valores abaixo dos quais fica a porcentagem pedida: o de 50% é a mediana.

Mediana resiste a valores extremos, a média não

Um único valor extremo arrasta a média para longe do que é "típico", e a mediana mal se mexe. Por isso, em dados de salário, preço de imóveis e tempo de resposta, a mediana costuma ser a medida honesta:

salarios = np.array([3000, 3200, 3100, 2900, 50000])
print(salarios.mean(), np.median(salarios))

12440.0 3100.0

Contar valores distintos

O np.unique com return_counts=True devolve os valores diferentes e quantas vezes cada um aparece, sem laço:

votos = np.array([2, 1, 2, 3, 2, 1])
valores, contagens = np.unique(votos, return_counts=True)
print(valores, contagens)

[1 2 3] [2 3 1]

Agregar com `keepdims`

Mantendo a dimensão reduzida, o resultado ainda "encaixa" no array original por broadcasting. É o que o capítulo anterior usou para normalizar linhas:

print(dados.sum(axis=1, keepdims=True))
print(dados / dados.sum(axis=1, keepdims=True))

[[15]
 [32]]
[[0.26666667 0.6        0.13333333]
 [0.34375    0.1875     0.46875   ]]

def resumo_por_coluna(dados):
    return {
        "minimo": dados.min(axis=0),
        "maximo": dados.max(axis=0),
        "media": dados.mean(axis=0),
    }


r = resumo_por_coluna(np.array([[4, 9, 2], [11, 6, 15]]))
assert r["minimo"].tolist() == [4, 6, 2]
assert r["maximo"].tolist() == [11, 9, 15]
assert r["media"].tolist() == [7.5, 7.5, 8.5]
print("ok")

ok

Matemática sobre o array inteiro

Uma ufunc (função universal) aplica uma operação a cada elemento. Os operadores +, -, * e / são, por baixo, np.add, np.subtract, np.multiply e np.divide:

import numpy as np

a = np.array([1.0, 4.0, 9.0])
print(np.sqrt(a), np.exp(np.array([0, 1])), np.log(np.array([1, np.e])))
print(np.add(a, 1), np.multiply(a, a))
print(np.abs(np.array([-3, 2])), np.round(np.array([1.234, 5.678]), 1))
print(np.clip(np.array([-5, 3, 12]), 0, 10))

[1. 2. 3.] [1.         2.71828183] [0. 1.]
[ 2.  5. 10.] [ 1. 16. 81.]
[3 2] [1.2 5.7]
[ 0  3 10]

O np.clip(x, mínimo, máximo) "aperta" os valores para dentro de um intervalo: tudo abaixo do mínimo vira o mínimo, e tudo acima do máximo vira o máximo. É muito usado para limpar dados com valores impossíveis.

Resultados impossíveis viram `nan` e `inf`

Uma conta impossível (raiz de número negativo, divisão por zero) não derruba o programa: o NumPy emite um aviso e devolve um valor especial, nan ("não é um número") ou inf (infinito). A execução continua:

import warnings

with warnings.catch_warnings(record=True) as avisos:
    warnings.simplefilter("always")
    resultado = np.sqrt(np.array([4.0, -4.0]))
print(resultado, avisos[0].category.__name__)

[ 2. nan] RuntimeWarning

Esse comportamento é conveniente, e também perigoso: o nan se espalha. Qualquer conta com um nan produz nan, e o erro aparece longe de onde nasceu. O capítulo 25 trata dos dados ausentes. Se você sabe o que está fazendo, o np.errstate silencia os avisos em um trecho:

with np.errstate(invalid="ignore", divide="ignore"):
    print(np.array([1.0, 0.0]) / np.array([0.0, 0.0]))

[inf nan]

Escolher onde calcular: `where` e `out`

As ufuncs aceitam dois parâmetros que economizam memória e evitam contas indesejadas. O where diz em quais posições calcular, e o out diz onde guardar o resultado, em vez de criar um array novo. Onde o where é falso, o valor de out permanece:

x = np.array([1.0, 4.0, 9.0, 16.0])
saida = np.zeros_like(x)
np.sqrt(x, out=saida, where=x > 5)
print(saida)

[0. 0. 3. 4.]

Divisão inteira e resto

Os operadores // e % também são ufuncs, e a divisão inteira arredonda para baixo, mesmo com números negativos:

print(np.array([7, -7]) // 2, np.mod(-7, 3))

[ 3 -4] 2

def sigmoide(x):
    return 1 / (1 + np.exp(-x))


assert sigmoide(0) == 0.5
assert np.allclose(sigmoide(np.array([-1.0, 1.0])), [0.26894142, 0.73105858])
assert (np.diff(sigmoide(np.linspace(-5, 5, 20))) > 0).all()
print("ok")

ok

A forma moderna: `default_rng`

Você cria um gerador com uma semente (seed) e pede números a ele. A mesma semente produz sempre a mesma sequência, e é isso que torna um resultado reproduzível:

import numpy as np

rng = np.random.default_rng(42)
print(rng.random(3))
print(rng.integers(0, 10, size=5))
print(rng.normal(loc=0, scale=1, size=3))
print(rng.choice([10, 20, 30], size=2, replace=False))

[0.77395605 0.43887844 0.85859792]
[0 6 2 0 5]
[ 0.1278404  -0.31624259 -0.01680116]
[30 20]

a = np.random.default_rng(7).random(3)
b = np.random.default_rng(7).random(3)
print(np.array_equal(a, b))

True

Por que não `np.random.seed`

A API antiga (np.random.seed(42) seguida de np.random.rand(3)) ainda funciona e aparece em muito material, inclusive em cursos, mas usa um estado global escondido no módulo:

np.random.seed(42)
print(np.random.rand(3))

[0.37454012 0.95071431 0.73199394]

O problema do estado global é que qualquer código, em qualquer lugar, pode consumir números dele e mudar a sua sequência sem você perceber, o que quebra a reprodutibilidade e torna os testes frágeis. Com o Generator, o estado mora em um objeto que você controla. E, por isso, funções que usam aleatoriedade devem receber o gerador como parâmetro:

def sorteio(n, rng):
    return rng.integers(1, 7, size=n)


print(sorteio(5, np.random.default_rng(1)).tolist())

[3, 4, 5, 6, 1]

Com essa assinatura, o teste da função passa um gerador com semente fixa e verifica um resultado exato, e a produção passa um gerador qualquer.

Embaralhar: no lugar ou em cópia

O rng.shuffle embaralha o próprio array e devolve None. Escrever x = rng.shuffle(x) joga fora o array. Para obter um embaralhado e manter o original, use permutation:

vetor = np.arange(1, 6)
resultado = rng.shuffle(vetor)
print(resultado, sorted(vetor.tolist()))

original = np.arange(1, 6)
embaralhado = rng.permutation(original)
print(original, sorted(embaralhado.tolist()))

None [1, 2, 3, 4, 5]
[1 2 3 4 5] [1, 2, 3, 4, 5]

Gerar um conjunto sintético

Dados sintéticos com características conhecidas permitem testar uma análise sabendo a resposta. Aqui, mil alturas com média 170 e desvio 10: a média amostral fica perto de 170, mas não exatamente nela, porque é uma amostra:

alturas = rng.normal(170, 10, size=1000)
print(abs(alturas.mean() - 170) < 1, abs(alturas.std() - 10) < 1)

sorteados = rng.choice(["cara", "coroa"], size=1000, p=[0.9, 0.1])
print((sorteados == "cara").mean() > 0.8)

True True
True

def simular_dados(n, rng):
    lancamentos = rng.integers(1, 7, size=n)
    return np.bincount(lancamentos, minlength=7)[1:]


contagens = simular_dados(6000, np.random.default_rng(0))
assert contagens.shape == (6,)
assert contagens.sum() == 6000
assert (contagens > 800).all()
print("ok")

ok

As dimensões, de zero a cinco

import numpy as np

escalar = np.array(7)
vetor = np.array([255, 0, 0])
matriz = np.zeros((28, 28))
imagem = np.zeros((224, 224, 3), dtype=np.uint8)
lote = np.zeros((32, 224, 224, 3), dtype=np.float32)

for nome, t in [("escalar", escalar), ("vetor", vetor), ("matriz", matriz), ("imagem", imagem), ("lote", lote)]:
    print(f"{nome:<8} ndim={t.ndim} shape={t.shape}")
print(round(lote.nbytes / 1e6, 1), "MB")

escalar  ndim=0 shape=()
vetor    ndim=1 shape=(3,)
matriz   ndim=2 shape=(28, 28)
imagem   ndim=3 shape=(224, 224, 3)
lote     ndim=4 shape=(32, 224, 224, 3)
19.3 MB

O último número é um aviso: um lote de 32 imagens de 224 por 224 em float32 ocupa cerca de 19 MB. Tensores crescem rápido, e o dtype e o tamanho do lote decidem se o seu computador aguenta.

Ler uma forma como um idioma

Quando um tutorial mostrar uma forma como (32, 224, 224, 3), leia como lote × altura × largura × canais. Esse hábito de leitura resolve a maior parte da confusão dos primeiros dias com aprendizado profundo.

Existem duas convenções para o lugar dos canais. O NumPy e o TensorFlow costumam usar canais no fim (altura, largura, canais), e o PyTorch usa canais antes (canais, altura, largura). O np.moveaxis troca de uma para a outra:

chw = np.moveaxis(imagem, -1, 0)
print(chw.shape, np.transpose(chw, (1, 2, 0)).shape)

(3, 224, 224) (224, 224, 3)

A mesma regra de sempre, em mais dimensões

Indexação, fatiamento, broadcasting e agregação funcionam exatamente como antes. Com axis você escolhe sobre quais dimensões resumir. Para calcular a média de cada canal de cor em um lote inteiro, resuma lote, altura e largura de uma vez, e sobram os canais:

print(lote.mean(axis=(0, 1, 2)).shape)

(3,)

Uma imagem de verdade, em miniatura

Uma imagem colorida é um tensor. Em uma imagem de 2 por 2 pixels, cada pixel tem três canais (vermelho, verde, azul), de 0 a 255:

pixels = np.array(
    [[[255, 0, 0], [0, 255, 0]], [[0, 0, 255], [255, 255, 255]]],
    dtype=np.uint8,
)
print(pixels.shape)
print(pixels[0, 1])
print(pixels[..., 0])
cinza = pixels.mean(axis=-1).astype(np.uint8)
print(cinza)

(2, 2, 3)
[  0 255   0]
[[255   0]
 [  0 255]]
[[ 85  85]
 [ 85 255]]

Os quatro pixels são vermelho, verde, azul e branco. pixels[..., 0] extrai o canal vermelho, e a média dos três canais produz uma versão em tons de cinza (um jeito simples, e não o mais fiel à percepção humana, de converter).

def para_lote(imagens):
    return np.stack(imagens)


imagens = [np.zeros((8, 8, 3)) for _ in range(4)]
assert para_lote(imagens).shape == (4, 8, 8, 3)
print("ok")

ok

O `timeit` e o melhor de vários

Um tempo medido uma vez é ruído: o processador pode estar ocupado, o cache pode estar frio. O módulo timeit repete a medida, e a regra é usar o menor dos tempos, que é o menos afetado pelo que o seu computador estava fazendo ao mesmo tempo:

import timeit

import numpy as np

a = np.arange(1_000_000)
tempos = timeit.repeat(lambda: a.sum(), number=20, repeat=5)
print("melhor tempo por chamada abaixo de 10 ms:", min(tempos) / 20 < 0.01)

melhor tempo por chamada abaixo de 10 ms: True

Em um notebook, %timeit a.sum() faz tudo isso por você. Em um script, use timeit.repeat e olhe o min.

Quanto custa começar: tamanhos pequenos não contam a história

Chamar uma função do NumPy tem um custo fixo (verificar tipos, alocar o array de saída). Para poucos elementos, esse custo domina, e o array pode até perder para uma lista. A vantagem só aparece, e cresce, quando os dados crescem:

def razao(n):
    lista = list(range(n))
    array = np.arange(n)
    t_lista = min(timeit.repeat(lambda: [x * 2 for x in lista], number=50, repeat=5))
    t_array = min(timeit.repeat(lambda: array * 2, number=50, repeat=5))
    return t_lista / t_array


print("a vantagem cresce com o tamanho:", razao(100_000) > razao(100))
print("em 100 mil elementos o array ganha:", razao(100_000) > 5)

a vantagem cresce com o tamanho: True
em 100 mil elementos o array ganha: True

Por isso, medir com 10 elementos e concluir que "o NumPy não é mais rápido" é um erro clássico. O teste tem de ser feito no tamanho real dos seus dados.

Medir memória

Tempo é só uma metade. Muitas análises falham por falta de memória, e não por lentidão. O tracemalloc mede a memória alocada pelo Python, e o NumPy reporta os seus arrays a ele:

import tracemalloc

tracemalloc.start()
x = np.zeros(1_000_000)
atual, pico = tracemalloc.get_traced_memory()
tracemalloc.stop()
print(round(atual / 1e6, 1), "MB")

8.0 MB

Um milhão de float64 ocupa 8 MB, exatamente size * itemsize. Para um array que já existe, o nbytes dá a resposta sem medir nada.

Armadilhas de medição

def comparar(f, g, repeticoes=5):
    t_f = min(timeit.repeat(f, number=10, repeat=repeticoes))
    t_g = min(timeit.repeat(g, number=10, repeat=repeticoes))
    return t_f / t_g


razao_medida = comparar(lambda: sum(range(10000)), lambda: np.arange(10000).sum())
assert razao_medida > 1
print("ok")

ok

Juntar ao longo de um eixo

O np.concatenate une arrays ao longo de um eixo que já existe. Todas as outras dimensões precisam ser iguais:

import numpy as np

a = np.array([[1, 2], [3, 4]])
b = np.array([[5, 6]])
print(np.concatenate([a, b], axis=0))
print(np.concatenate([a, b.T], axis=1))

[[1 2]
 [3 4]
 [5 6]]
[[1 2 5]
 [3 4 6]]

Quando as dimensões não batem, o erro diz exatamente qual:

try:
    np.concatenate([a, b], axis=1)
except ValueError as erro:
    print(erro)

all the input array dimensions except for the concatenation axis must match exactly, but along dimension 0, the array at index 0 has size 2 and the array at index 1 has size 1

Empilhar em um eixo novo

O np.stack cria um eixo novo e coloca cada array em uma posição dele. É o que você quer para juntar vários vetores de mesma forma em uma matriz, ou várias imagens em um lote. Os atalhos vstack e hstack empilham na vertical e na horizontal:

x = np.array([1, 2, 3])
y = np.array([4, 5, 6])
print(np.stack([x, y]).shape, np.stack([x, y], axis=1).shape)
print(np.vstack([x, y]), np.hstack([x, y]))

(2, 3) (3, 2)
[[1 2 3]
 [4 5 6]] [1 2 3 4 5 6]

Dividir

O np.split divide em partes iguais (e reclama se não der), o np.array_split aceita partes de tamanhos diferentes, e uma lista de posições diz onde cortar:

dados = np.arange(10)
print(np.split(dados, 5))
print(np.array_split(dados, 3))
print(np.split(dados, [2, 7]))

[array([0, 1]), array([2, 3]), array([4, 5]), array([6, 7]), array([8, 9])]
[array([0, 1, 2, 3]), array([4, 5, 6]), array([7, 8, 9])]
[array([0, 1]), array([2, 3, 4, 5, 6]), array([7, 8, 9])]

O uso mais comum é separar treino e teste: embaralhe as posições (capítulo 16) e corte. Embaralhar antes é essencial, porque dados guardados em ordem (por data, por classe) dariam um teste enviesado:

rng = np.random.default_rng(0)
ordem = rng.permutation(10)
treino, teste = np.split(ordem, [8])
print(len(treino), len(teste), sorted(np.concatenate([treino, teste]).tolist()) == list(range(10)))

8 2 True

Repetir

print(np.repeat([1, 2], 3), np.tile([1, 2], 3))

[1 1 1 2 2 2] [1 2 1 2 1 2]