NumPyMiddleTechnical

Что такое универсальные функции (ufuncs) в NumPy и почему они эффективны?

Ufuncs — скомпилированные C-функции, применяемые поэлементно к массивам. Они эффективны за счёт векторизованных SIMD-инструкций, работы без Python overhead для каждого элемента и поддержки broadcasting.

Универсальные функции (ufuncs) в NumPy

Ufunc (Universal Function) — это функция, которая работает поэлементно над массивами NumPy и реализована на C. Они являются основой векторизованных вычислений: вместо Python-цикла по элементам выполняется один вызов скомпилированного кода с SIMD-оптимизациями.

Встроенные ufuncs: примеры использования

import numpy as np

arr = np.array([1.0, 4.0, 9.0, 16.0])

# Математические ufuncs
print(np.sqrt(arr))    # [1. 2. 3. 4.]
print(np.exp(arr))     # [2.718... 54.6... 8103... 8886111...]
print(np.log(arr))     # [0. 1.386... 2.197... 2.772...]

# Тригонометрические
angles = np.array([0, np.pi/6, np.pi/4, np.pi/2])
print(np.sin(angles))  # [0. 0.5 0.707... 1.]

# Бинарные ufuncs (два операнда)
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
print(np.add(a, b))       # [5 7 9]
print(np.multiply(a, b))  # [4 10 18]
print(np.maximum(a, b))   # [4 5 6]
print(np.power(a, b))     # [1 32 729]

Почему ufuncs быстрее Python-циклов

import numpy as np
import time

N = 10_000_000
arr = np.random.rand(N)

# Python-цикл: медленно
start = time.perf_counter()
result_py = [x ** 0.5 for x in arr]
print(f"Python loop: {time.perf_counter() - start:.3f}s")

# NumPy ufunc: быстро
start = time.perf_counter()
result_np = np.sqrt(arr)
print(f"NumPy ufunc: {time.perf_counter() - start:.3f}s")
# NumPy в 50-200x быстрее!

# Причины скорости:
# 1. Компилированный C-код без Python GIL overhead
# 2. SIMD инструкции (AVX/SSE) — обрабатывают 4-8 float за такт
# 3. Оптимальная работа с кэшем CPU
# 4. Нет создания Python-объектов для каждого элемента

Broadcasting в ufuncs

import numpy as np

# Ufuncs автоматически применяют broadcasting
a = np.array([[1, 2, 3], [4, 5, 6]])  # Shape (2, 3)
b = np.array([10, 20, 30])             # Shape (3,)

print(np.add(a, b))  # Broadcasting: b растягивается до (2, 3)
# [[11 22 33]
#  [14 25 36]]

# Outer product через метод .outer()
print(np.multiply.outer([1, 2, 3], [10, 20]))
# [[10 20]
#  [20 40]
#  [30 60]]

Методы ufunc: reduce, accumulate, reduceat

import numpy as np

arr = np.array([1, 2, 3, 4, 5])

# reduce: применяет функцию пошагово ко всем элементам
print(np.add.reduce(arr))       # 15 (1+2+3+4+5)
print(np.multiply.reduce(arr))  # 120 (1*2*3*4*5)

# accumulate: промежуточные результаты reduce
print(np.add.accumulate(arr))   # [1 3 6 10 15] — аналог cumsum

# reduceat: reduce по сегментам
print(np.add.reduceat(arr, [0, 2, 4]))  # [3, 7, 5] — суммы [0:2], [2:4], [4:]

# at: применяет к выбранным индексам (с повторением!)
arr2 = np.zeros(5, dtype=int)
np.add.at(arr2, [0, 1, 1, 2], 1)  # Индекс 1 добавляется дважды
print(arr2)  # [1 2 1 0 0]

Создание кастомного ufunc через np.frompyfunc и np.vectorize

import numpy as np

# np.frompyfunc: обёртка Python-функции в ufunc
# Возвращает object dtype — медленнее настоящих ufuncs
def clamp(x, lo, hi):
    return max(lo, min(hi, x))

clamp_ufunc = np.frompyfunc(clamp, 3, 1)  # 3 входа, 1 выход
result = clamp_ufunc(np.array([1, 5, 10, 15]), 3, 8)
print(result)  # [3 5 8 8] (dtype=object)

# Конвертация в нужный dtype
result_float = result.astype(np.float64)

# Настоящие ufuncs создаются через Cython, C-extension или Numba
import numba

@numba.vectorize(['float64(float64, float64)'], nopython=True)
def fast_custom(x, y):
    return x * x + y * y  # Компилируется в LLVM, настоящий ufunc

arr1 = np.array([1.0, 2.0, 3.0])
arr2 = np.array([4.0, 5.0, 6.0])
print(fast_custom(arr1, arr2))  # [17. 29. 45.]

out параметр: запись без выделения памяти

import numpy as np

a = np.array([1.0, 2.0, 3.0])
b = np.array([4.0, 5.0, 6.0])
result = np.empty(3)  # Предварительно выделенный буфер

np.add(a, b, out=result)  # Нет выделения новой памяти
print(result)  # [5. 7. 9.]

# Полезно в tight loops для снижения давления на GC

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

  • np.frompyfunc возвращает object dtype: обёрнутые Python-функции работают с object-массивами, что медленнее и требует явного .astype().
  • np.vectorize — не настоящий ufunc: это удобная обёртка, но производительность примерно как у Python-цикла. Подробнее — в вопросе про np.vectorize.
  • Нет поддержки Python exceptions в ufuncs: ошибки в C-коде могут проявляться как нарушения памяти, а не как Python исключения.
  • Broadcasting создаёт временные массивы: при сложных broadcasting-операциях NumPy может создавать промежуточные массивы, увеличивая потребление памяти.
  • np.add.at медленнее чем np.add: метод .at() не векторизован — он обрабатывает повторяющиеся индексы корректно, но без SIMD оптимизаций.
  • Точность floating-point в reduce: np.add.reduce() суммирует слева направо. Для больших массивов float32 это накапливает ошибку. Используйте np.sum() с алгоритмом попарного суммирования.

Common mistakes

  • Объяснять ufuncs vectorized kernels только синтаксисом без shape, dtype, состояния или режима выполнения.
  • Игнорировать leakage, воспроизводимость, пустые входы и скрытые копии данных.
  • Не проверять production-симптомы: latency, память, ретраи, дрейф качества и несовпадение версий.

What the interviewer is testing

  • Может ли связать ufuncs vectorized kernels с реальным контрактом входов и выходов.
  • Упоминает ли тесты, метрики, reproducibility и диагностику ошибок.
  • Видит ли различие между demo-кодом в ноутбуке и production-пайплайном.

Sources

Related topics