Numpy:在任意方向的数组切片上应用函数



是否有方法将函数应用于多维数组的一般切片?

例如,给定表示彩色视频[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)

最新更新