如何通过向量化方式从矩阵中提取多组指定行列的子矩阵？

问题背景 #

首先生成示例矩阵A：

import numpy as np

n, d = 5, 100
X = np.random.normal(size=(n, d))
A = X @ X.T / d

输出的A为：

array([[ 0.8728073 ,  0.03398258, -0.00736437, -0.03662121, -0.01160379],
       [ 0.03398258,  1.02946313,  0.08282249, -0.10650558, -0.23883738],
       [-0.00736437,  0.08282249,  1.10018812,  0.00730896, -0.27687929],
       [-0.03662121, -0.10650558,  0.00730896,  1.07049803,  0.02455436],
       [-0.01160379, -0.23883738, -0.27687929,  0.02455436,  1.1720245 ]])

使用np.ix_可以提取单个子矩阵：

idx = np.array([0,1])
sub_matrix = A[np.ix_(idx, idx)]

得到结果：

array([[0.8728073 , 0.03398258],
       [0.03398258, 1.02946313]])

但传入多组行列索引时，np.ix_会报错：

idx = np.array([[0,1], [0,2], [1,2]])
A[np.ix_(idx, idx)]

报错信息：

ValueError: Cross index must be 1 dimensional

需要实现向量化方法，得到如下三维数组结果：

array([[[0.8728073 , 0.03398258],
        [0.03398258, 1.02946313]],

       [[0.8728073 , -0.00736437],
        [-0.00736437, 1.10018812]],

       [[1.02946313, 0.08282249],
        [0.08282249, 1.10018812]]
])

解决方案 #

利用numpy的广播索引实现向量化提取，无需循环：

idx = np.array([[0,1], [0,2], [1,2]])
# 扩展行索引维度为(3,2,1)，列索引维度为(3,1,2)，通过广播匹配为(3,2,2)
row_idx = idx[:, :, np.newaxis]
col_idx = idx[:, np.newaxis, :]
result = A[row_idx, col_idx]

验证结果：

print(result)

输出与目标结果完全一致。

原理说明 #

• row_idx将每组子矩阵的行索引扩展为三维数组，形状为(组数, 子矩阵行数, 1)
• col_idx将每组子矩阵的列索引扩展为三维数组，形状为(组数, 1, 子矩阵列数)
• numpy会自动广播这两个索引数组到相同的形状(组数, 子矩阵行数, 子矩阵列数)，从而一次性提取所有子矩阵

内容的提问来源于stack exchange，提问作者Sengiley