three-body-problem/three_body_problem/examples/figure8.py

"""
8字形轨道示例 - 著名的稳定三体轨道
"""

import numpy as np
import matplotlib.pyplot as plt
import sys
import os

# 添加父目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from three_body_problem import ThreeBodySolver, ThreeBodyConfig, ThreeBodyVisualizer


def run_figure8_example():
    """运行8字形轨道示例"""
    print("=" * 60)
    print("8字形轨道示例")
    print("=" * 60)
    
    # 创建8字形轨道配置
    particles = ThreeBodyConfig.create_figure8_config()
    
    # 打印配置摘要
    ThreeBodyConfig.print_config_summary(particles)
    
    # 创建求解器
    dt = 0.001  # 时间步长（年）
    solver = ThreeBodySolver(particles, dt=dt)
    
    # 模拟10年
    total_time = 10.0
    print(f"\n开始模拟，总时间: {total_time}年，时间步长: {dt}年")
    
    solver.simulate(total_time=total_time, progress_interval=2000)
    
    # 计算守恒误差
    momentum_error, angular_momentum_error, energy_error = solver.get_conservation_errors()
    print(f"\n守恒定律误差:")
    print(f"  动量误差: {momentum_error:.6e}")
    print(f"  角动量误差: {angular_momentum_error:.6e}")
    print(f"  能量相对误差: {energy_error:.6e}")
    
    # 可视化
    print("\n生成可视化图形...")
    visualizer = ThreeBodyVisualizer(figsize=(14, 10))
    
    # 创建3D轨迹图
    plt.figure(figsize=(14, 10))
    
    # 3D轨迹
    ax1 = plt.subplot(2, 2, 1, projection='3d')
    trajectories = solver.get_trajectories()
    colors = ['red', 'green', 'blue']
    
    for i, (traj, particle) in enumerate(zip(trajectories, solver.particles)):
        color = particle.color if particle.color else colors[i % len(colors)]
        label = particle.name if particle.name else f"质点 {i+1}"
        ax1.plot(traj[:, 0], traj[:, 1], traj[:, 2], 
                color=color, alpha=0.7, linewidth=1.5, label=label)
        ax1.scatter(traj[-1, 0], traj[-1, 1], traj[-1, 2], 
                   color=color, s=100, edgecolors='black', linewidth=1.5)
    
    # 绘制质心
    com = solver.get_center_of_mass()
    ax1.scatter(com[0], com[1], com[2], 
               color='black', marker='x', s=200, label='质心', linewidth=2)
    
    ax1.set_xlabel('X (AU)')
    ax1.set_ylabel('Y (AU)')
    ax1.set_zlabel('Z (AU)')
    ax1.set_title('8字形轨道 - 3D视图')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # XY平面投影
    ax2 = plt.subplot(2, 2, 2)
    for i, (traj, particle) in enumerate(zip(trajectories, solver.particles)):
        color = particle.color if particle.color else colors[i % len(colors)]
        label = particle.name if particle.name else f"质点 {i+1}"
        ax2.plot(traj[:, 0], traj[:, 1], color=color, alpha=0.7, linewidth=1.5, label=label)
        ax2.scatter(traj[-1, 0], traj[-1, 1], color=color, s=100, edgecolors='black', linewidth=1.5)
    
    ax2.scatter(com[0], com[1], color='black', marker='x', s=200, label='质心', linewidth=2)
    ax2.set_xlabel('X (AU)')
    ax2.set_ylabel('Y (AU)')
    ax2.set_title('XY平面投影')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    ax2.set_aspect('equal', adjustable='box')
    
    # XZ平面投影
    ax3 = plt.subplot(2, 2, 3)
    for i, (traj, particle) in enumerate(zip(trajectories, solver.particles)):
        color = particle.color if particle.color else colors[i % len(colors)]
        label = particle.name if particle.name else f"质点 {i+1}"
        ax3.plot(traj[:, 0], traj[:, 2], color=color, alpha=0.7, linewidth=1.5, label=label)
        ax3.scatter(traj[-1, 0], traj[-1, 2], color=color, s=100, edgecolors='black', linewidth=1.5)
    
    ax3.scatter(com[0], com[2], color='black', marker='x', s=200, label='质心', linewidth=2)
    ax3.set_xlabel('X (AU)')
    ax3.set_ylabel('Z (AU)')
    ax3.set_title('XZ平面投影')
    ax3.legend()
    ax3.grid(True, alpha=0.3)
    ax3.set_aspect('equal', adjustable='box')
    
    # YZ平面投影
    ax4 = plt.subplot(2, 2, 4)
    for i, (traj, particle) in enumerate(zip(trajectories, solver.particles)):
        color = particle.color if particle.color else colors[i % len(colors)]
        label = particle.name if particle.name else f"质点 {i+1}"
        ax4.plot(traj[:, 1], traj[:, 2], color=color, alpha=0.7, linewidth=1.5, label=label)
        ax4.scatter(traj[-1, 1], traj[-1, 2], color=color, s=100, edgecolors='black', linewidth=1.5)
    
    ax4.scatter(com[1], com[2], color='black', marker='x', s=200, label='质心', linewidth=2)
    ax4.set_xlabel('Y (AU)')
    ax4.set_ylabel('Z (AU)')
    ax4.set_title('YZ平面投影')
    ax4.legend()
    ax4.grid(True, alpha=0.3)
    ax4.set_aspect('equal', adjustable='box')
    
    plt.suptitle('8字形三体轨道', fontsize=16, fontweight='bold')
    plt.tight_layout()
    
    # 保存图形
    output_file = "figure8_orbit.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"\n图形已保存到: {output_file}")
    
    # 显示图形
    plt.show()
    
    return solver


