使用数组提高Python代码的效率

q3qa4bjr  于 2023-03-31  发布在  Python
关注(0)|答案(1)|浏览(131)

我有下面的python代码:

import matplotlib.pyplot as plt
import numpy as np

class bifurcation_diagram(object):
    def __init__(self):
        self.omega = []
        self.theta = []
        self.dt = (2 * np.pi / (2.0 / 3.0)) / 600
        self.time = []
        self.theta_in_bifurcation_diagram = []
        self.F_D = np.arange(1.35,1.5,0.001)
        self.theta_versus_F_D = []
    def calculate(self):
        l = 9.8
        g = 9.8
        q = 0.5
        Omega_D = 2.0 / 3.0
        for f_d in self.F_D:
            self.omega.append([0])
            self.theta.append([0.2])
            self.time.append([0])
            for i in range(600000):
                k1_theta = self.dt * self.omega[-1][-1]
                k1_omega = self.dt * ((-g / l) * np.sin(self.theta[-1][-1]) - q * self.omega[-1][-1] + f_d * np.sin(Omega_D * self.time[-1][-1]))
                k2_theta = self.dt * (self.omega[-1][-1] + 0.5 * k1_omega)
                k2_omega = self.dt * ((-g / l) * np.sin(self.theta[-1][-1] + 0.5 * k1_theta) - q * (self.omega[-1][-1] + 0.5 * k1_omega) + f_d * np.sin(Omega_D * (self.time[-1][-1] + 0.5 * self.dt)))
                k3_theta = self.dt * (self.omega[-1][-1] + 0.5 * k2_omega)
                k3_omega = self.dt * ((-g / l) * np.sin(self.theta[-1][-1] + 0.5 * k2_theta) - q * (self.omega[-1][-1] + 0.5 * k2_omega) + f_d * np.sin(Omega_D * (self.time[-1][-1] + 0.5 * self.dt)))
                k4_theta = self.dt * (self.omega[-1][-1] + k3_omega)
                k4_omega = self.dt * ((-g / l) * np.sin(self.theta[-1][-1] + k3_theta) - q * (self.omega[-1][-1] + k3_omega) + f_d * np.sin(Omega_D * (self.time[-1][-1] + self.dt)))
                temp_theta = self.theta[-1][-1] + (1.0 / 6.0) * (k1_theta + 2 * k2_theta + 2 * k3_theta + k4_theta)
                temp_omega = self.omega[-1][-1] + (1.0 / 6.0) * (k1_omega + 2 * k2_omega + 2 * k3_omega + k4_omega)
                while temp_theta > np.pi:
                    temp_theta -= 2 * np.pi
                while temp_theta < -np.pi:
                    temp_theta += 2 * np.pi
                self.omega[-1].append(temp_omega)
                self.theta[-1].append(temp_theta)
                self.time[-1].append(self.dt * i)
            for i in range(500,1000):
                n = i * 600
                self.theta_in_bifurcation_diagram.append(self.theta[-1][n])
                self.theta_versus_F_D.append(f_d)
    def show_results(self):
        plt.plot(self.theta_versus_F_D,self.theta_in_bifurcation_diagram,'.')
        plt.title('Bifurcation diagram' + '\n' + r'$\theta$ versus $F_D$')
        plt.xlabel(r'$F_D$')
        plt.ylabel(r'$\theta$ (radians)')
        plt.xlim(1.35,1.5)
        plt.ylim(1,3)
        plt.show()

        
bifurcation = bifurcation_diagram()
bifurcation.calculate()
bifurcation.show_results()

我想把它做得更紧凑,做一个彩色的图,沿着提高它的效率。在这方面的任何帮助都将是真正有益的。
我期望代码运行得快,但它需要超过15分钟的运行时间。

7gcisfzg

7gcisfzg1#

Python不是为这样的代码而设计的,因为它通常使用CPython进行解释(这意味着要为粘合代码和脚本完成)和CPython执行(几乎)没有对代码进行优化,因此重复的表达式会被重新计算。加快此类代码速度的常用方法是使用Numpy函数对大型数组进行操作(不是列表)。这被称为矢量化。也就是说,这段代码非常昂贵,使用Numpy可能还不够,在这里使用Numpy也是一个不可忽视的工作。另一种解决方案是使用Numbaso编译这段代码注意,Numba旨在加速主要使用Numpy和基本数值计算的代码(而不是像绘图这样的通用代码)。
首先要做的是摆脱列表并将其替换为Numpy数组。数组可以预先分配合适的大小,以避免对append的昂贵调用。然后,计算功能应该从类中移走,以便Numba易于使用。然后,两个while循环可以替换为np.remainder。最后,你可以使用多个线程并行计算数组的每一行。
下面是生成的代码(几乎没有测试过):

