python - 如何在 python 中向量化包含 eigvalsh 的复杂代码

Question

我有以下代码(抱歉，它不是太小，我已经尝试从原来的代码中减少它)。

基本上，我在运行 eval_s() 方法/函数时遇到性能问题，我在其中:

1) 使用 eigvalsh()

找到 4x4 厄密矩阵的 4 个特征值

2) 将特征值的倒数求和到一个变量result

3) 然后我对许多由 x,y,z 参数化的矩阵重复步骤 1 和 2，将累积和存储在 result 中。

我在第 3 步中重复计算(求特征值和求和)的次数取决于我代码中的变量 ksep，我需要这个数字来增加我的实际代码(即， ksep 必须减少)。但是 eval_s() 中的计算在 x,y,z 上有一个 for 循环，我猜这确实会减慢速度。[试试 ksep=0.5 看看我的意思。]

有没有办法对我的示例代码中指示的方法进行矢量化(或者通常，涉及查找参数化矩阵的特征值的函数)？

代码:

import numpy as np
import sympy as sp
import itertools as it
from sympy.abc import x, y, z


class Solver:
    def __init__(self, vmat):
        self._vfunc = sp.lambdify((x, y, z),
                                  expr=vmat,
                                  modules='numpy')
        self._q_count, self._qs = None, []  # these depend on ksep!

    ################################################################
    # How to vectorize this?
    def eval_s(self, stiff):
        assert len(self._qs) == self._q_count, "Run 'populate_qs' first!"
        result = 0
        for k in self._qs:
            evs = np.linalg.eigvalsh(self._vfunc(*k))
            result += np.sum(np.divide(1., (stiff + evs)))
        return result.real - 4 * self._q_count
    ################################################################

    def populate_qs(self, ksep: float = 1.7):
        self._qs = [(kx, ky, kz) for kx, ky, kz
                    in it.product(np.arange(-3*np.pi, 3.01*np.pi, ksep),
                                  np.arange(-3*np.pi, 3.01*np.pi, ksep),
                                  np.arange(-3*np.pi, 3.01*np.pi, ksep))]
        self._q_count = len(self._qs)


def test():
    vmat = sp.Matrix([[1, sp.cos(x/4+y/4), sp.cos(x/4+z/4), sp.cos(y/4+z/4)],
                      [sp.cos(x/4+y/4), 1, sp.cos(y/4-z/4), sp.cos(x/4 - z/4)],
                      [sp.cos(x/4+z/4), sp.cos(y/4-z/4), 1, sp.cos(x/4-y/4)],
                      [sp.cos(y/4+z/4), sp.cos(x/4-z/4), sp.cos(x/4-y/4), 1]]) * 2
    solver = Solver(vmat)
    solver.populate_qs(ksep=1.7)  # <---- Performance starts to worsen (in eval_s) when ksep is reduced!
    print(solver.eval_s(0.65))


if __name__ == "__main__":
    import timeit
    print(timeit.timeit("test()", setup="from __main__ import test", number=100))

附注代码的 sympy 部分可能看起来很奇怪，但它在我的原始代码中有用。

Answer 1

你可以，方法如下:

def eval_s_vectorized(self, stiff):
    assert len(self._qs) == self._q_count, "Run 'populate_qs' first!"
    mats = np.stack([self._vfunc(*k) for k in self._qs], axis=0)
    evs = np.linalg.eigvalsh(mats)
    result = np.sum(np.divide(1., (stiff + evs)))
    return result.real - 4 * self._q_count

这仍然留下未向量化的 Sympy 表达式的计算。这部分向量化有点棘手，主要是因为输入矩阵中的 1。您可以通过修改 Solver 来制作代码的完全矢量化版本，以便它用 vmat 中的数组常量替换标量常量:

import itertools as it
import numpy as np
import sympy as sp
from sympy.abc import x, y, z
from sympy.core.numbers import Number
from sympy.utilities.lambdify import implemented_function

xones = implemented_function('xones', lambda x: np.ones(len(x)))
lfuncs = {'xones': xones}

def vectorizemat(mat):
    ret = mat.copy()
    # get the first element of the set of symbols that mat uses
    for x in mat.free_symbols: break
    for i,j in it.product(*(range(s) for s in mat.shape)):
        if isinstance(mat[i,j], Number):
            ret[i,j] = xones(x) * mat[i,j]
    return ret

class Solver:
    def __init__(self, vmat):
        self._vfunc = sp.lambdify((x, y, z),
                                  expr=vectorizemat(vmat),
                                  modules=[lfuncs, 'numpy'])
        self._q_count, self._qs = None, []  # these depend on ksep!

    def eval_s_vectorized_completely(self, stiff):
        assert len(self._qs) == self._q_count, "Run 'populate_qs' first!"
        evs = np.linalg.eigvalsh(self._vfunc(*self._qs.T).T)
        result = np.sum(np.divide(1., (stiff + evs)))
        return result.real - 4 * self._q_count

    def populate_qs(self, ksep: float = 1.7):
        self._qs = np.array([(kx, ky, kz) for kx, ky, kz
                    in it.product(np.arange(-3*np.pi, 3.01*np.pi, ksep),
                                  np.arange(-3*np.pi, 3.01*np.pi, ksep),
                                  np.arange(-3*np.pi, 3.01*np.pi, ksep))])
        self._q_count = len(self._qs)

测试/计时

对于小型 ksep，矢量化版本比原始版本快 2 倍，完全矢量化版本快 20 倍:

# old version for ksep=.3
import timeit
print(timeit.timeit("test()", setup="from __main__ import test", number=10))
-85240.46154500882
-85240.46154500882
-85240.46154500882
-85240.46154500882
-85240.46154500882
-85240.46154500882
-85240.46154500882
-85240.46154500882
-85240.46154500882
-85240.46154500882
118.42847006605007

# vectorized version for ksep=.3
import timeit
print(timeit.timeit("test()", setup="from __main__ import test", number=10))
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
64.95763925800566

# completely vectorized version for ksep=.3
import timeit
print(timeit.timeit("test()", setup="from __main__ import test", number=10))
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
5.648927717003971

矢量化版本结果中的舍入误差与原始版本略有不同。这大概是由于 result 中的总和的计算方式不同所致。