Taking a Square Submatrix of a Square Matrix?

Question

Is there a more Pythonic / better way of doing this?

def square_take( m, idx ):
    """ Take a square submatrix of m """
    return m.take( idx, axis=0 ).take( idx, axis=1 )

m = np.eye(3)
idx = np.array( [0, 1] )
print square_take( m, idx )

It looks ugly, and I was wondering if I could do it using a single command.

Update:

I have compared the two solutions with the original method:

def square_take( m, idx ):
    """ Take a square submatrix of m """
    return m.take( idx, axis=0 ).take( idx, axis=1 )

def square_take2( m, idx ):
    """ Take a square submatrix of m """
    return m[ np.ix_( idx, idx ) ]

def square_take3( m, idx ):
    """ Take a square submatrix of m """
    return m[idx][:,idx]

N = 10000
cv = np.random.randn( 30, 30 )

idx = np.array( [ 1, 2, 5, 6 ] )
with Timer( 'Square Take 1 ' ):
    for _ in range( N ):
        r1 = helpers.square_take( cv, idx )

with Timer( 'Square Take 2 ' ):
    for _ in range( N ):
        r2 = helpers.square_take2( cv, idx )

with Timer( 'Square Take 3 ' ):
    for _ in range( N ):
        r3 = helpers.square_take3( cv, idx )

np.testing.assert_array_equal( r1, r2 )
np.testing.assert_array_equal( r1, r3 )

• Square Take 1 Took 0.08s seconds

• Square Take 2 Took 0.30s seconds

• Square Take 3 Took 0.10s seconds

Unfortunately, while the two solutions look more elegant, they are slower than the original.

Answer 1

Use the numpy.ix_ function:

m[np.ix_(idx, idx)]

Answer 2

 m[idx][:,idx]

You can index like this just because they are numpy arrays.