Numpy：在任意方向的数组切片上应用函数
Numpy: apply function over arbitrarily oriented slices of array
有没有办法将函数应用于多维数组的一般切片？
例如，给定一个表示彩色视频的 4D 输入数组 [frame, y, x, color_channel]，我们想将 2D 图像过滤器应用于 [y, x] 中的所有 2D 切片。
这可以表示为如下的一般操作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)
只是为了测试我过去的观察，apply_along_axis 很方便，但不快（呃）：
定义一个简单的一维函数：
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（即传递大小为 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)
这更接近 apply 中使用的逻辑，但出于某种原因仍然快了很多。 apply，更一般，必须有更多的开销，包括测试评估以确定 return 数组的大小。
相同的逻辑可以应用于需要二维数组的 foo。
这是一个解决方案：
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.
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)
有没有办法将函数应用于多维数组的一般切片？
例如，给定一个表示彩色视频的 4D 输入数组 [frame, y, x, color_channel]，我们想将 2D 图像过滤器应用于 [y, x] 中的所有 2D 切片。
这可以表示为如下的一般操作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)
只是为了测试我过去的观察，apply_along_axis 很方便，但不快（呃）：
定义一个简单的一维函数：
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（即传递大小为 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)
这更接近 apply 中使用的逻辑，但出于某种原因仍然快了很多。 apply，更一般，必须有更多的开销，包括测试评估以确定 return 数组的大小。
相同的逻辑可以应用于需要二维数组的 foo。
这是一个解决方案：
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.
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)