优化Python分岔图代码：提升效率、简化结构及实现彩色绘图

我编写了如下用于绘制分岔图的Python代码，但运行耗时超过15分钟，希望简化代码结构、提升运行效率，并实现彩色绘图。

原始代码 #

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

优化方案与代码 #

核心优化点 #

• 用NumPy数组替代列表存储迭代数据，避免频繁append的内存开销
• 移除不必要的time数组，直接通过迭代次数计算时间项
• 预分配内存空间，减少动态扩容的性能损耗
• 简化角度归一化逻辑，用np.mod替代循环判断
• 实现彩色绘图，根据F_D值映射不同颜色

优化后的代码 #

import matplotlib.pyplot as plt
import numpy as np

def generate_bifurcation_diagram():
    # 参数设置
    l = 9.8
    g = 9.8
    q = 0.5
    Omega_D = 2.0 / 3.0
    dt = (2 * np.pi / Omega_D) / 600  # 每个驱动周期600步
    F_D = np.arange(1.35, 1.5, 0.001)
    total_steps = 600000
    sample_start = 500 * 600  # 跳过前500个驱动周期的暂态
    sample_end = 1000 * 600   # 采样后500个驱动周期的稳态

    # 存储采样结果
    theta_samples = []
    fd_values = []

    for f_d in F_D:
        # 初始化当前FD下的状态
        theta = 0.2
        omega = 0.0

        # 迭代计算，跳过暂态并采样稳态
        for i in range(total_steps):
            t = i * dt
            # Runge-Kutta四阶计算
            k1_theta = dt * omega
            k1_omega = dt * ((-g/l)*np.sin(theta) - q*omega + f_d*np.sin(Omega_D*t))
            
            k2_theta = dt * (omega + 0.5*k1_omega)
            k2_omega = dt * ((-g/l)*np.sin(theta + 0.5*k1_theta) - q*(omega + 0.5*k1_omega) + f_d*np.sin(Omega_D*(t + 0.5*dt)))
            
            k3_theta = dt * (omega + 0.5*k2_omega)
            k3_omega = dt * ((-g/l)*np.sin(theta + 0.5*k2_theta) - q*(omega + 0.5*k2_omega) + f_d*np.sin(Omega_D*(t + 0.5*dt)))
            
            k4_theta = dt * (omega + k3_omega)
            k4_omega = dt * ((-g/l)*np.sin(theta + k3_theta) - q*(omega + k3_omega) + f_d*np.sin(Omega_D*(t + dt)))
            
            # 更新状态
            theta += (k1_theta + 2*k2_theta + 2*k3_theta + k4_theta)/6
            omega += (k1_omega + 2*k2_omega + 2*k3_omega + k4_omega)/6
            
            # 角度归一化到[-π, π]
            theta = np.mod(theta + np.pi, 2*np.pi) - np.pi

            # 仅在采样周期记录数据
            if i >= sample_start and i < sample_end and i % 600 == 0:
                theta_samples.append(theta)
                fd_values.append(f_d)

    # 彩色绘图
    plt.figure(figsize=(10,6))
    scatter = plt.scatter(fd_values, theta_samples, c=fd_values, cmap='viridis', s=1)
    plt.title(r'Bifurcation Diagram: $\theta$ versus $F_D$', fontsize=12)
    plt.xlabel(r'$F_D$', fontsize=10)
    plt.ylabel(r'$\theta$ (radians)', fontsize=10)
    plt.xlim(1.35, 1.5)
    plt.ylim(-np.pi, np.pi)  # 归一化后显示完整范围，也可保留原1-3范围
    plt.colorbar(scatter, label=r'$F_D$')
    plt.tight_layout()
    plt.show()

if __name__ == '__main__':
    generate_bifurcation_diagram()

优化效果说明 #

1. 性能提升：

   • 移除了大量列表append操作，改用单个变量跟踪状态，减少内存操作开销
   • 跳过暂态时不再存储所有中间数据，仅在采样点记录结果，内存占用大幅降低
   • 用np.mod简化角度归一化，比原循环判断更快

2. 结构简化：

   • 去掉冗余的类结构，改用函数式编程，逻辑更清晰
   • 参数集中定义，便于修改和维护

3. 彩色绘图：

   • 使用plt.scatter的c参数将F_D值映射为颜色，配合颜色条可直观观察分岔随驱动强度的变化
   • 选择viridis配色方案，色彩均匀且对色盲友好

内容的提问来源于stack exchange，提问作者user86346