Ufuncs from the inside

The methods of a ufunc

Every binary ufunc has methods that generalize the operation. The reduce applies the operation along an axis until only one value remains. The accumulate stores the partial results. The outer applies the operation to all pairs:

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]]

The np.add.reduce is the sum, the np.multiply.reduce is the product, and the np.add.accumulate is the cumsum. These "shortcut" functions are just convenient names for the methods. The advantage of knowing the methods is that they work with any ufunc: np.maximum.accumulate (the running maximum from chapter 28), np.logical_and.reduce (all true?), np.multiply.outer (a multiplication table).

The trap of repeated indices

This is one of the most famous surprises of NumPy. When an index appears more than once, the assignment a[indices] += 1 does not accumulate: each element receives the result of a single read. The reason is that the expression reads a[indices], adds 1, and writes back, all in three separate steps:

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

[1 1 1]

What would be expected is [2 1 3] (index 0 appeared twice, index 2, three times). The np.add.at performs the operation without buffering, and accumulates correctly. And, to count, the bincount is the tool made for that:

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]

The np.add.at is used for general cases (summing values at repeated positions), and the bincount is almost always the fastest option for counts and sums by groups of small integers.

`np.vectorize` does not make anything faster

The np.vectorize transforms a function that works with a value into one that accepts arrays. It is an interface convenience: underneath, it still calls the Python function once per element, and therefore does not have the speed of true vectorization:

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

A function with while and deviations (like the one from Collatz) does not have simple vectorized form, and the np.vectorize just makes it more convenient to call. For real speed, the options are to rewrite the algorithm in a vectorized way, or compile the loop with an external tool (like Numba).

`__array_ufunc__`: how other classes enter the game

When you call np.sqrt(objeto) with an object that does not is an array, NumPy asks the object what to do, calling its method __array_ufunc__ . This is how libraries like pandas, Dask, and CuPy allow NumPy functions to work on their own structures. A minimal example, a wrapper that "wraps" the result back:

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

The np.sqrt handed the object the ufunc, the method ("__call__" here) and the inputs, and the object decided: open, calculate, and wrap again. This is the contract that makes the ecosystem interoperable.