Vectorization

The loop and the operation on the array

In pure Python, to double all elements you go through one by one. In NumPy, you write the operation on the array, and it applies to all elements:

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]

Notice that the loop returned a list (of NumPy scalars), and not an array. Those who write for x in array: lista.append(...) abandon the advantages of NumPy: they lose the type, lose speed, and still need to convert the list back with np.array(resultado).

Whole formulas, no loop

The formulas are written as in mathematics, and apply to one element or a million:

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

The second account is the total of an order: price times quantity, element by element, and then the sum. Two lines, no loop.

Why it's faster

A Python loop pays a small tax to each iteration: checking the type, managing the loop, returning to the interpreter. For 5 elements it is invisible, and for millions it dominates the time. In the vectorized version, the loop exists, but runs in compiledcode, without this tax. The difference is easy to measure:

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

The result is identical, and the vectorized version is, in this case, much faster (with a loop over an array, it is worse than over a list, because each a[i] creates a scalar object from NumPy).

The error of using functions that are not from NumPy

The functions from the module math expect a number, not an array. NumPy has a vectorized version of each one:

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

When the loop is the right answer

Not everything can be vectorized. When each step depends on the result of the previous one (a balance that earns interest and receives contributions, a simulation, an iterative algorithm), the loop remains the natural form. What matters is the question: "is there a vectorized version?":

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]

The rule of thumb

When writing a for over elements of an array, I stop and look for the vectorized form first: arithmetic operators, comparisons, np.where, aggregations with axis, universal functions. When they do not solve it, it means the problem has dependencies between the steps, and then the loop is legitimate (or it's time to look at cumsum, accumulate and chapter 34).