NumPy применяются вдоль N-пробелов
У меня есть 4d массив, и я хотел бы применить функцию к каждому 2d срезу, взятому путем итерации по двум последним измерениям. А именно, примените f(2d_array) к (x,y,0,0) и f(2d_array) к (x,y,0,1) и т. Д. И т. Д. Моя функция работает с массивом на месте, поэтому размеры будут будет таким же, но общее решение вернет массив формы (x',y',w,z), где w и z - два последних измерения исходного массива.
Это, очевидно, может быть обобщено на срезы mD по массиву nD.
Есть ли встроенная функциональность, которая делает это?
3 ответа
"Базовая" модель применения вдоль оси состоит в том, чтобы выполнять итерацию по одной оси и передавать другую в вашу функцию:
In [197]: def foo(x): # return same size
...: return x*2
...: np.array([foo(x) for x in np.arange(12).reshape(3,4)])
...:
Out[197]:
array([[ 0, 2, 4, 6],
[ 8, 10, 12, 14],
[16, 18, 20, 22]])
In [198]: def foo(x):
...: return x.sum() # return one less dim
...: np.array([foo(x) for x in np.arange(12).reshape(3,4)])
...:
Out[198]: array([ 6, 22, 38])
In [199]: def foo(x):
...: return x.sum(keepdims=True) # condense the dim
...: np.array([foo(x) for x in np.arange(12).reshape(3,4)])
...:
Out[199]:
array([[ 6],
[22],
[38]])
Ваша проблема 4d может быть обработана, чтобы соответствовать этому.
In [200]: arr_4d = np.arange(24).reshape(2,3,2,2)
In [201]: arr_2d = arr_4d.reshape(6,4).T
In [202]: res = np.array([foo(x) for x in arr_2d])
In [203]: res
Out[203]:
array([[60],
[66],
[72],
[78]])
In [204]: res.reshape(2,2)
Out[204]:
array([[60, 66],
[72, 78]])
что эквивалентно выполнению:
In [205]: arr_4d[:,:,0,0].sum()
Out[205]: 60
In [206]: foo(arr_4d[:,:,0,0].ravel())
Out[206]: array([60])
apply_along_axis
требуется функция, которая принимает массив 1d, но может применяться следующим образом:
In [209]: np.apply_along_axis(foo,0,arr_4d.reshape(6,2,2))
Out[209]:
array([[[60, 66],
[72, 78]]])
foo
может изменить свой ввод в 2d и передать его функции, которая принимает 2d. apply_along_index
использования np.ndindex
генерировать индексы для осей итерации.
In [212]: list(np.ndindex(2,2))
Out[212]: [(0, 0), (0, 1), (1, 0), (1, 1)]
np.vectorize
обычно работает с функцией, которая принимает скаляр. Но последние версии имеют signature
параметр, который, я считаю, может быть использован для работы с вашим делом. Может потребоваться транспонировать входные данные, чтобы выполнить итерации по первым двум осям, передавая последние две функции. Смотрите мой ответ на /questions/17378736/kak-primenit-funktsiyu-kotoraya-vozvraschaet-vektor-k-kazhdomu-elementu-massiva-numpy-i-poluchit-massiv-s-bolshej-razmernostyu/17378749#17378749.
Ни один из этих подходов не дает преимущества в скорости.
Без изменения формы или замены, я могу перебрать с помощью ndindex
,
Определите функцию, которая ожидает 2d ввода:
def foo2(x):
return x.sum(axis=1, keepdims=True) # 2d
Индексный итератор за последние 2 месяца arr_4d
:
In [260]: idx = np.ndindex(arr_4d.shape[-2:])
Сделайте тест calc, чтобы определить форму возврата. vectorize
а также apply...
сделать этот вид теста.
In [261]: r1 = foo2(arr_4d[:,:,0,0]).shape
In [262]: r1
Out[262]: (2, 1)
Массив результатов:
In [263]: res = np.zeros(r1+arr_4d.shape[-2:])
In [264]: res.shape
Out[264]: (2, 1, 2, 2)
Теперь повторяем:
In [265]: for i,j in idx:
...: res[...,i,j] = foo2(arr_4d[...,i,j])
...:
In [266]: res
Out[266]:
array([[[[ 12., 15.],
[ 18., 21.]]],
[[[ 48., 51.],
[ 54., 57.]]]])
Я думаю, вы ищете что-то вроде numpy.apply_over_axes
в сочетании с for
цикл для итерации других переменных осей.
Я катал свой собственный. Мне было бы интересно узнать, есть ли какие-либо различия в производительности между этим и методом @ hpaulj, и есть ли основания полагать, что написание пользовательского модуля c принесет существенное улучшение. Конечно, метод @ hpaulj носит более общий характер, поскольку это связано с тем, что мне нужно просто выполнить операцию с массивом на месте.
def apply_along_space(f, np_array, axes):
# apply the function f on each subspace given by iterating over the axes listed in axes, e.g. axes=(0,2)
for slic in itertools.product(*map(lambda ax: range(np_array.shape[ax]) if ax in axes else [slice(None,None,None)], range(len(np_array.shape)))):
f(np_array[slic])
return np_array