太阳系引力模拟

丁宇泽

[Python] 纯文本查看 复制代码

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.widgets import Slider, Button, CheckButtons

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义常量
G = 6.67430e-11  # 万有引力常数

class CelestialBody:
    """天体类，包含质量、位置、速度等属性，以及更新方法"""
    def __init__(self, mass, position, velocity, color='b', size=5, name=''):
        self.mass = mass
        self.position = np.array(position, dtype=float)
        self.velocity = np.array(velocity, dtype=float)
        self.color = color
        self.size = size
        self.name = name
        self.trajectory = []
        
    def update(self, others, dt):
        """根据其他天体的引力更新自身位置和速度"""
        total_force = np.zeros(2)
        for other in others:
            if other is not self:
                r = other.position - self.position
                distance = np.linalg.norm(r)
                # 避免距离过小时引力过大的问题
                if distance > 1e6:
                    force = G * self.mass * other.mass * r / distance**3
                    total_force += force
        acceleration = total_force / self.mass
        self.velocity += acceleration * dt
        self.position += self.velocity * dt
        self.trajectory.append(self.position.copy())

def create_solar_system():
    """创建太阳系模型"""
    sun = CelestialBody(
        mass=1.989e30, position=[0, 0], 
        velocity=[0, 0], color='yellow', size=20, name='太阳'
    )
    
    mercury = CelestialBody(
        mass=3.301e23, position=[57.9e9, 0], 
        velocity=[0, 47.4e3], color='gray', size=3, name='水星'
    )
    
    venus = CelestialBody(
        mass=4.867e24, position=[108.2e9, 0], 
        velocity=[0, 35.0e3], color='orange', size=5, name='金星'
    )
    
    earth = CelestialBody(
        mass=5.972e24, position=[149.6e9, 0], 
        velocity=[0, 29.8e3], color='blue', size=5, name='地球'
    )
    
    mars = CelestialBody(
        mass=6.417e23, position=[227.9e9, 0], 
        velocity=[0, 24.1e3], color='red', size=4, name='火星'
    )
    
    return [sun, mercury, venus, earth, mars]

def init_animation():
    """初始化动画"""
    global lines, points, time_text, trail_length
    lines = []
    points = []
    
    for body in bodies:
        # 创建轨迹线
        line, = ax.plot([], [], lw=1, color=body.color, alpha=0.7)
        lines.append(line)
        
        # 创建表示天体的点
        point, = ax.plot([], [], 'o', ms=body.size, color=body.color, label=body.name)
        points.append(point)
    
    # 创建时间文本
    time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)
    
    # 初始化滑块和按钮
    init_ui()
    
    return lines + points + [time_text]

def init_ui():
    """初始化用户界面控件"""
    global play_pause_button, reset_button, trail_check, speed_slider, trail_slider
    
    # 调整主图位置，为控件留出空间
    plt.subplots_adjust(bottom=0.3, left=0.15)
    
    # 创建播放/暂停按钮
    ax_play = plt.axes([0.15, 0.22, 0.1, 0.04])
    play_pause_button = Button(ax_play, '暂停')
    
    # 创建重置按钮
    ax_reset = plt.axes([0.3, 0.22, 0.1, 0.04])
    reset_button = Button(ax_reset, '重置')
    
    # 创建轨迹显示复选框
    ax_trail = plt.axes([0.5, 0.22, 0.15, 0.04])
    trail_check = CheckButtons(ax_trail, ['显示轨迹'], [True])
    
    # 创建速度滑块
    ax_speed = plt.axes([0.15, 0.15, 0.7, 0.03])
    speed_slider = Slider(ax_speed, '模拟速度', 0.1, 10.0, valinit=1.0)
    
    # 创建轨迹长度滑块
    ax_trail_length = plt.axes([0.15, 0.1, 0.7, 0.03])
    trail_slider = Slider(ax_trail_length, '轨迹长度', 10, 1000, valinit=300)
    
    # 绑定按钮和滑块事件
    play_pause_button.on_clicked(toggle_animation)
    reset_button.on_clicked(reset_simulation)
    trail_check.on_clicked(toggle_trail)
    speed_slider.on_changed(update_speed)
    trail_slider.on_changed(update_trail_length)

def update(frame):
    """更新每一帧的动画"""
    global current_step, animation_speed, trail_length, paused
    
    if not paused:
        # 更新所有天体的位置
        for _ in range(int(animation_speed)):
            if current_step < num_steps:
                for body in bodies:
                    body.update(bodies, dt)
                current_step += 1
    
    # 更新绘图
    for i, body in enumerate(bodies):
        # 获取轨迹数据
        if show_trail:
            # 限制轨迹显示长度
            trail_data = body.trajectory[-int(trail_length):]
            x_trail = [pos[0] for pos in trail_data]
            y_trail = [pos[1] for pos in trail_data]
            lines[i].set_data(x_trail, y_trail)
        
        # 更新当前位置
        points[i].set_data(body.position[0], body.position[1])
    
    # 更新时间文本
    time_years = current_step * dt / (365 * 24 * 3600)
    time_text.set_text(f'模拟时间: {time_years:.2f} 年\n时间步长: {dt:.2f} 秒')
    
    return lines + points + [time_text]

def toggle_animation(event):
    """切换动画的播放/暂停状态"""
    global paused
    paused = not paused
    play_pause_button.label.set_text('播放' if paused else '暂停')

def reset_simulation(event):
    """重置模拟"""
    global current_step, bodies, initial_bodies
    
    # 重置当前步数
    current_step = 0
    
    # 重置所有天体的状态
    bodies = []
    for body in initial_bodies:
        new_body = CelestialBody(
            body.mass, body.position.copy(), 
            body.velocity.copy(), body.color, 
            body.size, body.name
        )
        new_body.trajectory = []
        bodies.append(new_body)
    
    # 更新初始状态
    initial_bodies = bodies.copy()

def toggle_trail(label):
    """切换轨迹显示状态"""
    global show_trail
    show_trail = not show_trail

def update_speed(val):
    """更新模拟速度"""
    global animation_speed
    animation_speed = val

def update_trail_length(val):
    """更新轨迹显示长度"""
    global trail_length
    trail_length = val

# 主程序
if __name__ == "__main__":
    # 创建太阳系模型
    bodies = create_solar_system()
    
    # 保存初始状态以便重置
    initial_bodies = bodies.copy()
    
    # 模拟参数
    dt = 86400  # 时间步长（秒），默认一天
    num_steps = 10000  # 总模拟步数
    current_step = 0
    animation_speed = 1.0  # 动画速度倍数
    trail_length = 300  # 轨迹显示长度
    show_trail = True  # 是否显示轨迹
    paused = False  # 是否暂停
    
    # 创建图形
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111)
    
    # 设置坐标轴范围
    max_distance = max(max(abs(body.position[0]), abs(body.position[1])) for body in bodies)
    ax.set_xlim(-max_distance * 1.5, max_distance * 1.5)
    ax.set_ylim(-max_distance * 1.5, max_distance * 1.5)
    
    # 设置坐标轴标签和标题
    ax.set_xlabel('X (米)')
    ax.set_ylabel('Y (米)')
    ax.set_title('太阳系引力模拟')
    ax.grid(True, alpha=0.3)
    
    # 添加图例
    ax.legend(loc='upper right')
    
    # 创建动画
    ani = FuncAnimation(
        fig, update, frames=range(num_steps),
        init_func=init_animation, blit=True, interval=20, repeat=False
    )
    
    plt.show()