optimization - 在 numpy 中矢量化 linalg.eig()

Question

我有一个 m*m*n numpy 数组（称为 A），我想找到A[:,:,n]该数组中每个子矩阵的特征值。我可以linalg.eig()相对容易地在一个循环中做到这一点，但确实应该有一种方法来矢量化它。类似 a 的东西ufunc，但它可以处理子向量而不是单个元素。这可能吗？

Answer 1

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

但是，如果您对许多非常小的矩阵执行此操作，它可能会变得太慢。有几个与此相关的问题，解决方案通常最终是一个编译的扩展。正如 enigmaticPhysicist 在评论中所说，以与 ufunc 相同的方式处理子向量和子矩阵的想法通常会很有用。这些被称为通用 ufunc，并且已经在 numpy 的开发版本中。我发现形状矩阵的速度快了大约 8 倍(1000, 3, 3)：

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