Ufunc за минуту: @vectorize и @stencil

1. @vectorize: своя универсальная функция

Универсальные функции (ufunc) — это np.exp, np.sqrt и прочие: они принимают массив любой формы, применяются поэлементно, дружат с broadcast. @vectorize позволяет добавить в этот клуб свою формулу — писав её для одного элемента:

from numba import vectorize

@vectorize(['float64(float64)'])          # сигнатура: из float64 в float64
def expit(x):                              # логистическая сигмоида
    return 1.0 / (1.0 + np.exp(-x))       # написана для ОДНОГО числа

expit(a)          # ...а работает с массивом любой формы
expit(0.0)        # и со скаляром

Чем это лучше numpy-выражения 1 / (1 + np.exp(-a))? Numpy выполняет его в три прохода: посчитать −a во временный массив, exp — во второй, деление — в третий; десять миллионов элементов трижды пройдут через всю память. Ufunc от @vectorize делает один проход без временных массивов, а target='parallel' раздаёт его по ядрам (тот же map из P02):

Бонусы клуба ufunc достаются бесплатно: broadcast (наша функция двух аргументов сама сообразит, что второй — скаляр), методы .reduce и .accumulate, работа с любыми формами массивов. Третий target — 'cuda' — оставим до P04: это та же строка кода, выполняющаяся на видеокарте.

@guvectorize — обобщение на случай «вход — строка, выход — строка»: сигнатура формы записывается строкой вида '(n),()->(n)' («из массива длины n и скаляра — массив длины n»), и функция автоматически применяется к каждой строке многомерного входа. Пример — скользящее среднее (проверено: от [0,1,2,3,4,5] с окном 3 получается [0, 0.5, 1, 2, 3, 4]):

from numba import guvectorize

@guvectorize(['(float64[:], int64, float64[:])'], '(n),()->(n)')
def moving_avg(row, w, out):          # out приходит готовым — заполняем его
    acc = 0.0
    for i in range(row.size):
        acc += row[i]
        if i >= w:
            acc -= row[i - w]
        out[i] = acc / min(i + 1, w)

2. @stencil: вычисления по окрестности

Второй вездесущий паттерн — «новое значение точки считается из её соседей»: размытие картинок, разностные схемы для уравнений физики, клеточные автоматы. Писать его циклами муторно из-за границ (у крайних точек нет соседей слева) и легко ошибиться в индексах. @stencil описывает только суть — ядро в относительных координатах:

from numba import stencil, njit

@stencil
def blur_kernel(m):
    # m[0, 0] — текущая точка; m[-1, 0] — сосед сверху, m[0, 1] — справа…
    return (m[-1, 0] + m[1, 0] + m[0, -1] + m[0, 1] + m[0, 0]) / 5.0

@njit                                     # стенсил вызывается из njit-функции
def blur(m):
    return blur_kernel(m)

Numba сама генерирует двойной цикл по всем точкам, а граничные точки, которым не хватает соседей, по умолчанию получают ноль (границу можно настроить). Стенсил внутри @njit(parallel=True) автоматически параллелится — это map по точкам.

Сквозной пример — игра «Жизнь» Конвея: клетка living=1/dead=0, у каждой 8 соседей; живая выживает при 2–3 соседях, мёртвая оживает ровно при 3. Весь автомат — один стенсил:

@stencil
def life_kernel(m):
    neigh = (m[-1,-1] + m[-1,0] + m[-1,1] +
             m[0,-1]            + m[0,1] +
             m[1,-1]  + m[1,0] + m[1,1])
    return np.uint8((neigh == 3) or (m[0,0] and neigh == 2))

@njit
def life_step(m):
    return life_kernel(m)

Проверка корректности — классическая фигура глайдер: пять клеток, которые за четыре поколения воспроизводят себя со сдвигом на одну клетку по диагонали. Тест запущен: через 4 шага поле в точности равно исходному, сдвинутому np.roll на (1, 1). Скорость — шаг поля 2000×2000 за 2 мс. Заметьте важное свойство: «Жизнь» — это scan по времени (поколение зависит от предыдущего — не распараллелишь), но map по пространству (все клетки поколения считаются одновременно). Разложение задачи на «последовательное время × параллельное пространство» — стандартная судьба физических симуляций.

3. Тренажёр: «Жизнь» вживую

Контрольные вопросы

• Почему один проход @vectorize быстрее трёх numpy-операций, если numpy тоже скомпилирован?
  Numpy вычисляет выражение по операциям: каждая создаёт временный массив и гонит все данные через память заново — у 1/(1+exp(−a)) три прохода и два временных массива. Ufunc от @vectorize вычисляет всю формулу за один проход без временных массивов; память — узкое место, поэтому один проход выигрывает, а target='parallel' добавляет все ядра.
• Что обещает автор функции, помечая её @vectorize, и что нарушает это обещание?
  Что функция описывает вычисление одного элемента из скалярных аргументов — чистый map без взгляда на соседей и без операций над массивами. Срезы, суммы по массиву, обращения к соседним элементам — нарушение; для строк есть @guvectorize, для окрестностей — @stencil.
• Как @stencil справляется с граничными точками и с ошибкой «обновление на месте»?
  Границы, где соседей не хватает, по умолчанию заполняются нулём (поведение настраивается). Результат стенсил всегда пишет в новый массив — двойная буферизация встроена, поэтому «сосед слева из нового поколения» невозможен. При ручной реализации циклом об обеих проблемах нужно помнить самому.
• В каком смысле «Жизнь» параллелится, а в каком — нет?
  По пространству — map: все клетки нового поколения считаются одновременно и независимо (из старого поля). По времени — рекуррентная цепочка: поколение t+1 требует полностью готового поколения t, шаги строго последовательны. Разложение «параллельное пространство × последовательное время» типично для симуляций.
• Чем глайдер удобен как тест корректности?
  Его эволюция известна точно: через 4 поколения фигура воспроизводит себя со сдвигом ровно на (1, 1). Сравнение поля с np.roll исходного — бинарный тест без допусков, который ловит и ошибку правила, и ошибку границ, и обновление на месте.

Источники

1. Creating NumPy universal functions (@vectorize, @guvectorize) // Numba Documentation : [сайт]. — URL: https://numba.readthedocs.io/en/latest/user/vectorize.html (дата обращения: 09.07.2026).
2. Using @stencil // Numba Documentation : [сайт]. — URL: https://numba.readthedocs.io/en/latest/user/stencil.html (дата обращения: 09.07.2026).
3. Conway's Game of Life // Wikipedia : [сайт]. — URL: https://en.wikipedia.org/wiki/Conway%27s_Game_of_Life (дата обращения: 09.07.2026).