Как использовать np.apply_along_axis() и когда это уместно по сравнению с векторизованными операциями?

Что делает np.apply_along_axis()

np.apply_along_axis(func1d, axis, arr) применяет функцию func1d к каждому одномерному срезу массива arr вдоль заданной оси. NumPy итерирует по всем осям кроме указанной, нарезает 1D-вектор и передаёт его в функцию. Результаты собираются в новый массив.

Ключевой момент: внутри apply_along_axis — обычный цикл Python, замаскированный под удобный API. Это означает медленную работу на больших массивах по сравнению с настоящей векторизацией через ufunc или broadcasting.

Пример: нормализация каждого столбца своим диапазоном

import numpy as np

data = np.array([
    [10.0, 200.0, 0.5],
    [20.0, 400.0, 1.5],
    [30.0, 100.0, 2.5],
    [40.0, 300.0, 0.0],
])
print("shape:", data.shape, "dtype:", data.dtype)  # (4, 3) float64

def minmax_scale(col: np.ndarray) -> np.ndarray:
    """Нормализует 1D-вектор в диапазон [0, 1]."""
    lo, hi = col.min(), col.max()
    if hi == lo:
        return np.zeros_like(col)
    return (col - lo) / (hi - lo)

# axis=0 → func получает каждый столбец (вектор длиной 4)
result = np.apply_along_axis(minmax_scale, axis=0, arr=data)
print(result)
# [[0.   0.333 0.5 ]
#  [0.333 1.   1.  ]
#  [0.667 0.   ...]
#  ...]
print("result shape:", result.shape)  # (4, 3) — сохранилась форма

axis=0 vs axis=1

# axis=0 → 1D срез вдоль строк (по каждому столбцу)
# data.shape = (4, 3) → func вызывается 3 раза, каждый раз с массивом (4,)

# axis=1 → 1D срез вдоль столбцов (по каждой строке)
# func вызывается 4 раза, каждый раз с массивом (3,)
row_stats = np.apply_along_axis(lambda row: [row.mean(), row.std()], axis=1, arr=data)
print(row_stats.shape)  # (4, 2) — func вернула вектор длиной 2

Когда apply_along_axis уместен

• Функция сложна и не раскладывается на ufunc-операции: сортировка, квантили, running statistics с состоянием.
• Прототип или разовая задача, где читаемость важнее скорости.
• Функция возвращает вектор переменной длины — тогда только Python-цикл и apply_along_axis.

Когда vectorized быстрее

# apply_along_axis — медленно: вызов Python на каждый срез
result_slow = np.apply_along_axis(minmax_scale, axis=0, arr=data)

# Векторизованный вариант — быстро: broadcasting без Python-цикла
lo = data.min(axis=0)       # shape (3,)
hi = data.max(axis=0)       # shape (3,)
result_fast = (data - lo) / (hi - lo)  # broadcasting: (4,3) - (3,) → (4,3)

# Для больших массивов разница в 10–100x
import timeit
N = 1000
big = np.random.rand(N, N)
print(timeit.timeit(lambda: np.apply_along_axis(lambda c: (c - c.min()) / (c.max() - c.min() + 1e-9), 0, big), number=5))
print(timeit.timeit(lambda: (big - big.min(0)) / (big.max(0) - big.min(0) + 1e-9), number=5))
# apply_along_axis: ~1.5s, broadcasting: ~0.005s

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

• Производительность: apply_along_axis — скрытый Python for-loop. На массиве (1000, 1000) с простой операцией в 100–300 раз медленнее broadcasting. Всегда профилируйте перед выбором.
• Форма результата зависит от возвращаемого значения func: если func возвращает скаляр — результат имеет меньше осей; если вектор — форма зависит от его длины. Неожиданный shape ломает downstream-код молча.
• Первый вызов для определения shape: NumPy вызывает func дважды на первом срезе (один раз для пробного запуска). Функции с побочными эффектами (logging, изменение состояния) отработают лишний раз.
• Не работает с разной длиной выходов между срезами: если func возвращает разное количество элементов для разных срезов — ValueError или ragged array. Используйте list comprehension с np.object_.
• Пустой массив: при arr.shape[axis] == 0 func не вызывается, но результат не всегда интуитивен — проверяйте shape явно.
• NaN-пропуски: функция получает NaN в срезе и может вернуть NaN или некорректный результат без предупреждения. Добавляйте np.isnan-проверки внутри func при работе с реальными данными.
• GIL и параллелизм: apply_along_axis не освобождает GIL, поэтому threading не ускорит его. Для параллельного применения — joblib.Parallel с явным list comprehension или переход на Numba/Cython.
• Альтернатива np.vectorize: np.vectorize тоже является Python-циклом, но работает поэлементно, а не срезами. Не путайте с настоящей векторизацией через ufunc.

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

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