在 Numpy 中通过张量乘法进行对角块乘法

Question

是否有一个表达式（可能使用np.tensordot）最能表达地捕获对角块的块矩阵和相应大小的向量之间的稀疏矩阵向量乘法？

我有一个工作实现，可以执行我想要的确切操作，但我使用了两个 python 循环（见下文）而不是一个合适的 numpy 命令，它可能存在。

例如：

import numpy as np

outer_size = 2
inner_size = 5
shape = (outer_size, outer_size, inner_size)
diag_block = np.arange(np.prod(shape)).reshape(shape)
true_diag = np.bmat([[np.diag(diag_block[i,j]) for j in range(shape[1])] for i in range(shape[0])]).A
x = np.arange(shape[1] * shape[2])

def sparse_prod(diags, x):
    outer_size = diags.shape[0]
    return np.hstack(sum(diags[i] * x.reshape(outer_size, -1)) for i in range(outer_size))

print(true_diag.dot(x))
print(sparse_prod(diag_block, x))

Answer 1

方法#1：你可以使用np.einsum -

np.einsum('ijk,jk->ik', diag_block, x.reshape(outer_size, -1)).ravel()

基本上，我们保持最后一个轴在两个输入之间对齐，并分别对第二个和第一个轴进行求和。

---

方法#2：根据输入数组的形状，您可能希望在循环中使用np.dot/np.tensordot，如this post 中详细讨论的那样。

这是一种带有循环的方法 -

m,_,n = diag_block.shape
x2D = x.reshape(outer_size, -1)
out = np.empty((m,n))
for i in range(n):
    out[:,i] = np.dot(diag_block[...,i], x2D[:,i])
out.shape = out.size  # Flatten