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
Другие вопросы по тегам