Ufuncs por dentro

Os métodos de uma ufunc

Toda ufunc binária tem métodos que generalizam a operação. O reduce aplica a operação ao longo de um eixo até sobrar um valor. O accumulate guarda os resultados parciais. O outer aplica a operação a todos os pares:

import numpy as np

v = np.array([1, 2, 3, 4])
print(np.add.reduce(v), np.multiply.reduce(v), np.add.accumulate(v), np.multiply.accumulate(v))
print(np.multiply.outer([1, 2, 3], [1, 2, 3]))

10 24 [ 1  3  6 10] [ 1  2  6 24]
[[1 2 3]
 [2 4 6]
 [3 6 9]]

O np.add.reduce é o sum, o np.multiply.reduce é o produto, e o np.add.accumulate é o cumsum. Essas funções "de atalho" são só nomes convenientes para os métodos. A vantagem de conhecer os métodos é que eles funcionam com qualquer ufunc: np.maximum.accumulate (o máximo corrido do capítulo 28), np.logical_and.reduce (todos verdadeiros?), np.multiply.outer (uma tabuada).

A armadilha dos índices repetidos

Esta é uma das surpresas mais famosas do NumPy. Quando um índice aparece mais de uma vez, a atribuição a[indices] += 1 não acumula: cada elemento recebe o resultado de uma só leitura. A razão é que a expressão lê a[indices], soma 1 e escreve de volta, tudo em três passos separados:

a = np.zeros(3, dtype=int)
indices = np.array([0, 0, 1, 2, 2, 2])
a[indices] += 1
print(a)

[1 1 1]

O esperado seria [2 1 3] (o índice 0 apareceu duas vezes, o 2, três). O np.add.at faz a operação sem buffer, e acumula corretamente. E, para contar, o bincount é a ferramenta feita para isso:

b = np.zeros(3, dtype=int)
np.add.at(b, indices, 1)
print(b)
print(np.bincount(indices, minlength=3))

[2 1 3]
[2 1 3]

O np.add.at serve para casos gerais (somar valores em posições repetidas), e o bincount é quase sempre a opção mais rápida para contagens e somas por grupo de inteiros pequenos.

`np.vectorize` não torna nada mais rápido

O np.vectorize transforma uma função que trabalha com um valor em uma que aceita arrays. É uma conveniência de interface: por baixo, ele continua chamando a função Python uma vez por elemento, e por isso não tem a velocidade de uma vetorização de verdade:

import timeit


def passos_collatz(n):
    passos = 0
    while n != 1:
        n = n // 2 if n % 2 == 0 else 3 * n + 1
        passos += 1
    return passos


vetorizada = np.vectorize(passos_collatz)
print(vetorizada(np.array([6, 7, 27])))

entradas = np.arange(1, 2000)
t_vec = min(timeit.repeat(lambda: vetorizada(entradas), number=1, repeat=3))
t_laco = min(timeit.repeat(lambda: [passos_collatz(int(n)) for n in entradas], number=1, repeat=3))
print("vectorize não é muito mais rápido que o laço:", t_vec > 0.5 * t_laco)

[  8  16 111]
vectorize não é muito mais rápido que o laço: True

Uma função com while e desvios (como a de Collatz) não tem forma vetorizada simples, e o np.vectorize só a deixa mais cômoda de chamar. Para velocidade de verdade, as opções são reescrever o algoritmo de forma vetorizada, ou compilar o laço com uma ferramenta externa (como o Numba).

`__array_ufunc__`: como outras classes entram no jogo

Quando você chama np.sqrt(objeto) com um objeto que não é um array, o NumPy pergunta ao objeto o que fazer, chamando o método __array_ufunc__ dele. É assim que bibliotecas como pandas, Dask e CuPy deixam as funções do NumPy funcionarem em estruturas próprias. Um exemplo mínimo, um invólucro que "embrulha" o resultado de volta:

class Embrulhado:
    def __init__(self, dados):
        self.dados = np.asarray(dados)

    def __array_ufunc__(self, ufunc, metodo, *entradas, **kwargs):
        abertas = [x.dados if isinstance(x, Embrulhado) else x for x in entradas]
        return Embrulhado(getattr(ufunc, metodo)(*abertas, **kwargs))


r = np.sqrt(Embrulhado([4, 9]))
print(type(r).__name__, r.dados)

Embrulhado [2. 3.]

O np.sqrt entregou ao objeto a ufunc, o método ("__call__" aqui) e as entradas, e o objeto decidiu: abrir, calcular e embrulhar de novo. Esse é o contrato que torna o ecossistema interoperável.