时间戳

首页 > 解决方案 > 几个 2x2 矩阵序列的向量化乘法

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

问题描述

对于未矢量化的版本，我有一系列(2, 2)形状矩阵（即形状的 ndarray (n, 2, 2)），需要按顺序将它们相乘（矩阵乘法），这意味着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这样的矩阵序列数组（一个形状的ndarray (k, n, 2, 2)），并且需要为每个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. 这实际上大大减慢了速度

3x33）我可以看到矩阵在遥远的未来应用，但具体的解决方案2x2很好。无论哪种方式，所有矩阵都是方形的，并且具有相同的维度

标签： pythonnumpyvectorization

解决方案

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, :, :])

推荐阅读