Numpy: применить функцию к произвольно ориентированным фрагментам массива

Question

Есть ли способ применить функцию к общим слоям многомерного массива?

В качестве примера, учитывая входной массив 4D, представляющий цветное видео [frame, y, x, color_channel], мы хотели бы применить изображение 2D фильтровать все 2D-срезы в [y, x].

Может ли это быть выражено как общая операция apply_to_slices, как показано ниже?

video = np.random.rand(2, 3, 4, 3)  # 2 frames, each 3x4 pixels with 3 channels.

def filter_2d(image):  # example of simple 2D blur filter
  import scipy.signal
  kernel = np.ones((3, 3)) / 9.0
  return scipy.signal.convolve2d(image, kernel, mode='same', boundary='symm')


def apply_to_slices(func, array, axes):
  """Apply 'func' to each slice of 'array', where a slice spans 'axes'.

  Args:
    func: function expecting an array of rank len(axes) and returning a
      modified array of the same dimensions.
    array: input of arbitrary shape.
    axes: integer sequence specifying the slice orientation.
  """
  pass


def non_general_awkward_solution(func, video):
  new_video = np.empty_like(video)
  for frame in range(video.shape[0]):
    for channel in range(video.shape[3]):
      new_video[frame, ..., channel] = func(video[frame, ..., channel])
  return new_video


# new_video = apply_to_slices(filter_2d, video, axes=(1, 2))
new_video = non_general_awkward_solution(filter_2d, video)
print(video)
print(new_video)

Answer 1

Вот решение:

def apply_to_slices(func, a, axes):
  """Apply 'func' to each slice of array 'a', where a slice spans 'axes'.

  Args:
    func: function expecting an array of rank len(axes) and returning a
      modified array of the same shape.
    a: input array of arbitrary shape.
    axes: integer sequence specifying the slice orientation.
  """
  # The approach is to move the slice axes to the end of the array, reshape to
  # a 1-D array of slices, apply the user function to each slice, reshape back
  # to an outer array of slices, and finally move the slice axes back to their
  # original locations.  https://stackoverflow.com/a/61297133/
  assert len(axes) <= a.ndim
  outer_ndim = a.ndim - len(axes)
  a = np.moveaxis(a, axes, range(outer_ndim, a.ndim))
  outer_shape = a.shape[:outer_ndim]
  slice_shape = a.shape[outer_ndim:]
  a = a.reshape((-1,) + slice_shape)
  a = np.array([func(a_slice) for a_slice in a])
  a = a.reshape(outer_shape + slice_shape)
  a = np.moveaxis(a, range(outer_ndim, a.ndim), axes)
  return a

Проверка:

new_video = apply_to_slices(filter_2d, video, axes=(1, 2))
new_video2 = non_general_awkward_solution(filter_2d, video)
assert np.all(new_video == new_video2)

Answer 2

Просто для проверки моих прошлых наблюдений, что apply_along_axis удобно, но не быстро (er):

определяет простую 1d функцию:

In [683]: def foo(X): 
     ...:     assert(X.ndim==1) 
     ...:     return X 
     ...:      
     ...:                                                                                              
In [684]: foo(np.arange(3))                                                                            
Out[684]: array([0, 1, 2])
In [685]: foo(np.ones((3,2)))                                                                          
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)

Создайте многомерный массив (> 2d):

In [686]: arr = np.ones((2,3,4,5))                                                                     

применить foo вдоль первого (ie массивов с размером прохода 2 60 раз):

In [687]: np.apply_along_axis(foo, 0, arr);                                                            
In [688]: timeit np.apply_along_axis(foo, 0, arr);                                                     
293 µs ± 406 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

сделать эквивалент с изменением формы на (2,60) и перенести в (60,2). Выполните итерацию по первой оси:

In [689]: np.array([foo(x) for x in arr.reshape(2,-1).transpose()]).shape                              
Out[689]: (60, 2)
In [690]: np.array([foo(x) for x in arr.reshape(2,-1).transpose()]).transpose().reshape(arr.shape);    
In [691]: timeit np.array([foo(x) for x in arr.reshape(2,-1).transpose()]).transpose().reshape(arr.shape);                                                                                          
49.4 µs ± 20.4 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

, что значительно быстрее, чем apply.

Сделайте то же самое, но по последней оси, поэтому мне не требуется транспонирование (только 24 итерации):

In [692]: np.array([foo(x) for x in arr.reshape(-1,5)]).reshape(arr.shape);                            
In [693]: timeit np.array([foo(x) for x in arr.reshape(-1,5)]).reshape(arr.shape);                     
23.6 µs ± 23.2 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

и эквивалент применения:

In [694]: timeit np.apply_along_axis(foo, 3, arr);                                                     
156 µs ± 85.1 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

и 3-уровневая итерация (немного медленнее, чем изменение формы, но все же быстрее, чем apply:

In [695]: np.array([foo(arr[i,j,k,:]) for i in range(2) for j in range(3) for k in range(4)]);                                                                                              
In [696]: timeit np.array([foo(arr[i,j,k,:]) for i in range(2) for j in range(3) for k in range(4)]);  
32.5 µs ± 864 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Использование ndindex для генерации индексного кортежа (i,j,k):

In [701]: timeit np.array([foo(arr[i,j,k]) for i,j,k in np.ndindex((2,3,4))]).reshape(arr.shape);      
87.3 µs ± 218 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Это ближе к логике c, используемой в apply, хотя по некоторым причинам все еще довольно немного быстрее. apply, будучи более общим, должен иметь больше накладных расходов, включая тестовую оценку для определения размера возвращаемого массива.

Те же логики c могут быть применены к foo, который требуется 2d массив.