Универсальные функции (ufunc): поэлементная математика

Урок знакомит с ufunc — скомпилированными поэлементными функциями, на которых держится вся векторная математика NumPy.

ufunc (universal function) — функция NumPy, применяющая операцию к каждому элементу массива в скомпилированном цикле; именно ufunc делают арифметику и математику векторными.

Что такое ufunc и почему их так много

Когда вы пишете a + b или np.sqrt(a), под капотом вызывается ufunc — заранее скомпилированная функция, которая проходит по всем элементам массива в цикле на C. Арифметические операторы (+ - * / ** %) — это синтаксический сахар над ufunc np.add, np.subtract, np.multiply и так далее. Помимо арифметики, NumPy предоставляет сотни математических ufunc: np.sqrt, np.exp, np.log, np.sin, np.cos, np.abs, np.maximum, np.sign и многие другие.

Ключевое свойство: ufunc применяется к каждому элементу независимо и возвращает массив той же формы. Это и есть векторизация в чистом виде.

import numpy as np
a = np.array([1.0, 4.0, 9.0, 16.0])

print(np.sqrt(a))         # корень из каждого элемента
print(np.exp(np.array([0, 1, 2])))   # экспонента
print(a + 1)              # это np.add(a, 1)
print(np.maximum(a, 5))   # поэлементный максимум с 5

Вывод:

[1. 2. 3. 4.]
[1.         2.71828183 7.3890561 ]
[ 2.  5. 10. 17.]
[ 5.  5.  9. 16.]

Скалярные функции math против ufunc

Встроенный модуль math работает только с одним числом: math.sqrt([1,4,9]) выбросит ошибку. Чтобы применить его к коллекции, нужен цикл. ufunc снимает эту необходимость. Сравним подходы на чистом Python — сначала «скалярный» способ через цикл и math:

import math

a = [1.0, 4.0, 9.0, 16.0]

# Скалярная функция — только цикл
roots = [math.sqrt(x) for x in a]
print(roots)

Вывод:

[1.0, 2.0, 3.0, 4.0]

В NumPy цикл не нужен — np.sqrt(a) сразу обрабатывает весь массив. Концептуально ufunc — это «функция, поднятая на уровень массива»: она знает, как пройтись по всем элементам сама. Напишем такой «подъём» вручную, чтобы прочувствовать механику:

import math

def vectorize(func):
    # Превращаем скалярную функцию в «ufunc-подобную»
    def wrapped(arr):
        return [func(x) for x in arr]
    return wrapped

vsqrt = vectorize(math.sqrt)
vexp = vectorize(math.exp)

print(vsqrt([1, 4, 9, 16]))
print([round(x, 4) for x in vexp([0, 1, 2])])

Вывод:

[1.0, 2.0, 3.0, 4.0]
[1.0, 2.7183, 7.3891]

NumPy делает то же самое, но цикл выполняется в C, поэтому работает на порядки быстрее нашего Python-обёртки.

Унарные, бинарные и где они берут broadcasting

ufunc бывают унарными (один аргумент: np.sqrt, np.exp, np.abs) и бинарными (два аргумента: np.add, np.maximum, np.power). Бинарные ufunc — это в точности то, что стоит за арифметическими операторами, и именно они приносят broadcasting: когда вы складываете массивы разных форм, работает бинарная ufunc np.add, которая умеет выравнивать формы по правилам broadcasting (отдельный урок далее). Унарные ufunc проще — они просто применяют функцию к каждому элементу, сохраняя форму. Понимание, что вся поэлементная математика и арифметика — это вызовы ufunc, даёт цельную картину: одни и те же правила (broadcasting, повышение типов, параметры out и where) действуют единообразно, будь то сложение, возведение в степень или синус. Не нужно запоминать поведение каждой функции отдельно — достаточно знать общие правила ufunc.

Скорость ufunc и почему цикл Python проигрывает

Когда вы вызываете np.sqrt(a) для массива из миллиона чисел, происходит ровно один переход из Python в скомпилированный код, после чего весь цикл по элементам идёт на C без возврата в интерпретатор. Сравните с [math.sqrt(x) for x in a], где на каждый из миллиона элементов интерпретатор выполняет вызов функции, разбор байт-кода, упаковку результата в объект. Разница — не в алгоритме извлечения корня (он одинаков), а в накладных расходах вокруг него, которые ufunc устраняет почти полностью. Поэтому правило простое: для любой математики над массивами ищите готовую ufunc вместо цикла или list comprehension. Почти для всего, что есть в модуле math, в NumPy найдётся векторный аналог с тем же именем.

Параметр out=: запись результата без лишней копии

По умолчанию ufunc создаёт новый массив под результат. Но если у вас уже есть готовый буфер (например, заранее выделенный np.empty или сам входной массив), можно записать результат прямо в него через параметр out=. Это избавляет от выделения временной памяти — важная оптимизация в горячих циклах и при работе с большими массивами.

import numpy as np
a = np.array([1.0, 4.0, 9.0, 16.0])
buf = np.empty_like(a)

np.sqrt(a, out=buf)       # результат пишется в buf, нового массива нет
print(buf)

np.multiply(a, 2, out=a)  # можно писать и в сам входной массив (на месте)
print(a)

Вывод:

[1. 2. 3. 4.]
[ 2.  8. 18. 32.]

Запись в out=a (тот же входной массив) эквивалентна операции на месте и экономит память. Это особенно ценно в конвейерах вычислений, где иначе создавались бы десятки временных массивов.

