Python - 矩阵乘法与稀疏结果

假设我有两个密集矩阵 U (10000x50) 和 V(50x10000)，以及一个稀疏矩阵 A(10000x10000)。 A 中的每个元素都是 1 或 0。我希望找到 A*(UV)，注意 '*' 是逐元素乘法。为了解决这个问题，Scipy/numpy 会先计算一个密集矩阵 UV。但是紫外线又大又密(10000x10000)，所以速度很慢。

因为我只需要 A 表示的几个 UV 元素，如果只计算必要的元素而不是计算所有元素然后使用 A 过滤，应该可以节省大量时间。 有没有办法指示 scipy 这样做？

顺便说一句，我使用 Matlab 来解决这个问题，而且 Matlab 足够聪明，可以找到我想要做的事情并且有效地工作。

更新:
我发现 Matlab 完全像 scipy 一样计算了 UV。我的 scipy 安装实在是太慢了...

最佳答案

这是一个测试脚本和可能的加速。基本思想是使用 A 的非零坐标选择 U 和 V 的行和列，然后使用 einsum 执行可能点积的子集。

import numpy as np
from scipy import sparse

#M,N,d = 10,5,.1
#M,N,d = 1000,50,.1
M,N,d = 5000,50,.01   # about the limit for my memory

A=sparse.rand(M,M,d)
A.data[:] = 1   # a sparse 0,1 array
U=(np.arange(M*N)/(M*N)).reshape(M,N)
V=(np.arange(M*N)/(M*N)).reshape(N,M)

A1=A.multiply(U.dot(V))   # the direct solution
A2=np.einsum('ij,ik,kj->ij',A.A,U,V)

print(np.allclose(A1,A2))

def foo(A,U,V):
    # use A to select elements of U and V
    A3=A.copy()
    U1=U[A.row,:]
    V1=V[:,A.col]
    A3.data[:]=np.einsum('ij,ji->i',U1,V1)
    return A3

A3 = foo(A,U,V)

print(np.allclose(A1,A3.A))

3 个解决方案匹配。对于大型数组，foo 比直接解决方案快大约 2 倍。对于小尺寸，纯 einsum 是有竞争力的，但对于大数组则陷入困境。

在 dot 中使用 foo 会计算出太多的产品， ij,jk->ik 而不是 ij,ji->i 。