是否有方法将函数应用于多维数组的一般切片?
例如,给定表示彩色视频[frame, y, x, color_channel]
的4D输入阵列,我们希望将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
很方便,但不快(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(即通过大小为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)
其明显快于CCD_ 5。
做同样的事情,但在最后一个轴上,所以我不需要转置(只有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
更通用,必须有更多的开销,包括确定返回数组大小的测试评估。
相同的逻辑可以应用于需要2d阵列的CCD_ 11。
这里有一个解决方案:
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)