first

pc-1

three_body_problem/run_example.py

@@ -0,0 +1,199 @@
+#!/usr/bin/env python3
+"""
+三体问题求解器 - 快速示例
+运行此文件以查看三体问题的基本模拟
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import sys
+import os
+
+# 添加当前目录到路径
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from three_body_problem import ThreeBodySolver, Particle, ThreeBodyConfig, ThreeBodyVisualizer
+
+
+def simple_example():
+    """简单示例：自定义三体系统"""
+    print("=" * 60)
+    print("三体问题求解器 - 简单示例")
+    print("=" * 60)
+    
+    # 创建三个质点
+    print("\n1. 创建三个质点...")
+    particles = [
+        Particle(mass=1.0, position=[1.0, 0.0, 0.0], velocity=[0.0, 0.5, 0.0], name="Star A", color="red"),
+        Particle(mass=1.0, position=[-1.0, 0.0, 0.0], velocity=[0.0, -0.5, 0.0], name="Star B", color="green"),
+        Particle(mass=0.5, position=[0.0, 1.0, 0.0], velocity=[-0.3, 0.0, 0.0], name="Star C", color="blue")
+    ]
+    
+    for i, p in enumerate(particles):
+        print(f"   质点 {i+1}: {p.name}, 质量: {p.mass:.2f}, 位置: [{p.position[0]:.2f}, {p.position[1]:.2f}, {p.position[2]:.2f}]")
+    
+    # 创建求解器
+    print("\n2. 创建求解器...")
+    solver = ThreeBodySolver(particles, dt=0.001)
+    
+    # 模拟2年
+    print("\n3. 模拟2年运动...")
+    print("   开始积分...")
+    solver.simulate(total_time=2.0, progress_interval=500)
+    
+    # 分析结果
+    print("\n4. 分析结果...")
+    trajectories = solver.get_trajectories()
+    print(f"   模拟完成! 生成了 {len(trajectories[0])} 个轨迹点")
+    
+    com = solver.get_center_of_mass()
+    print(f"   系统质心位置: [{com[0]:.4f}, {com[1]:.4f}, {com[2]:.4f}] AU")
+    
+    momentum_error, angular_momentum_error, energy_error = solver.get_conservation_errors()
+    print(f"   守恒定律误差:")
+    print(f"     - 动量误差: {momentum_error:.2e}")
+    print(f"     - 角动量误差: {angular_momentum_error:.2e}")
+    print(f"     - 能量相对误差: {energy_error:.2e}")
+    
+    # 可视化
+    print("\n5. 生成可视化图形...")
+    visualizer = ThreeBodyVisualizer(figsize=(12, 8))
+    
+    # 创建3D轨迹图
+    visualizer.create_3d_plot()
+    visualizer.plot_trajectories(solver, title="三体系统运动轨迹")
+    
+    # 保存图形
+    output_file = "simple_example_3d.png"
+    visualizer.save_figure(output_file, dpi=200)
+    print(f"   3D轨迹图已保存到: {output_file}")
+    
+    # 创建2D投影图
+    fig, ax = visualizer.plot_2d_projection(solver, projection='xy', title="XY平面投影")
+    output_file_2d = "simple_example_2d.png"
+    fig.savefig(output_file_2d, dpi=200, bbox_inches='tight')
+    print(f"   2D投影图已保存到: {output_file_2d}")
+    
+    print("\n6. 显示图形...")
+    plt.show()
+    
+    return solver
+
+
+def figure8_example():
+    """8字形轨道示例"""
+    print("\n" + "=" * 60)
+    print("8字形轨道示例")
+    print("=" * 60)
+    
+    # 使用预置的8字形轨道配置
+    particles = ThreeBodyConfig.create_figure8_config()
+    
+    # 打印配置信息
+    ThreeBodyConfig.print_config_summary(particles)
+    
+    # 创建求解器
+    solver = ThreeBodySolver(particles, dt=0.001)
+    
+    # 模拟5年
+    print("\n模拟5年8字形轨道...")
+    solver.simulate(total_time=5.0, progress_interval=1000)
+    
+    # 可视化
+    visualizer = ThreeBodyVisualizer(figsize=(10, 8))
+    visualizer.create_3d_plot()
+    visualizer.plot_trajectories(solver, title="8字形三体轨道")
+    
+    output_file = "figure8_example.png"
+    visualizer.save_figure(output_file, dpi=200)
+    print(f"\n图形已保存到: {output_file}")
+    
+    plt.show()
+    
+    return solver
+
+
+def quick_test():
+    """快速测试所有模块"""
+    print("=" * 60)
+    print("快速测试所有模块")
+    print("=" * 60)
+    
+    # 测试1: 创建质点
+    print("\n1. 测试质点创建...")
+    p1 = Particle(mass=1.0, position=[1, 0, 0], velocity=[0, 1, 0], name="Test Star")
+    print(f"   创建质点: {p1.name}, 质量: {p1.mass}, 位置: {p1.position}")
+    
+    # 测试2: 创建配置
+    print("\n2. 测试配置创建...")
+    particles = ThreeBodyConfig.create_figure8_config()
+    print(f"   创建了 {len(particles)} 个质点")
+    
+    # 测试3: 创建求解器
+    print("\n3. 测试求解器创建...")
+    solver = ThreeBodySolver(particles, dt=0.01)
+    print(f"   求解器创建成功，时间步长: {solver.dt}")
+    
+    # 测试4: 单步积分
+    print("\n4. 测试单步积分...")
+    initial_positions = [p.position.copy() for p in particles]
+    solver.step()
+    for i, (old_pos, particle) in enumerate(zip(initial_positions, solver.particles)):
+        moved = not np.allclose(old_pos, particle.position)
+        print(f"   质点 {i+1} 位置变化: {'是' if moved else '否'}")
+    
+    # 测试5: 质心计算
+    print("\n5. 测试质心计算...")
+    com = solver.get_center_of_mass()
+    print(f"   系统质心: [{com[0]:.4f}, {com[1]:.4f}, {com[2]:.4f}]")
+    
+    # 测试6: 能量计算
+    print("\n6. 测试能量计算...")
+    energy = solver._calculate_total_energy()
+    print(f"   系统总能量: {energy:.6e}")
+    
+    print("\n所有测试通过!")
+    return True
+
+
+def main():
+    """主函数"""
+    print("三体问题求解器示例")
+    print("=" * 60)
+    print("选择要运行的示例:")
+    print("  1. 简单示例 (自定义三体系统)")
+    print("  2. 8字形轨道示例")
+    print("  3. 快速测试所有模块")
+    print("  4. 全部运行")
+    
+    try:
+        choice = input("\n请输入选择 (1-4): ").strip()
+        
+        if choice == "1":
+            simple_example()
+        elif choice == "2":
+            figure8_example()
+        elif choice == "3":
+            quick_test()
+        elif choice == "4":
+            quick_test()
+            simple_example()
+            figure8_example()
+        else:
+            print("无效选择，运行简单示例...")
+            simple_example()
+            
+    except KeyboardInterrupt:
+        print("\n\n用户中断")
+    except Exception as e:
+        print(f"\n错误: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    print("\n" + "=" * 60)
+    print("示例运行完成!")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()