[SOLVED] Numpyic way to take the first N rows and columns out of every M rows and columns from a square matrix

Numpyic way to take the first N rows and columns out of every M rows and columns from a square matrix

I have a 20 x 20 square matrix. I want to take the first 2 rows and columns out of every 5 rows and columns, which means the output should be a 8 x 8 square matrix. This can be done in 2 consecutive steps as follows:

import numpy as np

m = 5 
n = 2 
A = np.arange(400).reshape(20,-1)
B = np.asarray([row for i, row in enumerate(A) if i % m < n])
C = np.asarray([col for j, col in enumerate(B.T) if j % m < n]).T

However, I am looking for efficiency. Is there a more Numpyic way to do this? I would prefer to do this in one step.

Solution

You can use np.ix_ to retain the elements whose row / column indices are less than 2 modulo 5:

import numpy as np

m = 5
n = 2
A = np.arange(400).reshape(20,-1)
mask = np.arange(20) % 5 < 2
result = A[np.ix_(mask, mask)]
print(result)

This outputs:

[[  0   1   5   6  10  11  15  16]
 [ 20  21  25  26  30  31  35  36]
 [100 101 105 106 110 111 115 116]
 [120 121 125 126 130 131 135 136]
 [200 201 205 206 210 211 215 216]
 [220 221 225 226 230 231 235 236]
 [300 301 305 306 310 311 315 316]
 [320 321 325 326 330 331 335 336]]