几个2x2矩阵序列的向量化乘法

时间:2019-06-12 18:50:19

标签: python numpy vectorization

对于非向量化版本,我有一系列(2, 2)形状的矩阵(即形状为(n, 2, 2)的ndarray),需要顺序地将它们相乘(矩阵乘法),这意味着顺序为n矩阵乘法。 这就是一个最小的例子。

def get_matrix_product_eig_val(J):
    # J holds the sequence of matrices to multiply and has shape=(n, 2, 2)
    M = np.identity(2, dtype=np.double)
    for i in range(n_gens):
        M = np.matmul(M, J[i])
    eig_val, eig_vec = np.linalg.eig(M)  # eig_val is what I'm interested in
    return eig_val

现在我有一个k的矩阵序列数组(形状为(k, n, 2, 2)的ndarray),并且需要为每个k条目执行相同的顺序矩阵乘积。 天真的方法是这样做

# Now J has shape=(k, n, 2, 2)
for i in range(k):
    eig_vals[i] = get_matrix_product_eig_val(J[i, :, :, :])

是否有摆脱循环的方法,并以向量化的方式实现?

注意:

1)n的数量级应为〜100。 k的范围在〜100到〜10,000之间

2)有人建议用np.linalg.multi_dot替换内部循环。实际上,这会使速度大大降低

3)我可以看到将来3x3矩阵的应用很遥远,但是2x2的特定解决方案很好。无论哪种方式,所有矩阵都是正方形,并且具有相同的尺寸

1 个答案:

答案 0 :(得分:0)

for i in range(n - 1):
    J[:, i + 1, :, 0] = np.broadcast_to((J[:, i+1, 0, 0])[..., None], (k, 2)) * J[:, i, :, 0] + \
                        np.broadcast_to((J[:, i+1, 1, 0])[..., None], (k, 2)) * J[:, i, :, 1]
    J[:, i + 1, :, 1] = np.broadcast_to((J[:, i+1, 0, 1])[..., None], (k, 2)) * J[:, i, :, 0] + \
                        np.broadcast_to((J[:, i+1, 1, 1])[..., None], (k, 2)) * J[:, i, :, 1]

eig_vals, _ = np.linalg.eig(J[:, -1, :, :])