Запись в сам входной массив через out=a — это идиома операции «на месте», знакомая по предыдущему разделу. Она особенно полезна в длинных конвейерах преобразований, где иначе на каждом шаге создавался бы новый временный массив. Например, нормализовать большой массив изображения «на месте» — поделить на максимум и записать обратно в тот же буфер — значит избежать удвоения памяти. Цена та же, что у любых операций на месте: если массив является view на другие данные, вы измените и их, поэтому применяйте out= с тем же вниманием к разделяемой памяти.

Параметр where=: условное применение

Параметр where= принимает булеву маску и применяет ufunc только там, где маска True. В остальных позициях значение остаётся прежним (берётся из out, если он задан, иначе не определено). Это позволяет избирательно преобразовать часть массива без отдельной индексации.

import numpy as np
a = np.array([-4.0, 9.0, -1.0, 16.0])
result = a.copy()

# Корень только из положительных, отрицательные оставляем как есть
np.sqrt(a, out=result, where=(a > 0))
print(result)

Вывод:

[-4.  3. -1.  4.]

Здесь корень посчитан только для 9 и 16 (стали 3 и 4), а отрицательные -4 и -1 остались нетронутыми. Без where= попытка взять корень из отрицательного дала бы предупреждение и nan.

Где ufunc встречаются незаметно

Стоит осознать, насколько ufunc вездесущи, — тогда станет ясно, почему этому уроку отведено столько места. Каждый оператор арифметики и сравнения, каждая математическая функция NumPy, всё поэлементное преобразование данных — это ufunc. Когда в следующих разделах мы будем нормировать данные, считать расстояния, применять функции активации, фильтровать по условиям — под капотом всякий раз работают ufunc с их broadcasting, повышением типов и параметрами out/where. Иными словами, ufunc — это атом, из которого собрано почти всё вычислительное в NumPy. Выучив один раз общие правила их поведения, вы получаете предсказуемость: незнакомая функция, оказавшаяся ufunc, будет вести себя по тем же законам, что и знакомое сложение. Это резко снижает объём того, что нужно держать в памяти.

Бинарные ufunc и их методы

Бинарные ufunc (принимающие два аргумента, как np.add) имеют полезные методы. reduce сворачивает массив одной операцией (так np.add.reduce даёт сумму), accumulate возвращает накопленные значения (префиксные суммы). На практике чаще пользуются готовыми sum, cumsum, но полезно понимать, что это частные случаи reduce/accumulate.

import numpy as np
a = np.array([1, 2, 3, 4])

print(np.add.reduce(a))        # 1+2+3+4 = 10 (то же, что a.sum())
print(np.add.accumulate(a))    # [1, 3, 6, 10] — накопленная сумма
print(np.multiply.reduce(a))   # 1*2*3*4 = 24 (то же, что a.prod())

Вывод:

10
[ 1  3  6 10]
24

Идею accumulate легко воспроизвести на Python — это и есть «бегущая» свёртка:

def accumulate_add(arr):
    out = []
    total = 0
    for x in arr:
        total += x
        out.append(total)
    return out

print(accumulate_add([1, 2, 3, 4]))

Вывод:

[1, 3, 6, 10]

Подводные камни

• Применять math к массиву. math.sqrt(array) падает — используйте np.sqrt. Функции из math только скалярные.
• Игнорировать nan и предупреждения. Корень/логарифм от недопустимого даёт nan/-inf с предупреждением, а не исключение. Это легко пропустить.
• Лишние временные массивы. Цепочки вроде np.exp(np.sqrt(a) + 1) создают промежуточные буферы; в горячем коде используйте out=.
• Путать where= ufunc и np.where. Параметр where= — это маска применения; функция np.where(cond, x, y) — поэлементный выбор. Разные вещи.

Особые значения: nan и inf как результат ufunc

ufunc, в отличие от функций math, обычно не выбрасывают исключение на недопустимом входе — они возвращают специальные значения с плавающей точкой. Корень из отрицательного даёт nan («не число»), логарифм нуля — -inf, деление на ноль — inf или nan. При этом печатается предупреждение, но выполнение не прерывается. Это сделано осознанно: при обработке миллионов элементов останавливать всё из-за одного плохого значения непрактично; удобнее получить nan в проблемной позиции и разобраться с ним отдельно. Но отсюда вытекает дисциплина: после операций, которые могут породить nan/inf (логарифмы, деления, корни), проверяйте результат через np.isfinite и решайте, что делать с особыми значениями — отфильтровать, заменить, или это сигнал об ошибке в данных. Параметр where=, который мы только что разобрали, как раз позволяет заранее не применять опасную операцию там, где она дала бы nan.

Лучшие практики

• Используйте ufunc вместо циклов и math для любой поэлементной математики.
• В критичном по памяти коде задавайте out=, чтобы переиспользовать буферы.
• Для условных преобразований выбирайте между where= (применить ufunc выборочно) и булевой индексацией.
• Помните: операторы + - * / — это ufunc, поэтому к ним применимы те же правила broadcasting и типов.

Подводя черту: ufunc превращают NumPy из «контейнера чисел» в «вычислительную машину». Они дают единый, предсказуемый способ применять любую математику ко всему массиву разом, с общими правилами broadcasting, типов и записи результата. Освоив их, вы получаете язык, на котором описывается всё остальное в библиотеке.

Итог

• ufunc — скомпилированные поэлементные функции; арифметические операторы тоже ufunc.
• В отличие от math, ufunc применяется ко всему массиву сразу, без цикла Python.
• out= пишет результат в готовый буфер, экономя память; where= применяет операцию по маске.
• Методы reduce/accumulate сворачивают и накапливают; sum/cumsum — их частные случаи.