Numpy: apply function over arbitrarily oriented slices of array

Question

Is there a way to apply a function over general slices of a multidimensional array?

As an example, given a 4D input array representing a color video [frame, y, x, color_channel], we would like to apply a 2D image filter to all 2D slices in [y, x].

Can this be expressed as a general operation apply_to_slices as in the following?

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)

There is a `np.apply_along_axis` that takes a function and a n-dimensional array. It iterates over n-1 dimensions, passing 1d slices to the function. It can be convenient, but isn't any faster than doing the iterations explicitly. — hpaulj, Apr 18 '20 at 03:03
Does this answer your question? [python how to put argument to function with numpy aply\_along\_axis](https://stackoverflow.com/questions/28452716/python-how-to-put-argument-to-function-with-numpy-aply-along-axis) — Joe, Apr 18 '20 at 05:30
np.apply_along_axis() along supports 1-D slices, whereas I'd like more general slices. It would be convenient/pythonic to avoid having to manually create for-loop iterations. — Hugues, Apr 18 '20 at 06:16
`For loops` are very pythonic. Hiding them in your `non_general_awkward_solution` is also pythonic. I mention `apply_along_axis` to show the closest thing in `numpy` to your ideal. Feel free to study it, and extend it to work in your case. Based on my timings I don't think it's worth the effort. — hpaulj, Apr 18 '20 at 07:22

hpaulj · Answer 1 · 2020-04-18T07:25:53.900

Just to test my past observations that apply_along_axis is convenient, but not fast(er):

define a simple 1d function:

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)

Make a multidimensional array (>2d):

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

apply foo along the first (ie pass size 2 arrays 60 times):

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)

do the equivalent with a reshape to (2,60) and transpose to (60,2). Iterate on the first axis:

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)

that is significantly faster than apply.

Do the same but on the last axis, so I don't need the transpose (only 24 iterations):

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)

and the apply equivalent:

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)

and 3 level nested iteration (a bit slower than the reshape, but still faster than 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)

Using ndindex to generate the (i,j,k) indexing tuple:

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)

This is closer to the logic used in apply, though for some reason still quite a bit faster. apply, being more general, must have more overhead, including a test evaluation to determine the size of the return array.

The same logic can be applied to a foo that requires a 2d array.

This answer makes the great point that the current _implementation_ of `np.apply_along_axis()` seems inefficient -- and possibly should be replaced by the alternates proposed. However it doesn't answer the question. Is the suggestion to use a combination of `np.moveaxis()` and `reshape()` to allow a 1D `np.array([...for...])` followed by `reshape()` and `np.moveaxis()`? Demonstrating this in the placeholder `apply_to_slices()` would be awesome. — Hugues, Apr 18 '20 at 15:02

Hugues · Accepted Answer · 2020-04-19T15:27:03.977

Here is a solution:

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

Verification:

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)

Numpy: apply function over arbitrarily oriented slices of array

2 Answers2