Numpy将函数应用于数组中的每个项目

Question

假设我有一个二维数组。如何将函数应用于数组中的每个项目并用返回替换该项目？另外，函数的返回值会是一个元组，所以数组会变成 3d。

这是记住的代码。

def filter_func(item):
    if 0 <= item < 1:
        return (1, 0, 1)
    elif 1 <= item < 2:
        return (2, 1, 1)
    elif 2 <= item < 3:
        return (5, 1, 4)
    else:
        return (4, 4, 4)

myarray = np.array([[2.5, 1.3], [0.4, -1.0]])

# Apply the function to an array

print(myarray)

# Should be array([[[5, 1, 4],
#                   [2, 1, 1]],
#                  [[1, 0, 1],
#                   [4, 4, 4]]])

任何想法我该怎么做？一种方法是使用np.array(list(map(filter_func, myarray.reshape((12,))))).reshape((2, 2, 3))，但这很慢，尤其是当我需要在形状数组 (1024, 1024) 上执行此操作时。

我也看到人们使用 np.vectorize，但它以某种方式结束为 (array([[5, 2], [1, 4]]), array([[1, 1], [0, 4]]), array([[4, 1], [1, 4]])). 然后它的形状为 (3, 2, 2)。

Answer 1

无需更改函数中的任何内容。

只需将函数的矢量化版本应用于数组 和 stack 结果：

np.stack(np.vectorize(filter_func)(myarray), axis=2)

结果是：

array([[[5, 1, 4],
        [2, 1, 1]],

       [[1, 0, 1],
        [4, 4, 4]]])

Answer 2

你可以使用这个函数，通过矢量化实现

def func(arr):
    
    elements = np.array([
        [1, 0, 1],
        [2, 1, 1],
        [5, 1, 4],
        [4, 4, 4],
    ])
    
    arr  = arr.astype(int)
    mask = (arr != 0) & (arr != 1) & (arr != 2)

    arr[mask] = -1
    
    return elements[arr]

由于形状不匹配，您将无法重写数组 但你可以覆盖变量myarray

myarray = func(myarray)
myarray

>>>   [[[5, 1, 4],
        [2, 1, 1]],

       [[1, 0, 1],
        [4, 4, 4]]]

Answer 3

您的列表地图：

In [4]: np.array(list(map(filter_func, myarray.reshape((4,))))).reshape((2, 2, 3))                   
Out[4]: 
array([[[5, 1, 4],
        [2, 1, 1]],

       [[1, 0, 1],
        [4, 4, 4]]])

使用嵌套列表推导的变体：

In [5]: np.array([[filter_func(j) for j in row] for row in myarray])                                 
Out[5]: 
array([[[5, 1, 4],
        [2, 1, 1]],

       [[1, 0, 1],
        [4, 4, 4]]])

使用vectorize，结果是函数返回的每个元素对应一个数组。

In [6]: np.vectorize(filter_func)(myarray)                                                           
Out[6]: 
(array([[5, 2],
        [1, 4]]),
 array([[1, 1],
        [0, 4]]),
 array([[4, 1],
        [1, 4]]))

正如@Vladi 所示，这些可以与stack 组合（或np.array 后跟转置）：

In [7]: np.stack(np.vectorize(filter_func)(myarray),2)                                               
Out[7]: 
array([[[5, 1, 4],
        [2, 1, 1]],

       [[1, 0, 1],
        [4, 4, 4]]])

您的列表图是最快的。我从来没有发现vectorize 更快：

In [8]: timeit np.array(list(map(filter_func, myarray.reshape((4,))))).reshape((2, 2, 3))            
17.2 µs ± 47.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [9]: timeit np.array([[filter_func(j) for j in row] for row in myarray])                          
20.5 µs ± 78.1 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [10]: timeit np.stack(np.vectorize(filter_func)(myarray),2)                                       
75.2 µs ± 297 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

将np.vectorize(filter_func) 排除在计时循环之外会有所帮助。

frompyfunc 类似于 vectorize，但返回对象 dtype。它通常更快：

In [29]: timeit np.stack(np.frompyfunc(filter_func, 1,3)(myarray),2).astype(int)                     
28.7 µs ± 125 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

一般来说，如果您有一个只接受标量输入的函数，则很难比简单迭代做得更好。 vectorize/frompyfunc 不要改进。正如@Hammad 所演示的那样，numpy 的最佳使用需要重写函数以直接处理数组。

尽管有了这个小例子，即使是这个正确的numpy 解决方案也不会更快。我希望它会更好地扩展：

In [32]: timeit func(myarray)                                                                        
25 µs ± 60.8 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)