apply_along_axis

Коротко

• np.apply_along_axis(func, axis, a) вызывает func для каждого одномерного среза вдоль axis: при axis=1 — для каждой строки.
• Функция может вернуть число или массив — результаты собираются в новый массив.
• Внутри — обычный цикл Python, поэтому это не быстрее цикла. Выигрыш только в краткости.
• Если функция собрана из агрегаций NumPy (max, mean, sort…), вызовите их сразу с axis — это в разы быстрее.
• apply_along_axis оправдан для функций без параметра axis: bincount, interp, свои функции с ветвлениями.

Синтаксис

np.apply_along_axis(func1d, axis, arr, *args, **kwargs)
np.apply_along_axis(f, 1, a)      # f(строка) для каждой строки
np.apply_along_axis(f, 0, a)      # f(столбец) для каждого столбца

Параметры

Примеры

# оценки 3 учеников по 5 работам
scores = np.array([[4, 5, 3, 5, 4],
                   [3, 3, 4, 2, 5],
                   [5, 5, 5, 4, 5]])

def spread(row):
    return row.max() - row.min()

np.apply_along_axis(spread, 1, scores)       # функция получает каждую строку

array([2, 3, 1])

def top2(row):
    return np.sort(row)[-2:]                 # функция возвращает массив

np.apply_along_axis(top2, 1, scores)

array([[5, 5],
       [4, 5],
       [5, 5]])

def mode(row):
    return np.bincount(row).argmax()         # у bincount нет параметра axis

np.apply_along_axis(mode, 1, scores)         # самая частая оценка каждого ученика

array([4, 3, 5])

np.apply_along_axis(np.median, 0, scores)    # по столбцам — функция получает каждый столбец

array([4., 5., 4., 4., 5.])

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

scores.max(axis=1) - scores.min(axis=1)      # то же без apply_along_axis

array([2, 3, 1])