Universal functions and vectorised arithmetic

A Python function that works out a delivery fee handles one order at a time. Wrapped in np.vectorize it accepts an array of 100,000 orders, yet it still took about 15 ms in this lesson’s run, while the same rule written with NumPy’s universal functions took under 0.2 ms.

You will see what a ufunc is and how it picks a compiled loop, predict integer division, remainders and rounding, divide safely with out= and where=, turn any two-input ufunc into a reduction, a running total or a table with its methods, and measure why np.vectorize is a convenience rather than a speed-up.

Every operator on arrays calls a ufunc, a set of compiled loops

A universal functionufunc: A universal function: a NumPy function such as np.add or np.sqrt that works element by element on whole arrays, with one compiled loop for each combination of input types it supports. Operators such as + call ufuncs., or ufunc, is a NumPy function that, in the words of NumPy’s guide, “operates on ndarrays in an element-by-element fashion”, with broadcasting and type casting built in. Arithmetic, comparisons, np.exp, np.log, np.sin and their relatives are all ufuncs, and the operators call them: on arrays, a + b is np.add(a, b). Inside, a ufunc holds one compiled loop for each combination of types it supports:

"""Operators on arrays are ufuncs: each is a set of compiled loops, one for each combination of dtypes."""
import numpy as np

a = np.array([1, 2, 3], dtype=np.int64)
b = np.array([10, 20, 30], dtype=np.int64)
print("a + b equals np.add(a, b):", np.array_equal(a + b, np.add(a, b)))
print("np.add is a", type(np.add).__name__, "with", np.add.nin, "inputs,", np.add.nout, "output and", np.add.ntypes, "typed loops")
print("its first loops:", np.add.types[:6])
print("int64 + int64 ->", (a + b).dtype, "| int64 + float64 ->", (a + b.astype(np.float64)).dtype)
print("a > 1:", a > 1, "->", (a > 1).dtype)
print("np.sqrt(a):", np.sqrt(a).round(4), "| np.exp(0):", np.exp(0), "| np.log(np.e):", np.log(np.e))
print("np.sin(np.pi / 2):", np.sin(np.pi / 2), "| np.sin(np.pi):", np.sin(np.pi))
u = np.array([200, 100], dtype=np.uint8)
print("uint8 [200, 100] + itself:", u + u, "(it wraps at 256)")

a + b equals np.add(a, b): True
np.add is a ufunc with 2 inputs, 1 output and 22 typed loops
its first loops: ['??->?', 'bb->b', 'BB->B', 'hh->h', 'HH->H', 'ii->i']
int64 + int64 -> int64 | int64 + float64 -> float64
a > 1: [False  True  True] -> bool
np.sqrt(a): [1.     1.4142 1.7321] | np.exp(0): 1.0 | np.log(np.e): 1.0
np.sin(np.pi / 2): 1.0 | np.sin(np.pi): 1.2246467991473532e-16
uint8 [200, 100] + itself: [144 200] (it wraps at 256)

np.add takes 2 inputs, gives 1 output and holds 22 typed loops; its list starts with ??->? for booleans and bb->b for int8. Two int64 arrays use the int64 loop and give int64, an int64 and a float64 array give float64, comparisons give booleans, and np.sqrt of integers gives floats. np.sin(np.pi) is 1.2e-16 rather than 0, because np.pi is only the float nearest to π. And the uint8 loop keeps its result in uint8, so 200 + 200 wrapped round to 144 without a warning. Step through what happens between a + b and its result:

Every later section builds on this: the dtypes choose the loop, the loop decides how results overflow, and out= decides where they are written.

Integer division, remainders and rounding have rules of their own

Four operations surprise people coming from other languages or from school arithmetic. Here they are side by side:

"""Integer division, remainders, powers and rounding follow rules of their own."""
import numpy as np

