矢量化linalg.eig()在numpy

时间:2012-11-23 15:13:28

标签: optimization numpy linear-algebra

我有一个m * m * n numpy数组(称之为A),我想在这个数组中找到每个子矩阵A[:,:,n]的特征值。我可以在一个相对容易的循环中使用linalg.eig()来完成它,但实际上应该有一种方法来实现这一点。类似ufunc的东西,但可以处理子向量而不是单个元素。这可能吗?

1 个答案:

答案 0 :(得分:1)

特征值和特征向量的计算不能被矢量化,因为通常无法共享不同矩阵的工作。 np.linalg.eig(对于实际输入)只是dgeev的包装器,根据文档,每次调用只接受一个矩阵并且计算相当昂贵,因此对于不小的矩阵,一个python循环可以忽略不计。

但是,如果你为许多非常小的矩阵做这个,它可能会变得太慢。有several questions与此相关,解决方案通常最终成为编译扩展。正如神秘的物理学家在评论中所说,以与ufuncs相同的方式处理子向量和子矩阵的想法一般都是有用的。这些被称为广义ufuncs,已经在numpy的development version。对于形状为(1000, 3, 3)的矩阵,我发现它快8倍左右:

In [2]: np.__version__
Out[2]: '1.8.0.dev-dcf7cac'

In [3]: A = np.random.rand(1000, 3, 3)

In [4]: timeit np.linalg.eig(A)
P100 loops, best of 3: 9.65 ms per loop

In [5]: timeit [np.linalg.eig(Ai) for Ai in A]
10 loops, best of 3: 74.6 ms per loop

In [6]: a1 = np.linalg.eig(A)

In [7]: a2 = [np.linalg.eig(Ai) for Ai in A]

In [8]: all(np.allclose(a1[i][j], a2[j][i]) for j in xrange(1000) for i in xrange(2))
Out[8]: True