применить пользовательскую функцию в массиве 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

Я пытаюсь применить свою функцию к списку и получить только те индексы, которые соответствуют действительности.

мой ожидаемый результат.

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

Я могу сделать это, как list(filter(getSum,mylist)), как использовать это в numpy.

пробовал np.where не выдает ожидаемый результат.

Answer 1

Если вы хотите проверить, равны ли суммы цифр > 20, здесь чисто numpy решение ( здесь может найти способ разложения целого числа в его цифры):

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 не работает должным образом.
В любом случае, если мы проверяем их на более широком диапазоне входных данных, все меняется.
Вот тест, выполненный с perflot от @ NicoSchlömer.
Для ввода более 100 элементов все решения эквивалентны, в то время как чистая цифра быстрее:

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)

Oops! это немного медленнее, даже если мы удалим создание 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)

результат - объект 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 one.

С большим примером время может отличаться, но я не ожидаю каких-либо расстройств.

Answer 3

Я думаю, что есть петли, поэтому лучше использовать 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 4

Используя list comprehension, базовая концепция np.vectorize предназначена для цикла из документов (также не улучшает вашу производительность):

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

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

Answer 5

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