采用方矩阵的平方子矩阵?

时间:2015-03-20 20:34:34

标签: python numpy

有更多的Pythonic /更好的方法吗?

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 )

它看起来很难看,我想知道我是否可以使用一个命令来完成它。

更新

我将两种解决方案与原始方法进行了比较:

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花费0.08秒

  • Square Take 2花了0.30秒

  • Square Take 3 0.10秒秒

不幸的是,虽然这两个解决方案看起来更优雅,但它们比原始解决方案慢。

2 个答案:

答案 0 :(得分:4)

使用numpy.ix_功能:

m[np.ix_(idx, idx)]

答案 1 :(得分:2)

 m[idx][:,idx]

你可以像这样索引,因为它们是numpy数组。