python - numpy 中两个二维掩码数组的快速内积

Question

我的问题如下。我有两个形状为 n, p 的数组 X 和 Y，其中 p >> n (例如 n = 50，p = 10000)。

我还有一个关于p的掩码mask(大小为p的 bool 值的一维数组)，小密度(例如np.mean(mask)为0.05)。

我尝试尽快计算 X 和 Y 相对于 mask 的内积:输出 inner 是一个形状为 n, n 的数组，并且满足 inner[i, j] = np.sum(X[i, np.logic_not(mask)) ] * Y[j, np.logic_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(Y.T) 慢大约 5 倍...

首先，我认为这是一个问题，.dot不知道掩码仅相对于p，但如果可能的话我不知道使用此信息。

我正在寻找一种执行计算的方法，而又不会比天真的点慢很多。

非常感谢!

Answer 1

我们可以使用带或不带屏蔽版本的矩阵乘法，因为从完整版本中进行屏蔽减法会产生我们所需的输出 -

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

或者简单地使用反向掩码，但对于稀疏掩码来说会更慢 -

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