高效点积密集到稀疏

Question

我想高效地计算两个低秩矩阵A和B的点积，只针对特定的指数。通常，A和B的形状为(10000,100%)，我只希望获得A.T @B的1%的条目。

我已经尝试过使用Numba库，但它似乎比计算密集的点积np.dot(A.T，B)慢得多。

Answer 1

如果您想要一个点积的结果数组的子集，只需在使用dot之前获取该子集。也就是说，如果你想要输出的“左上角”10x10矩阵，只需这样做

np.dot(A.T[:10,:], B[:, :10])

如果你想要一些特殊的索引，你可以使用更花哨的索引。例如，如果您想要索引3、5和29，您可以这样做：

indices = np.array([3, 5, 29]).reshape(-1, 1)
inner_all = np.arange(A.shape[0]).reshape(-1, 1)
result = np.dot(A.T[indices, inner_all.T], B[inner_all, indices.T])

如果只需要第一行，则需要第1、3和5列：

rows = np.array([1]).reshape(-1, 1)
columns = rows = np.array([1, 3, 5]).reshape(-1, 1)
inner_all = np.arange(A.shape[0]).reshape(-1, 1)
result = np.dot(A.T[rows, inner_all.T], B[inner_all, columns.T])

Answer 2

您可以使用einsum

import numpy as np

def direct():
    return (A.T@B)[iy, ix]

def einsum():
    return np.einsum('ij,ij->j',A[:,iy],B[:,ix])

def make(M,N,p):
    global A,B,iy,ix
    A,B = (np.random.randint(0,10,(M,N)) for _ in "AB")
    iy,ix = np.random.randint(0,N,(2,int(N*N*p)))

from time import time

make(100,10000,0.01)
T=[]
T.append(time())
d = direct()
T.append(time())
e = einsum()
T.append(time())
print("results equal: ",np.allclose(d,e))
print("direct {:.6f} einsum {:.6f}".format(*np.diff(T)))

示例运行：

results equal:  True
direct 18.463711 einsum 1.947425