在Numba中计算numpy数组中非零值的数量

Question

非常简单。我正在尝试使用Numba编译的NumPy jit来计算数组中非零值的数量（njit()）。以下方法在Numba中不被允许：

1. a[a != 0].size
2. np.count_nonzero(a)
3. len(a[a != 0])
4. len(a) - len(a[a == 0])

如果有更快、更Pythonic和更优雅的方法，我不想使用for循环。

对于那些想要看到完整代码示例的评论者...

import numpy as np
from numba import njit

@njit()
def n_nonzero(a):
    return a[a != 0].size

- SARose

---

请展示至少一段实际的、完整的代码，包括import语句、装饰器和样本测试框架。 - Mark Setchell

[a 中的 l 的筛选列表(filter((0).ne, l))] - Cleve Green

@MarkSetchell 当然可以... - SARose

Answer 1

你也可以考虑，嗯，统计非零值：

import numba as nb

@nb.njit()
def count_loop(a):
    s = 0
    for i in a:
        if i != 0:
            s += 1
    return s

我知道这看起来很奇怪，但请耐心听我说：

import numpy as np
import numba as nb

@nb.njit()
def count_loop(a):
    s = 0
    for i in a:
        if i != 0:
            s += 1
    return s

@nb.njit()
def count_len_nonzero(a):
    return len(np.nonzero(a)[0])

@nb.njit()
def count_sum_neq_zero(a):
    return (a != 0).sum()

np.random.seed(100)
a = np.random.randint(0, 3, 1000000000, dtype=np.uint8)
c = np.count_nonzero(a)
assert count_len_nonzero(a) == c
assert count_sum_neq_zero(a) == c
assert count_loop(a) == c

%timeit count_len_nonzero(a)
# 5.94 s ± 141 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit count_sum_neq_zero(a)
# 848 ms ± 80.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit count_loop(a)
# 189 ms ± 4.41 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

事实上，它比np.count_nonzero更快，后者由于某些原因可能会变得非常慢：

%timeit np.count_nonzero(a)
# 4.36 s ± 69.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

- jdehesa

---

当Numba看到一个可以优化的循环时，它确实表现出色。+1 - MSeifert

Answer 2

如果你需要快速处理大数组，甚至可以使用numbas prange 来并行处理计数（对于小数组来说，由于并行处理的开销，它会更慢）。

import numpy as np
from numba import njit, prange

@njit(parallel=True)
def parallel_nonzero_count(arr):
    flattened = arr.ravel()
    sum_ = 0
    for i in prange(flattened.size):
        sum_ += flattened[i] != 0
    return sum_

请注意，当您使用numba时，通常希望编写循环，因为这正是numba擅长优化的地方。

实际上，我使用我的Python模块simple_benchmark对其进行了时间测量。

重现代码：

import numpy as np
from numba import njit, prange

@njit
def n_nonzero(a):
    return a[a != 0].size

@njit
def count_non_zero(np_arr):
    return len(np.nonzero(np_arr)[0])

@njit() 
def methodB(a): 
    return (a!=0).sum()

@njit(parallel=True)
def parallel_nonzero_count(arr):
    flattened = arr.ravel()
    sum_ = 0
    for i in prange(flattened.size):
        sum_ += flattened[i] != 0
    return sum_

@njit()
def count_loop(a):
    s = 0
    for i in a:
        if i != 0:
            s += 1
    return s

from simple_benchmark import benchmark

args = {}
for exp in range(2, 20):
    size = 2**exp
    arr = np.random.random(size)
    arr[arr < 0.3] = 0.0
    args[size] = arr

b = benchmark(
    funcs=(n_nonzero, count_non_zero, methodB, np.count_nonzero, parallel_nonzero_count, count_loop),
    arguments=args,
    argument_name='array size',
    warmups=(n_nonzero, count_non_zero, methodB, np.count_nonzero, parallel_nonzero_count, count_loop)
)

- MSeifert

---

在并行循环迭代之间共享 sum_ 是安全的吗？（我不太了解并行化 Numba 的保证） - jdehesa

是的，numba有一些可以安全并行化的缩减操作。求和和乘法就是其中之一。这是因为numba意识到它可以使用sum_ = 0并行处理每个过程，然后在每个过程完成后将它们相加。我还针对np.count_nonzero进行了一致性检查。 - MSeifert

Answer 3

你可以使用 np.nonzero 来获取它的长度：

@njit
def count_non_zero(np_arr):
    return len(np.nonzero(np_arr)[0])

count_non_zero(np.array([0,1,0,1]))
# 2

- Chris

---

那个 [0] 似乎就是做到这件事的东西。非常感谢！ - SARose

有趣的是，np.nonzero 在第二次迭代时使用 np.count_nonzero（在 c-api 级别上）来确定它将填充的数组大小。我认为使用 numba 的整个意义就在于能够毫不顾虑地进行迭代 :) - hpaulj

Answer 4

我不确定是否犯了错误，但这似乎快了6倍：

# Make something worth checking
a=np.random.randint(0,3,1000000000,dtype=np.uint8)  

In [41]: @njit() 
    ...: def methodA(a): 
    ...:     return len(np.nonzero(a)[0])                                                                                           

# Call and check result
In [42]: methodA(a)                                                                                 
Out[42]: 666644445

In [43]: %timeit methodA(a)                                                                         
4.65 s ± 28.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [44]: @njit() 
    ...: def methodB(a): 
    ...:     return (a!=0).sum()                                                                                         

# Call and check result    
In [45]: methodB(a)                                                                                 
Out[45]: 666644445

In [46]: %timeit methodB(a)                                                                         
724 ms ± 14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

- Mark Setchell