Get multiple submatrices given row and column ids-CodePudding

I want to get a vectorized way to get multiple submatrices of a matrix given their row and column indices.

import numpy as np

n, d = 5, 100

X = np.random.normal(size=(n, d))
A = X @ X.T / d
print(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 ]])

I can extract a single submatrix via np.ix_

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

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

but passing multiple row/column indices doesn't work

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

ValueError: Cross index must be 1 dimensional

How to achieve that? So I need a vectorized way to obtain

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

CodePudding user response：

This should select the sub-matrices how you want:

A[idx.reshape(-1, 2, 1), idx.reshape(-1, 1, 2)]

Example:

A = np.arange(9).reshape(-1, 3)

[[0 1 2]
 [3 4 5]
 [6 7 8]]

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

[[0 1]
 [0 2]
 [1 2]]

A[idx.reshape(-1, 2, 1), idx.reshape(-1, 1, 2)]

[[[0 1]
  [3 4]]

 [[0 2]
  [6 8]]

 [[4 5]
  [7 8]]]