python - 2D 稀疏矩阵乘以 3D 数组

Question

有什么方法可以将 2D 稀疏矩阵乘以 3D numpy 数组吗？ 例如我有这个功能

def myFun(x, p):
    r = 2
    out = x * np.log(p) + r * np.log(1-p)
    return out

其中 x 是维度 3500, 90 的数组，p 是维度 3500, 90, 70 的另一个数组。目前，x 和 p 都是密集数组，我只是在调用该函数时进行广播:

out = myFun(x[..., None], p)

但是数组 x 非常稀疏，只有 7% 的元素非零。另一方面，p 没有任何零，仅在零和一之间 float 。 我希望通过稀疏矩阵(可能来自 scipy.sparse )我会看到速度的提高。但是，我不知道如何执行此操作，或者这是否更有效。

我使用的是 python 3。

非常感谢

Answer 1

您可以使用 where 关键字来利用 x 的稀疏性。

def sprse(x, p):
    r = 2
    out = x * np.log(p, where=x.astype(bool)) + r * np.log(1-p)
    return out

from timeit import timeit
x = np.random.uniform(-13, 1, (3500, 90, 1)).clip(0, None)
p = np.random.random((3500, 90, 70))
assert np.all(sprse(x, p)==myFun(x, p))
def f():
    return myFun(x, p)
print(timeit(f, number=3))
def f():
    return sprse(x, p)
print(timeit(f, number=3))

示例运行:

5.171174691990018
3.2122434769989923