import matplotlib.pyplot as plt
import numpy as np
import numba as nb

@nb.njit('(float64[::1], float64)', parallel=True)
def compute_numba(F_D, dt):
    l = 9.8
    g = 9.8
    q = 0.5
    Omega_D = 2.0 / 3.0
    size = 600000
    loop2_start, loop2_end = 500, 1000
    loop2_stride = 600
    omega = np.empty((F_D.size, size+1), dtype=np.float64)
    theta = np.empty((F_D.size, size+1), dtype=np.float64)
    time = np.empty((F_D.size, size+1), dtype=np.float64)
    theta_in_bifurcation_diagram = np.empty((loop2_end-loop2_start) * F_D.size, dtype=np.float64)
    theta_versus_F_D = np.empty((loop2_end-loop2_start) * F_D.size, dtype=np.float64)
    for i in nb.prange(F_D.size):
        f_d = F_D[i]
        omega[i, 0] = 0.0
        theta[i, 0] = 0.2
        time[i, 0] = 0.0
        for j in range(size):
            cur_omega = omega[i,j]
            cur_theta = theta[i,j]
            cur_time = time[i,j]
            tmp = np.sin(Omega_D * (cur_time + 0.5 * dt))
            k1_theta = dt * cur_omega
            k1_omega = dt * ((-g / l) * np.sin(cur_theta) - q * cur_omega + f_d * np.sin(Omega_D * cur_time))
            k2_theta = dt * (cur_omega + 0.5 * k1_omega)
            k2_omega = dt * ((-g / l) * np.sin(cur_theta + 0.5 * k1_theta) - q * (cur_omega + 0.5 * k1_omega) + f_d * tmp)
            k3_theta = dt * (cur_omega + 0.5 * k2_omega)
            k3_omega = dt * ((-g / l) * np.sin(cur_theta + 0.5 * k2_theta) - q * (cur_omega + 0.5 * k2_omega) + f_d * tmp)
            k4_theta = dt * (cur_omega + k3_omega)
            k4_omega = dt * ((-g / l) * np.sin(cur_theta + k3_theta) - q * (cur_omega + k3_omega) + f_d * np.sin(Omega_D * (cur_time + dt)))
            temp_theta = cur_theta + (1.0 / 6.0) * (k1_theta + 2 * k2_theta + 2 * k3_theta + k4_theta)
            temp_omega = cur_omega + (1.0 / 6.0) * (k1_omega + 2 * k2_omega + 2 * k3_omega + k4_omega)
            temp_theta = np.remainder(temp_theta + np.pi, 2 * np.pi) - np.pi
            omega[i,j+1] = temp_omega
            theta[i,j+1] = temp_theta
            time[i,j+1] = dt * j
        for k, j in enumerate(range(loop2_start, loop2_end)):
            n = j * loop2_stride
            offset = (loop2_end - loop2_start) * i + k
            theta_in_bifurcation_diagram[offset] = theta[i,n]
            theta_versus_F_D[offset] = f_d
    return (omega, theta, time, theta_in_bifurcation_diagram, theta_versus_F_D)

class bifurcation_diagram(object):
    def __init__(self):
        self.omega = np.array([])
        self.theta = np.array([])
        self.dt = (2 * np.pi / (2.0 / 3.0)) / 600
        self.time = np.array([])
        self.theta_in_bifurcation_diagram = np.array([])
        self.F_D = np.arange(1.35,1.5,0.001)
        self.theta_versus_F_D = np.array([])
    def calculate(self):
        self.omega, self.theta, self.time, self.theta_in_bifurcation_diagram, self.theta_versus_F_D = compute_numba(self.F_D, self.dt)
    def show_results(self):
        plt.plot(self.theta_versus_F_D,self.theta_in_bifurcation_diagram,'.')
        plt.title('Bifurcation diagram' + '\n' + r'$\theta$ versus $F_D$')
        plt.xlabel(r'$F_D$')
        plt.ylabel(r'$\theta$ (radians)')
        plt.xlim(1.35,1.5)
        plt.ylim(1,3)
        plt.show()

        
bifurcation = bifurcation_diagram()
bifurcation.calculate()
bifurcation.show_results()

在我的6核i5- 9600 KF处理器上,这个代码需要1.07秒,而不是最初的19分20秒!因此,这个新代码大约快150倍
此外,我还发现生成的代码更容易阅读。
我建议你检查一下结果,虽然它们乍一看很好。事实上,这是我到目前为止得到的结果图:

相关问题