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

Question

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

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()

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

更新：至少 hamiltonian 创建（以前是瓶颈）现在是用 numba 完成的。 然而，它仍然是循环中最慢的 function，正如性能工具的 output 中所见。

Answer 1

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

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

• 委托给 numpy 所有可以用矢量方式完成的东西，比如“Ntime = Nt1 + np.sign(Nt2 - Nt1) * Ntr”。
• 使用临时数据结构（如列表）可以更好地完成文件写入（I/0 瓶颈）。 file.write('\t'.join(lot_of_values))
• 可以记住所有重复的 function 调用（如@AKX 所述）。
• 最后，您可以使用cython或像numba这样的自动编译器在 C 中编写一些内容。 对于 function 和有限的 scope，这些可以帮助你加速很多。