高效地将numpy数组与其自身按元素进行比较 [英] Comparing numpy array with itself by element efficiently

问题描述

我正在执行大量这些计算:

I am performing a large number of these calculations:

A == A[np.newaxis].T

其中A是一个密集的numpy数组，通常具有相同的值.

where A is a dense numpy array which frequently has common values.

出于基准测试目的，我们可以使用:

For benchmarking purposes we can use:

n = 30000
A = np.random.randint(0, 1000, n)
A == A[np.newaxis].T

执行此计算时，我遇到了内存问题.我相信这是因为输出不是采用更有效的bitarray或np.packedbits格式.第二个需要考虑的问题是，由于生成的布尔数组是对称的，因此我们执行的比较次数是必要的两倍.

When I perform this calculation, I run into memory issues. I believe this is because the output isn't in more efficient bitarray or np.packedbits format. A secondary concern is we are performing twice as many comparisons as necessary, since the resulting Boolean array is symmetric.

我的问题是:

1. 是否可以在不牺牲速度的情况下以更高效的内存方式生成布尔值numpy数组输出?我所知道的选项是bitarray和np.packedbits，但是我只知道在创建大型布尔数组之后如何应用它们.
2. 我们可以利用计算的对称性将处理的比较数量减半，而又不牺牲速度吗?

我将需要能够执行&和|对布尔数组输出的操作.我已经尝试过bitarray，对于这些按位运算，它是超快速的.但是打包np.ndarray-> bitarray然后解压缩bitarray-> np.ndarray是很慢的.

I will need to be able to perform & and | operations on Boolean arrays output. I have tried bitarray, which is super-fast for these bitwise operations. But it is slow to pack np.ndarray -> bitarray and then unpack bitarray -> np.ndarray.

推荐答案

这里是带有numba的一个，用于为我们提供一个NumPy布尔数组作为输出-

Here's one with numba to give us a NumPy boolean array as output -

from numba import njit

@njit
def numba_app1(idx, n, s, out):
    for i,j in zip(idx[:-1],idx[1:]):
        s0 = s[i:j]
        c = 0
        for p1 in s0[c:]:
            for p2 in s0[c+1:]:
                out[p1,p2] = 1
                out[p2,p1] = 1
            c += 1
    return out

def app1(A):
    s = A.argsort()
    b = A[s]
    n = len(A)
    idx = np.flatnonzero(np.r_[True,b[1:] != b[:-1],True])
    out = np.zeros((n,n),dtype=bool)
    numba_app1(idx, n, s, out)
    out.ravel()[::out.shape[1]+1] = 1
    return out

时间-

In [287]: np.random.seed(0)
     ...: n = 30000
     ...: A = np.random.randint(0, 1000, n)

# Original soln
In [288]: %timeit A == A[np.newaxis].T
1 loop, best of 3: 317 ms per loop

# @Daniel F's soln-1 that skips assigning lower diagonal in output
In [289]: %timeit sparse_outer_eq(A)
1 loop, best of 3: 450 ms per loop

# @Daniel F's soln-2 (complete one)
In [291]: %timeit sparse_outer_eq(A)
1 loop, best of 3: 634 ms per loop

# Solution from this post
In [292]: %timeit app1(A)
10 loops, best of 3: 66.9 ms per loop

这篇关于高效地将numpy数组与其自身按元素进行比较的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！