def analyze_figure8_stability():
    """分析8字形轨道的稳定性"""
    print("\n" + "=" * 60)
    print("8字形轨道稳定性分析")
    print("=" * 60)
    
    # 创建8字形轨道配置
    particles = ThreeBodyConfig.create_figure8_config()
    
    # 测试不同时间步长
    time_steps = [0.01, 0.005, 0.001, 0.0005]
    total_time = 5.0
    
    results = []
    
    for dt in time_steps:
        print(f"\n测试时间步长: {dt}")
        
        solver = ThreeBodySolver([p.copy() for p in particles], dt=dt)
        solver.simulate(total_time=total_time, progress_interval=10000)
        
        # 计算守恒误差
        momentum_error, angular_momentum_error, energy_error = solver.get_conservation_errors()
        
        results.append({
            'dt': dt,
            'momentum_error': momentum_error,
            'angular_momentum_error': angular_momentum_error,
            'energy_error': energy_error
        })
        
        print(f"  能量相对误差: {energy_error:.6e}")
    
    # 绘制误差随步长变化
    plt.figure(figsize=(10, 6))
    
    dts = [r['dt'] for r in results]
    energy_errors = [r['energy_error'] for r in results]
    
    plt.loglog(dts, energy_errors, 'o-', linewidth=2, markersize=8)
    plt.xlabel('时间步长 (年)', fontsize=12)
    plt.ylabel('能量相对误差', fontsize=12)
    plt.title('8字形轨道数值误差分析', fontsize=14, fontweight='bold')
    plt.grid(True, alpha=0.3, which='both')
    
    # 添加参考线（四阶精度）
    ref_dt = np.array(dts)
    ref_error = 1e-4 * (ref_dt / 0.001)**4
    plt.loglog(ref_dt, ref_error, 'r--', alpha=0.7, label='四阶精度参考线')
    
    plt.legend()
    plt.tight_layout()
    
    output_file = "figure8_stability_analysis.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"\n稳定性分析图形已保存到: {output_file}")
    plt.show()


if __name__ == "__main__":
    # 运行主示例
    solver = run_figure8_example()
    
    # 运行稳定性分析（可选）
    # analyze_figure8_stability()