Numpy中张量乘法的对角块乘法

时间:2017-01-13 16:26:25

标签: python numpy

是否有表达式(可能使用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))

1 个答案:

答案 0 :(得分:2)

方法#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