numpy中兩個二維屏蔽數組的快速內積

Question

我的問題如下。 我有兩個形狀為 n, p 的數組X和Y ，其中p >> n （例如 n = 50，p = 10000）。

我還有一個相對於p小密度的掩碼mask （大小為p的布爾值的一維數組）（例如np.mean(mask)為 0.05）。

我嘗試盡可能快地計算X和Y相對於mask的內積：輸出inner是一個形狀為n, n的數組，並且使inner[i, j] = np.sum(X[i, np.logical_not(mask)] * Y[j, np.logical_not(mask)]) 。

我曾嘗試使用numpy.ma庫，但我的使用速度很慢：

import numpy as np
import numpy.ma as ma
n, p = 50, 10000
density = 0.05
mask = np.array(np.random.binomial(1, density, size=p), dtype=np.bool_)
mask_big = np.ones(n)[:, None] * mask[None, :]
X = np.random.randn(n, p)
Y = np.random.randn(n, p)
X_ma = ma.array(X, mask=mask_big)
Y_ma = ma.array(Y, mask=mask_big)

但是，在我的機器上， X_ma.dot(Y_ma.T)比X.dot(YT)慢大約 5 倍......

首先，我認為.dot不知道掩碼僅與p是一個問題，但如果可能使用此信息，我不知道。

我正在尋找一種方法來執行計算而不會比天真的點慢得多。

非常感謝 ！

Answer 1

我們可以使用帶有和不帶有掩碼版本的matrix-multiplication ，因為從完整版本中進行的掩碼減法會產生我們想要的輸出 -

inner = X.dot(Y.T)-X[:,mask].dot(Y[:,mask].T)

或者簡單地使用反向掩碼，雖然對於稀疏mask會更慢 -

inner = X[:,~mask].dot(Y[:,~mask].T)

時間——

In [34]: np.random.seed(0)
    ...: p,n = 10000,50
    ...: X = np.random.rand(n,p)
    ...: Y = np.random.rand(n,p)
    ...: mask = np.random.rand(p)>0.95

In [35]: mask.mean()
Out[35]: 0.0507

In [36]: %timeit X.dot(Y.T)-X[:,mask].dot(Y[:,mask].T)
100 loops, best of 3: 2.54 ms per loop

In [37]: %timeit X[:,~mask].dot(Y[:,~mask].T)
100 loops, best of 3: 4.1 ms per loop

In [39]: %%timeit
    ...: inner = np.empty((n,n))
    ...: for i in range(X.shape[0]):
    ...:     for j in range(X.shape[0]):
    ...:         inner[i, j] = np.sum(X[i, ~mask] * Y[j, ~mask])
1 loop, best of 3: 302 ms per loop