Numpy:用于复合格式的块对角矩阵的高效线性代数

时间:2016-12-20 13:17:48

标签: numpy multidimensional-array vectorization linear-algebra

我有一个线性系统,其中所有矩阵都是块对角线。它们的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, m)(N, n, p)的两个矩阵中生成形状为(N, p, m)的矩阵?
  • 有没有办法用numpy内置(可能更快)的功能执行这些操作?像np.linalg.inv这样的函数让我认为numpy旨在支持块对角矩阵的压缩格式。

2 个答案:

答案 0 :(得分: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是一种用于复杂线性操作的令人难以置信的多功能结构,并且通常使用两个操作数非常有效。我们的想法是以索引方式重写您的操作:您需要的是为每个M1[i,j,k]乘以M2[i,k,l]i,j,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是"块向量",结果将是块向量。

答案 1 :(得分:1)

从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