a = np.array([-7, 7, -7, 7])
b = np.array([2, 2, -2, -2])
print("a // b:", a // b, "| a % b:", a % b, "| np.fmod(a, b):", np.fmod(a, b))
print("Python: -7 // 2 =", -7 // 2, "and -7 % 2 =", -7 % 2)
try:
    np.array([2, 3]) ** -1
except ValueError as error:
    print("integer array ** -1: ValueError:", error)
print("np.array([2, 3]) ** -1.0:", np.array([2, 3]) ** -1.0)
print("np.round([0.5, 1.5, 2.5, 3.5, -0.5, -1.5]):", np.round([0.5, 1.5, 2.5, 3.5, -0.5, -1.5]))
print("Python round(2.5) and round(0.5):", round(2.5), round(0.5))
print("np.round(2.675, 2):", np.round(2.675, 2), "| round(2.675, 2):", round(2.675, 2))
print("np.floor, np.ceil and np.trunc of -2.5:", np.floor(-2.5), np.ceil(-2.5), np.trunc(-2.5))

a // b: [-4  3  3 -4] | a % b: [ 1  1 -1 -1] | np.fmod(a, b): [-1  1 -1  1]
Python: -7 // 2 = -4 and -7 % 2 = 1
integer array ** -1: ValueError: Integers to negative integer powers are not allowed.
np.array([2, 3]) ** -1.0: [0.5        0.33333333]
np.round([0.5, 1.5, 2.5, 3.5, -0.5, -1.5]): [ 0.  2.  2.  4. -0. -2.]
Python round(2.5) and round(0.5): 2 0
np.round(2.675, 2): 2.68 | round(2.675, 2): 2.67
np.floor, np.ceil and np.trunc of -2.5: -3.0 -2.0 -2.0

// is floor divisionfloor division: Division that rounds the quotient down, towards minus infinity: -7 // 2 is -4. Its remainder, -7 % 2, is 1, which has the sign of the divisor, so the quotient times the divisor plus the remainder gives -7 back., as in Python: -7 // 2 rounds −3.5 down to −4, and the remainder makes up the rest, -7 % 2 is 1, because NumPy’s reference says the remainder “has the same sign as the divisor”. np.fmod truncates instead, as C and MATLAB’s rem do, so its remainders take the dividend’s sign. An integer array raised to a negative integer power raises a ValueError, since the answer would not be an integer; a float exponent, ** -1.0, gives 0.5 and 0.333….

np.round uses rounding half to evenrounding half to even: A rounding rule that sends a value exactly halfway between two candidates to the even one: 0.5 and -0.5 to 0, 1.5 and 2.5 to 2. np.round and Python's round both use it.: 0.5 and −0.5 go to 0, 1.5 and 2.5 to 2, 3.5 to 4, and Python’s own round does the same. The last line shows a different trap. np.round(2.675, 2) gave 2.68 while Python gave 2.67: 2.675 is stored as 2.67499999…, which Python rounds correctly down, but NumPy’s reference describes its method as “fast but sometimes inexact”: it multiplies by 100, and 2.675 × 100 rounds to exactly 267.5 in floating point, which goes to the even 268.

out= reuses memory and where= skips positions

Four shops report their revenue and their number of orders, and two of them had no orders. Dividing one array by the other gives inf and nan for those two, with a warning for each. A ufunc can instead compute only where a condition holds and write into an array you supply. Step through np.divide with both:

The run divides the plain way first, then with out= and where=, and then shows two more uses of the arguments:

"""Divide without dividing by zero: where= chooses the positions, out= holds what the others keep."""
import warnings

import numpy as np

sales = np.array([3.0, 5.0, 0.0, 4.0])  # revenue of four shops
orders = np.array([2.0, 0.0, 0.0, 8.0])  # their numbers of orders
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    plain = sales / orders
print("sales / orders:", plain)
for w in caught:
    print("  ", f"{w.category.__name__}: {w.message}")
out = np.full(4, -1.0)  # what a skipped position keeps
result = np.divide(sales, orders, out=out, where=orders != 0)
print("np.divide(sales, orders, out=out, where=orders != 0):", result, "| result is out:", result is out)
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    np.divide(sales, orders, where=orders != 0)
print("where= without out=:", f"{caught[0].category.__name__}: {caught[0].message}")
with np.errstate(divide="ignore", invalid="ignore"):
    quiet = sales / orders
print("inside np.errstate:", quiet, "| np.isfinite:", np.isfinite(quiet))
total = np.zeros(4)
np.add(total, sales, out=total)  # the result goes into total: no new array
print("np.add(total, sales, out=total):", total)
small = np.array([200, 100], dtype=np.uint8)
print("np.add(small, small):", np.add(small, small), "| with dtype=np.int16:", np.add(small, small, dtype=np.int16))

sales / orders: [1.5 inf nan 0.5]
   RuntimeWarning: divide by zero encountered in divide
   RuntimeWarning: invalid value encountered in divide
np.divide(sales, orders, out=out, where=orders != 0): [ 1.5 -1.  -1.   0.5] | result is out: True
where= without out=: UserWarning: 'where' used without 'out', expect uninitialized memory in output. If this is intentional, use out=None.
inside np.errstate: [1.5 inf nan 0.5] | np.isfinite: [ True False False  True]
np.add(total, sales, out=total): [3. 5. 0. 4.]
np.add(small, small): [144 200] | with dtype=np.int16: [400 200]

The plain division printed [1.5 inf nan 0.5] and two RuntimeWarnings. With out=np.full(4, -1.0) and where=orders != 0, nothing was divided by zero and the two shops without orders kept −1.0: NumPy’s ufunc reference says a False position means “leave the value in the output alone”. Without out=, those positions would hold whatever the new memory contained, which is why NumPy has warned since 2.4 when where= comes without out=. np.errstate only silences the warnings; the inf and nan stay in the data, which np.isfinite shows.

out= is useful without where= as well: np.add(total, sales, out=total) writes into total and makes no new array, which is what total += sales does. dtype= sets the type the loop computes in: the uint8 values 200 and 100 doubled gave 144 and 200, and with dtype=np.int16 they gave 400 and 200.

Methods turn a ufunc into a reduction, a running total or a table

Every ufunc that takes two inputs and gives one output has five methods. NumPy’s guide lists them as four that work like reductions, reduce, accumulate, reduceat and outer, and one for in-place work, at. The run tries each on numbers small enough to check by hand:

"""The ufunc methods reduce, accumulate, outer, reduceat and at, on numbers small enough to check."""
import numpy as np

print("np.add.reduce([1, 2, 3, 4]):", np.add.reduce([1, 2, 3, 4]), "| np.add.accumulate:", np.add.accumulate([1, 2, 3, 4]))
print("np.multiply.accumulate(np.arange(1, 6)):", np.multiply.accumulate(np.arange(1, 6)))
print("np.maximum.accumulate([1, 3, 2, 5, 4]):", np.maximum.accumulate([1, 3, 2, 5, 4]))
print("np.multiply.outer([1, 2, 3], [10, 20]):")
print(np.multiply.outer([1, 2, 3], [10, 20]))
print("np.add.reduceat(np.arange(8), [0, 4, 6]):", np.add.reduceat(np.arange(8), [0, 4, 6]))
v = np.zeros(4, dtype=np.int64)
np.add.at(v, [0, 0, 2], 5)
print("np.add.at(v, [0, 0, 2], 5):", v)

monthly = np.full(120, 0.08 / 12)  # 8 % a year, added monthly, for 10 years
balance = 100_000 * np.multiply.accumulate(1 + monthly)
print(f"100,000 after 1 year: {balance[11]:,.2f} | after 10 years: {balance[-1]:,.2f}")
print(f"closed form, 100,000 * (1 + 0.08 / 12) ** 120: {100_000 * (1 + 0.08 / 12) ** 120:,.2f}")
monthly[60:] = 0.06 / 12  # the rate falls to 6 % after five years
balance = 100_000 * np.multiply.accumulate(1 + monthly)
print(f"with 6 % from year 6: {balance[-1]:,.2f}")

np.add.reduce([1, 2, 3, 4]): 10 | np.add.accumulate: [ 1  3  6 10]
np.multiply.accumulate(np.arange(1, 6)): [  1   2   6  24 120]
np.maximum.accumulate([1, 3, 2, 5, 4]): [1 3 3 5 5]
np.multiply.outer([1, 2, 3], [10, 20]):
[[10 20]
 [20 40]
 [30 60]]
np.add.reduceat(np.arange(8), [0, 4, 6]): [ 6  9 13]
np.add.at(v, [0, 0, 2], 5): [10  0  5  0]
100,000 after 1 year: 108,299.95 | after 10 years: 221,964.02
closed form, 100,000 * (1 + 0.08 / 12) ** 120: 221,964.02
with 6 % from year 6: 200,957.86

reduce combines all the values with the ufunc, so np.add.reduce is a sum, 10. accumulate keeps every intermediate result: the running sum [1 3 6 10], the running product [1 2 6 24 120] and the running maximum [1 3 3 5 5]. outer applies the ufunc to every pair, a 3 by 2 table. reduceat reduces the segments that start at the positions you give, so [0, 4, 6] sums 0 to 3, 4 to 5 and 6 to 7, [6 9 13]. at is the unbuffered operation that the lesson on boolean masks and fancy indexing used for repeated indices, and here it added 5 twice at position 0. A ufunc with one input has none of this: np.sin.reduce raises a ValueError.

accumulate is a growth model in one line. A deposit of 100,000 at 8 % a year, with interest added monthly, grows by a factor of 1 + 0.08 / 12 each month, and the running product of 120 factors gives every month’s balance: 108,299.95 after one year and 221,964.02 after ten, equal to the closed form. Because each month has a factor of its own, a rate that falls to 6 % after five years needs only changed factors, and gives 200,957.86. The model below starts at the run’s values; change the rate or the years and watch the balance:

Its balance after 10 years, 221,964, is the program’s last value rounded to a whole number.

np.vectorize is a Python loop with an array interface

np.vectorize turns a Python function of one number into one that accepts arrays. The run times it against a list comprehension and against the same fee rule written with ufuncs, on 100,000 seeded distances:

"""np.vectorize is a Python loop behind an array interface: timed against a ufunc expression."""
import platform
import timeit

import numpy as np

rng = np.random.default_rng(7)  # seeded: every run prices the same 100,000 orders
km = rng.uniform(0, 10, 100_000)  # delivery distances in km


def fee(d):
    """The delivery fee of one order: 40 up to 2 km, then 12 a km more."""
    return 40.0 if d <= 2 else 40.0 + 12.0 * (d - 2)


vfee = np.vectorize(fee)
ways = [
    ("list comprehension", lambda: np.array([fee(d) for d in km.tolist()])),
    ("np.vectorize(fee)", lambda: vfee(km)),
    ("np.where(...)", lambda: np.where(km <= 2, 40.0, 40.0 + 12.0 * (km - 2))),
]
results = {}
for name, fn in ways:
    results[name] = fn()
    best = min(timeit.repeat(fn, number=1, repeat=7))
    print(f"{name:<19} {best * 1000:8.3f} ms")
same = all(np.array_equal(results["np.where(...)"], r) for r in results.values())
print("the three give the same fees:", same, "| total", round(float(results["np.where(...)"].sum()), 2))
print("np.vectorize(fee) is a ufunc:", isinstance(vfee, np.ufunc))
print(f"CPython {platform.python_version()}, NumPy {np.__version__}, {platform.machine()}")

list comprehension    13.348 ms
np.vectorize(fee)     15.161 ms
np.where(...)          0.176 ms
the three give the same fees: True | total 7844441.05
np.vectorize(fee) is a ufunc: False
CPython 3.14.8, NumPy 2.5.3, arm64

The list comprehension and np.vectorize both call fee once per order, 100,000 Python calls, and took about 13 and 15 ms; the np.where expression computed both branches for every order with ufuncs and kept one per element, in about 0.18 ms, about 86 times faster on this run. All three gave the same fees. NumPy’s reference is plain about it: vectorize is provided “primarily for convenience, not for performance”, and its implementation “is essentially a for loop”.

It has a second trap: it learns the output dtype from the first result.

"""np.vectorize takes its output dtype from the first result, unless otypes says otherwise."""
import numpy as np


def discount(price):
    """10 % off above 100, nothing off up to 100."""
    return 0 if price <= 100 else price * 0.1


prices = np.array([80.0, 155.0, 99.5, 129.0])
print("np.vectorize(discount):                ", np.vectorize(discount)(prices))
print("np.vectorize(discount, otypes=[float]):", np.vectorize(discount, otypes=[float])(prices))
print("np.where(prices > 100, prices * 0.1, 0.0):", np.where(prices > 100, prices * 0.1, 0.0))

np.vectorize(discount):                 [ 0 15  0 12]
np.vectorize(discount, otypes=[float]): [ 0.  15.5  0.  12.9]
np.where(prices > 100, prices * 0.1, 0.0): [ 0.  15.5  0.  12.9]

discount(80.0) returned the integer 0, so every later result was cast to an integer and 15.5 became 15. Passing otypes=[float] fixes it, and the np.where version never had the problem.

Practise: divide safely and keep running maxima

Two small functions in plain Python reproduce what the ufunc arguments and methods of this lesson do: a division that skips zero divisors, as out= with where= does, and running maxima that either spread NaN or skip it. The sample tests run here:

def safe_ratio(a, b, fill):
    """Return a[i] / b[i] where b[i] != 0 and float(fill) elsewhere, as floats (np.divide with out= and where=)."""
    return [x / y if y > 0 else fill for x, y in zip(a, b)]

import math

from ratio import safe_ratio


def test_the_lessons_shops():
    """gives the ratios the lesson's run printed, with the fill where there were no orders"""
    assert safe_ratio([3.0, 5.0, 0.0, 4.0], [2.0, 0.0, 0.0, 8.0], -1.0) == [1.5, -1.0, -1.0, 0.5]


def test_floats_everywhere():
    """returns floats, also for integer inputs and an integer fill"""
    result = safe_ratio([7, 9, 1], [2, 3, 0], 0)
    assert result == [3.5, 3.0, 0.0]
    assert all(isinstance(x, float) for x in result)


def test_negative_divisors_and_negative_zero():
    """divides by negative numbers, and treats -0.0 as zero"""
    assert safe_ratio([6.0, 1.0], [-3.0, -0.0], 9.0) == [-2.0, 9.0]


def test_nan_divisor():
    """divides by NaN, because NaN != 0, and so gives NaN"""
    assert math.isnan(safe_ratio([1.0], [math.nan], 0.0)[0])


def test_lengths_must_match():
    """refuses lists of different lengths"""
    try:
        safe_ratio([1.0, 2.0], [1.0], 0.0)
    except ValueError:
        return
    assert False, "lists of different lengths must raise ValueError"

def running_max(x):
    """Return the running maximum of x, as np.maximum.accumulate gives it (NaN spreads)."""
    return [max(x[: i + 1]) for i in range(len(x))]


def running_fmax(x):
    """Return the running maximum of x, as np.fmax.accumulate gives it (NaN is skipped)."""
    return [max(x[: i + 1]) for i in range(len(x))]

import math

from runmax import running_fmax, running_max

nan = math.nan


def same(got, expected):
    """True when the lists agree position by position, counting NaN as equal to NaN."""
    return len(got) == len(expected) and all((math.isnan(g) and math.isnan(e)) or g == e for g, e in zip(got, expected))


def test_the_lessons_maxima():
    """gives the running maximum the lesson's run printed"""
    assert running_max([1, 3, 2, 5, 4]) == [1, 3, 3, 5, 5]
    assert running_fmax([1, 3, 2, 5, 4]) == [1, 3, 3, 5, 5]
    assert running_max([-3, -1, -2]) == [-3, -1, -1]


def test_nan_spreads_in_running_max():
    """makes every value NaN from the first NaN on, as np.maximum does"""
    assert same(running_max([2.0, nan, 1.0, 5.0]), [2.0, nan, nan, nan])
    assert same(running_max([nan, 1.0]), [nan, nan])


def test_nan_is_skipped_in_running_fmax():
    """skips NaN when there is a number to keep, as np.fmax does"""
    assert same(running_fmax([2.0, nan, 1.0, 5.0]), [2.0, 2.0, 2.0, 5.0])
    assert same(running_fmax([nan, 1.0, nan, 3.0]), [nan, 1.0, 1.0, 3.0])


def test_empty():
    """gives an empty list for no values"""
    assert running_max([]) == []
    assert running_fmax([]) == []

The questions that follow cover all five ideas, and the first of them comes from the idea your answers so far show as weakest:

Interview questions

Warm-up (fresher): What is a ufunc, and why is np.add(a, b) faster than a Python loop over the elements?

A ufunc, or universal function, is a NumPy function that works element by element on whole arrays, with broadcasting and type casting built in; np.add, np.sqrt and np.maximum are ufuncs, and operators such as + call them. Each ufunc holds compiled loops, one for each combination of input types it supports, and picks one from the dtypes of its inputs, for example the float64 loop for two float64 arrays. That loop then runs over every element in compiled code. A Python loop instead creates and inspects a Python object for every value and runs the interpreter for every addition, which is why, in the lesson’s timed run, a delivery fee written with ufuncs ran about 86 times faster than one Python call per order. Ufuncs also take out= to reuse memory and where= to skip positions, and have methods such as reduce and accumulate.

Two interview questions apply these rules to real data, a conversion report that divides by zero and image arithmetic that wraps round; their first steps are on the module’s question page:

Key takeaways

• Operators on arrays call ufuncs, and a ufunc picks a compiled loop from its inputs’ dtypes, so the dtypes decide the result type and whether a result wraps round, as uint8 200 + 200 = 144 did.
• // and % follow floor rules with the remainder’s sign from the divisor, integers to negative integer powers raise, and np.round rounds halves to even with a fast method that can differ from Python’s round.
• where= computes only the positions you choose and out= keeps every other one, which divides safely by zero; np.errstate only silences warnings.
• reduce, accumulate, outer, reduceat and at turn a two-input ufunc into a total, a running result, a table, segment totals or an unbuffered update.
• np.vectorize calls the Python function once per element and takes its dtype from the first result, so rewrite the rule with ufuncs and np.where when speed matters.