python中嵌套for循环的时间相关krylov时间演化太慢

Question

我已经为量子系统实现了在两个时间点 Nt1 和 Nt2 评估的相关函数。最后，我们必须对所有可能的 Nt1 和 Nt2 进行积分，因此我们必须计算所有可能的 Nt1、Nt2 在 0、...、Nt_max 中的相关函数。计算是使用 64 个处理器的 1 个 hpc 完成的。不同 Nt1 的计算是并行的，使用 bash 例程将作业提交给其中一个处理器。一项工作包括以下例程，以计算 k 相关相关函数：

def propagater(Nt1, Nt_max, phi, dim_krylov):

    # phi is the ground state of the system
    v1 = phi
    G_t2_k = np.zeros((Nt_max, len(k_all)), dtype=complex)

    for Ntr in range(Nt1):
        # create Ntr dependent Hamiltonian 1
        H_spin = make_H_spin(Ntr)
        # get eigenvalues, eigenvectors, and lanczos vectors with a lanczos method
        E_spin, V_spin, Q_T_spin = lanczos_full(H_spin, v1, dim_krylov)
        # apply time propagation for timestep dt: v1 = e^(-ij*H_spin*dt)*v1
        v1 = expm_lanczos(E_spin, V_spin, Q_T_spin, a=-1j * dt)

    # manipulate state, that implies as basis change
    v2 = createHole(v1)

    # propagation for all possible Nt2
    for Nt2 in range(Nt_max):
        v3 = v2
        # propagate v3 to t2
        # if t2<t1 we have to propagate backward in time, hence the np.sign
        for Ntr in range(np.abs(Nt2 - Nt1)):
            Ntime = Nt1 + np.sign(Nt2 - Nt1) * Ntr
            # create Ntr dependent Hamiltonian 2
            H_tJ = make_H_tJ(Ntime)
            E_tJ, V_tJ, Q_T = lanczos_full(H_tJ, v3, dim_krylov)
            v3 = expm_lanczos(E_tJ, V_tJ, Q_T, a=-np.sign(Nt2 - Nt1) * 1j * dt)

        # propagate <phi| to t2 (corresponds to v3 backward in time)
        v4 = phi
        for Ntr in range(Nt2):
            H_spin = make_H_spin(Ntr)
            E_spin, V_spin, Q_T_spin = lanczos_full(H_spin, v4, dim_krylov)
            v4 = expm_lanczos(E_spin, V_spin, Q_T_spin, a=-1j * dt)

        # now store the expectation value
        file = open("prop_" + str(Nt1) + "_bash.txt", "a")
        # manipulation of state implies a k dependence as well
        for k_id, k in enumerate(k_all):
            phi_pj = createHole(v4)
            G_t2_k_r = np.real(np.vdot(phi_pj, v3))
            G_t2_k_im = np.imag(np.vdot(phi_pj, v3))
            file.write(str(G_t2_k_r) + "\t" + str(G_t2_k_im) + "\t")
            G_t2_k[Nt2, k_id] = np.vdot(phi_pj, v3)
        file.write("\n")
        file.close()

我不想在这里详述所有细节，但我希望大体思路清晰。但是，此代码运行速度太慢，无法获得一些不错的结果。我现在的问题是，如何加快速度。一般来说，我认为最关键的点是循环（我知道它们在 python 中非常慢）。但我不知道如何对它们进行矢量化。我考虑使用 numba，但不幸的是它不支持我需要存储系统的 hamiltonian 的 csr 格式。另一种选择是 pypy ，直到现在我还没有检查过。当我读到那里的循环更快时，我还考虑用 C++ 重写代码，但我也不确定优势有多大。因此，欢迎所有帮助！

更新：至少 hamiltonian 创建（以前是瓶颈）现在是用 numba 完成的。然而，正如在performance tool 的输出中所见，它仍然是循环中最慢的函数。

Answer 1

第一步是确定在哪里花费的时间最多，为此您可以使用分析器。如果单个函数太慢或被调用太多次，此代码从 O(n^2) 开始，因此请注意此处如何使用该函数。 对于这些情况，我写了一个包perf_tool它可以指导你正确的方式。

如果研究结果没有任何惊人之处，一般来说：

• 委托 numpy 所有可以通过矢量完成的东西，例如“Ntime = Nt1 + np.sign(Nt2 - Nt1) * Ntr”。
• 使用临时数据结构（如列表）可以更好地完成文件写入（I/0 瓶颈）。 file.write('\t'.join(lot_of_values))
• 可以记住所有重复的函数调用（如 @AKX 所述）。
• 最后，您可以使用cython 或numba 等自动编译器在C 中编写一些代码。对于范围有限的功能，这些可以帮助您大大加快速度。