压缩格式的块对角矩阵的高效线性代数

Question

我有一个线性系统，其中所有矩阵都是块对角线。它们具有形状相同的N 块。

矩阵以压缩格式存储为形状为(N, n, m) 的numpy 数组，而向量的形状为(N, m)。

我目前将矩阵向量乘积实现为

import numpy as np

def mvdot(m, v):
    return (m * np.expand_dims(v, -2)).sum(-1)

由于广播规则，如果矩阵的所有块都相同，我只需将其存储在形状为 (1, n, m) 的数组中一次：带有向量 (N, m) 的乘积仍然给出正确的 (N, n) 向量.

我的问题是：

• 如何实现高效的矩阵-矩阵乘积，从形状为(N, n, p) 和(N, p, m) 的两个矩阵中生成形状为(N, n, m) 的矩阵？
• 有没有办法使用numpy 内置（可能更快）函数来执行这些操作？ np.linalg.inv 之类的函数让我觉得 numpy 旨在支持块对角矩阵的这种压缩格式。

Answer 1

如果我正确理解您的问题，您有两个形状分别为 (N,n,p) 和 (N,p,m) 的数组，它们的乘积应该是形状 (N,n,m) 其中元素 [i,:,:] 是 M1[i,:,:] 和的矩阵乘积M2[i,:,:]。这可以使用numpy.einsum 来实现：

import numpy as np
N = 7
n,p,m = 3,4,5
M1 = np.random.rand(N,n,p)
M2 = np.random.rand(N,p,m)
Mprod = np.einsum('ijk,ikl->ijl',M1,M2)

# check if all the submatrices are what we expect
all([np.allclose(np.dot(M1[k,...],M2[k,...]),Mprod[k,...]) for k in range(N)])
# True

Numpy 的einsum 是一种用于复杂线性运算的非常通用的构造，并且通常使用两个操作数非常有效。想法是以索引方式重写您的操作：您需要将每个i,j,l 的M1[i,j,k] 与M2[i,k,l] 相乘，然后将k 相加。这正是上面对einsum 的调用所做的：它折叠索引k，并按照给定顺序沿剩余维度执行必要的乘积和分配。

矩阵-向量积可以类似地完成：

M = np.random.rand(N,n,m)
v = np.random.rand(N,m)
Mvprod = np.einsum('ijk,ik->ij',M,v)

可能numpy.dot 可以通过适当的转置和维度技巧强制执行您想要的操作，但我无法做到。

通过在einsum 中允许隐含的维数，可以在同一个函数调用中完成上述两个操作：

def mvdot(M1,M2):
    return np.einsum('ijk,ik...->ij...',M1,M2)

Mprod = mvdot(M1,M2)
Mvprod = mvdot(M,v)

如果输入参数M2 是块矩阵，则会在结果中附加一个前导维度，从而创建一个块矩阵。如果M2 是“块向量”，则结果将是块向量。

Answer 2

从 Python 3.5 及更高版本开始，可以使用矩阵乘法运算符 @ (numpy.matmul) 简化前面的示例，它将这种情况视为驻留在最后两个索引中的矩阵堆栈并相应地广播：

import numpy as np
N = 7
n,p,m = 3,4,5
M1 = np.random.rand(N,n,p)
M2 = np.random.rand(N,p,m)
Mprod = M1 @ M2  # similar to np.matmul(M1, M2)

all([np.allclose(np.dot(M1[k,...],M2[k,...]),Mprod[k,...]) for k in range(N)])
#True