在 numpy 数组中应用自定义函数

Question

我有一个清单，

mylist=np.array([120,3,10,33,5,54,2,23,599,801])

还有一个功能：

def getSum(n): 
    n=n**2
    sum = 0
    while (n != 0): 

        sum = sum + int(n % 10) 
        n = int(n/10) 
    if sum <20:
        return True
    return False

我正在尝试将我的函数应用于 mylist 并仅检索那些索引为真。

我的预期输出是。

[120, 3, 10, 33, 5, 54, 2, 23, 801]

我可以像 list(filter(getSum,mylist)) 那样做，如何在 numpy 中使用它。

尝试np.where 没有产生预期的输出。

Answer 1

如果要检查数字的总和是否为> 20，这里有一个纯numpy 解决方案（here 可以找到如何分解其数字中的整数）：

import numpy as np


mylist=np.array([120,3,10,33,5,54,2,23,599,801])

mylist = mylist**2
max_digits = np.ceil(np.max(np.log10(mylist)))  # max number of digits in mylist
digits = mylist//(10**np.arange(max_digits)[:, None])%10  # matrix of digits
digitsum = np.sum(digits, axis=0)  # array of sums
mask = digitsum < 20
mask
# array([True, True, True, True, True, True, True, True, False, True])

更新：速度比较

@hpaulj 对（几乎）所有提议的解决方案进行了很好的时间比较。
获胜者是 filter，输入为纯 list，而我的 pure numpy 解决方案表现不佳。
无论如何，如果我们针对更广泛的输入对它们进行测试，情况就会发生变化。
这是使用来自@NicoSchlömer 的perflot 执行的测试。
对于 100 多个元素的输入，所有解决方案都是等效的，而纯 numpy 更快：

Answer 2

函数和测试数组：

In [22]: def getSum(n):  
    ...:     n=n**2 
    ...:     sum = 0 
    ...:     while (n != 0):  
    ...:  
    ...:         sum = sum + int(n % 10)  
    ...:         n = int(n/10)  
    ...:     if sum <20: 
    ...:         return True 
    ...:     return False 
    ...:                                                                        
In [23]: mylist=np.array([120,3,10,33,5,54,2,23,599,801])                       

您的filter 解决方案：

In [51]: list(filter(getSum, mylist))                                           
Out[51]: [120, 3, 10, 33, 5, 54, 2, 23, 801]

还有一个采样时间：

In [52]: timeit list(filter(getSum, mylist))                                    
32.8 µs ± 185 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

由于这会返回一个列表并进行迭代，如果mylist 是一个列表而不是一个数组，它应该会更快：

In [53]: %%timeit alist=mylist.tolist() 
    ...: list(filter(getSum, alist))                                                                        
18.4 µs ± 378 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

替代方案

您提议使用 np.vectorize：

In [56]: f = np.vectorize(getSum); mylist[f(mylist)]                            
Out[56]: array([120,   3,  10,  33,   5,  54,   2,  23, 801])
In [57]: timeit f = np.vectorize(getSum); mylist[f(mylist)]                     
63.4 µs ± 151 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [58]: timeit mylist[f(mylist)]                                               
57.6 µs ± 920 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

哎呀！即使我们从计时循环中删除 f 创建，这也会慢一些。 vectorize 很漂亮，但不保证速度。

我发现 frompyfunc 比 np.vectorize 快（尽管它们是相关的）：

In [59]: g = np.frompyfunc(getSum, 1,1)                                         
In [60]: g(mylist)                                                              
Out[60]: 
array([True, True, True, True, True, True, True, True, False, True],
      dtype=object)

结果是object dtype，在这种情况下必须转换为bool：

In [63]: timeit mylist[g(mylist).astype(bool)]                                  
25.5 µs ± 233 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

这比您的 filter 更好 - 但仅适用于数组，而不是列表。

@Saandeep 提出了一个列表理解：

In [65]: timeit mylist[[getSum(i) for i in mylist]]                             
40.7 µs ± 1.21 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

这比你的 filter 慢一点。

使用列表推导的更快方法是：

 [i for i in mylist if getSum(i)]

这与您的 filter 相同 - 对于数组和列表版本（我丢失了我正在计时的会话）。

纯numpy

​​>

@lante 提出了一个纯粹的numpy 解决方案，聪明但有点晦涩难懂。我还没有弄清楚逻辑：

def lante(mylist):
    max_digits = np.ceil(np.max(np.log10(mylist)))  # max number of digits in mylist
    digits = mylist//(10**np.arange(max_digits)[:, None])%10  # matrix of digits
    digitsum = np.sum(digits, axis=0)  # array of sums
    mask = digitsum > 20
    return mask

不幸的是不是速度恶魔：

In [69]: timeit mylist[~lante(mylist)]                                          
58.9 µs ± 757 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

我没有安装numba，所以无法为@jezrael's 解决方案计时。

所以你原来的filter 是一个很好的解决方案，特别是如果你从一个列表而不是一个数组开始。尤其是在考虑转换时间时，一个好的 Python 列表解决方案通常比 numpy 更好。

对于一个大的例子来说，时间可能会有所不同，但我预计不会有任何不安。

Answer 3

使用list comprehension，np.vectorize 的基本概念是文档中的 for 循环（也不会提高您的性能）：

mylist[[getSum(i) for i in mylist]]

array([120,   3,  10,  33,   5,  54,   2,  23, 801])

Answer 4

我认为有循环，所以最好使用numba：

from numba import jit
@jit(nopython=True)
def get_vals(arr):
    out = np.zeros(arr.shape[0], dtype=bool)
    for i, n in enumerate(arr):

        n=n**2
        sum1 = 0
        while (n != 0): 
            sum1 = sum1 + int(n % 10) 
            n = int(n/10) 
        if sum1 <20:
            out[i] = True
    return arr[out]

print(get_vals(mylist))

Answer 5

vec=np.vectorize(getSum)
mylist[vec(mylist)]
out[]:
array([120,   3,  10,  33,   5,  54,   2,  23, 801])