如何矢量化此二维矩阵运算?

时间:2019-10-08 19:02:13

标签: python numpy vectorization

我有2个布尔掩码数组,我想对两个掩码的每个组合计算一个运算。

  

慢速版本

N = 10000
M = 580
masksA = np.array(np.random.randint(0,2, size=(N,M)), dtype=np.bool)
masksB = np.array(np.random.randint(0,2, size=(N,M)), dtype=np.bool)

result = np.zeros(shape=(N,N), dtype=np.float)
for i in range(N):
    for j in range(N):
        result[i,j] = np.float64(np.count_nonzero(masksB[i,:]==masksB[j,:])) / M
  

更快的版本

for i in range(N):
    result[i,:] = np.array(np.count_nonzero(masksB[i]==masksA, axis=1), dtype=np.float64) / M

可以通过一个循环更快地完成此操作吗?

2 个答案:

答案 0 :(得分:3)

基本上是外部平等比较,将这些遮罩轴中的第一个对齐。我们可以利用matrix-multiplication的思想,即掩码自身及其同时求反的掩码的矩阵相乘会导致我们得出那些外部相等求和。因此,对于相等的总和,只需-

out = (~masksB).astype(int).dot(~masksA.T) + masksB.astype(int).dot(masksA.T)

或者,获取int版本并使用它来获取否定版本-

mB = masksB.astype(int)
out = (1-mB).dot(~masksA.T) + mB.dot(masksA.T)

最终输出为out/float(M)。或者,我们可以将这些int转换替换为float,然后将输出除以M

计时(N取为1000)以获得等式求和

# Setup
In [39]: np.random.seed(0)
    ...: N = 1000
    ...: M = 580
    ...: masksA = np.array(np.random.randint(0,2, size=(N,M)), dtype=np.bool)
    ...: masksB = np.array(np.random.randint(0,2, size=(N,M)), dtype=np.bool)

# Approach #1 with int dtype converison
In [40]: %timeit (~masksB).astype(int).dot(~masksA.T) + masksB.astype(int).dot(masksA.T)
1 loop, best of 3: 870 ms per loop

# Approach #1 with float dtype converison
In [41]: %timeit (~masksB).astype(float).dot(~masksA.T) + masksB.astype(float).dot(masksA.T)
10 loops, best of 3: 74 ms per loop

# Approach #2 with int dtype converison
In [42]: %%timeit
    ...: mB = masksB.astype(int)
    ...: out = (1-mB).dot(~masksA.T) + mB.dot(masksA.T)
1 loop, best of 3: 882 ms per loop

# Approach #2 with float dtype converison
In [43]: %%timeit
    ...: mB = masksB.astype(float)
    ...: out = (1-mB).dot(~masksA.T) + mB.dot(masksA.T)
10 loops, best of 3: 59.3 ms per loop

因此,首选方法是采用第二种方法的浮点转换并将输出除以M

将一个矩阵乘法与堆栈一起使用

由于等式求和本质上将是int个数字,因此我们可以对矩阵乘法使用较低精度的dtype。另外,还有一个想法是将原始蒙版及其取反的版本堆叠在一起,然后执行矩阵乘法。因此,对于等式求和,我们将有-

m1 = np.hstack((masksA,~masksA))
m2 = np.hstack((masksB,~masksB))
out = m2.astype(np.float32).dot(m1.T)

时间(与之前的设置相同)-

In [49]: %%timeit
    ...: m1 = np.hstack((masksA,~masksA))
    ...: m2 = np.hstack((masksB,~masksB))
    ...: out = m2.astype(np.float32).dot(m1.T)
10 loops, best of 3: 36.8 ms per loop

答案 1 :(得分:2)

我们仅需使用一个矩阵乘法就可以有效地将@Divakar的(当前)最佳时间减少为两个:

import numpy as np

N = 1000
M = 580
masksA = np.array(np.random.randint(0,2, size=(N,M)), dtype=np.bool)
masksB = np.array(np.random.randint(0,2, size=(N,M)), dtype=np.bool)

def f_pp():
    return np.where(masksB,1.0,-1.0)@np.where(masksA,0.5/M,-0.5/M).T+0.5

def f_pp_2():
    return (np.where(masksB,1.0,-1.0)@np.where(masksA,0.5,-0.5).T+0.5*M)*(1/M)

def f_dv_2():
    mB = masksB.astype(float)
    return ((1-mB).dot(~masksA.T) + mB.dot(masksA.T))*(1/M)

assert np.allclose(f_pp(),f_dv_2())
assert np.all(f_pp_2()==f_dv_2())

from timeit import timeit

print("PP     ",timeit(f_pp,number=100)*10,"ms")
print("PP 2   ",timeit(f_pp_2,number=100)*10,"ms")
print("Divakar",timeit(f_dv_2,number=100)*10,"ms")

样品运行:

PP      31.41063162125647 ms
PP 2    31.757128546014428 ms
Divakar 63.400797033682466 ms

它如何工作?通过将True和False映射到相反的数字x-x(两个因子可以使用不同的x;假设a代表maskA和b对于masksB),我们在点积中对于掩码一致的每个位置获得一个术语ab,对于每个掩码不同的位置获得一个术语-ab。如果k是一对向量之间的相等位数,则其点积将为kab-(M-k)ab = 2kab - Mab。选择ab,使2Mab = 1变成k/M - 1/2,这与我们期望的结果有一定的偏移量。