利用 Numpy 数组中的单调性

Question

我有一个大的 Numpy 数组，如果元素为正，我想按元素应用一些函数，否则用 -inf 替换。该数组基本上是两个向量的外部差异（如果不清楚，请参见示例），因此它具有特定的结构：每一行从正数开始，并在某些时候变为负数，例如，当我减去 [1,. ..,10] 来自 [3,4,5] 我得到矩阵

2  1  0 -1 -2 -3 ...
3  2  1  0 -1 -2 ...
4  3  2  1  0 -1 ...

（当然，在实际应用中向量并不是均匀的space整数）

我想得到

f(2)  f(1)  f(0)  -inf -inf -inf ...
f(3)  f(2)  f(1)  f(0) -inf -inf ...
f(4)  f(3)  f(2)  f(1) f(0) -inf ...

有一个简单的解决方案：

import numpy as np

x = np.arange(2000,dtype=np.float64)
y = np.arange(4000,dtype=np.float64)

M = x[:,None] - y[None,:]

i_pos = M>=0

M[i_pos] = np.sqrt(M[i_pos])
M[~i_pos] = -np.inf

它有效，但看起来效率很低，因为寻找正元素并不能解释单调性（如果沿着一行行进，我们遇到负数，我们不应该走得更远）。

更一般的例子

x = np.arange(10000,dtype=np.float64).reshape(100,100)
y = np.arange(4000,dtype=np.float64)

M = x[:,:,None] - y[None,None,:]

i_pos = M>0

M[i_pos] = np.sqrt(M[i_pos])
M[~i_pos] = -np.inf

需要一秒以上。

有没有办法合理地利用这个结构来提高速度？我试图用一个循环来计算一个 Numba 函数，但它似乎并没有快多少（1.38 秒对 1.18 秒 + 编译时间）。

Answer 1

内存分配是限制部分

如果您使用 sqrt、exp、sin、cos 等函数，请确保您已安装 Intel SVML。还要添加您到目前为止尝试过的所有事情以及您是如何获得时间的。

您的 Numpy 解决方案使用临时数组，因此应该有相当大的改进空间。

例子

import numpy as np
import numba as nb

def calc_M(x,y):
    M = x[:,:,None] - y[None,None,:]

    i_pos = M>0

    M[i_pos] = np.sqrt(M[i_pos])
    M[~i_pos] = -np.inf
    return M

@nb.njit(parallel=True)
def calc_M_nb_1(x,y):
    res=np.empty((x.shape[0],x.shape[1],y.shape[0]))
    for i in nb.prange(x.shape[0]):
        for j in range(x.shape[1]):
            for k in range(y.shape[0]):
                val= x[i,j]-y[k]
                if val>0:
                    res[i,j,k]=np.sqrt(val)
                else:
                    res[i,j,k]=-np.inf
    return res

@nb.njit(parallel=True)
def calc_M_nb_2(x,y,res):
    for i in nb.prange(x.shape[0]):
        for j in range(x.shape[1]):
            for k in range(y.shape[0]):
                val= x[i,j]-y[k]
                if val>0:
                    res[i,j,k]=np.sqrt(val)
                else:
                    res[i,j,k]=-np.inf
    return res

时间

x = np.arange(10000,dtype=np.float64).reshape(100,100)
y = np.arange(4000,dtype=np.float64)
res=np.empty((x.shape[0],x.shape[1],y.shape[0]))

%timeit res_1=calc_M(x,y)
469 ms ± 3.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res_2=calc_M_nb_1(x,y)
79.5 ms ± 671 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
res=np.empty((x.shape[0],x.shape[1],y.shape[0]))
%timeit res_3=calc_M_nb_2(x,y,out)
25.5 ms ± 204 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

如果比较 calc_M_nb_2 和 calc_M_nb_1 的时间可以看出，内存分配比整个计算花费的时间更多，这在简单函数中并不罕见。当然，并不总是可以避免内存分配，但是如果你想将此函数应用于很多相似的输入（数组大小），你可以获得相当大的加速。

利用 Numpy 数组中的单调性

Exploiting Monotonicity in Numpy Array

python

arrays

numpy

numba

内存分配是限制部分