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

"""
拉格朗日点示例 - 三体问题的稳定点
"""

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_lagrange_example(lagrange_point: int = 4, total_time: float = 100.0):
    """
    运行拉格朗日点示例
    
    参数:
        lagrange_point: 拉格朗日点编号 (4=L4, 5=L5)
        total_time: 总模拟时间（年）
    """
    point_name = "L4" if lagrange_point == 4 else "L5"
    print("=" * 60)
    print(f"拉格朗日点 {point_name} 示例")
    print("=" * 60)
    
    # 创建拉格朗日点配置
    particles = ThreeBodyConfig.create_lagrange_point_config(lagrange_point=lagrange_point)
    
    # 打印配置摘要
    ThreeBodyConfig.print_config_summary(particles)
    
    # 创建求解器
    dt = 0.001  # 时间步长（年）
    solver = ThreeBodySolver(particles, dt=dt)
    
    print(f"\n开始模拟，总时间: {total_time}年，时间步长: {dt}年")
    print(f"模拟拉格朗日点 {point_name} 的稳定性")
    
    solver.simulate(total_time=total_time, progress_interval=20000)
    
    # 计算守恒误差
    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}")
    
    # 分析测试质点的轨道稳定性
    test_particle = solver.particles[2]  # 测试质点是第三个
    trajectory = test_particle.get_trajectory()
    
    # 计算与L4/L5点的距离变化
    if lagrange_point == 4:
        lagrange_position = np.array([0.5, np.sqrt(3)/2, 0.0])
    else:  # L5
        lagrange_position = np.array([0.5, -np.sqrt(3)/2, 0.0])
    
    distances = np.linalg.norm(trajectory - lagrange_position, axis=1)
    time_points = np.arange(len(distances)) * dt
    
    print(f"\n测试质点稳定性分析:")
    print(f"  初始距离L{lagrange_point}点: {distances[0]:.6e} AU")
    print(f"  最终距离L{lagrange_point}点: {distances[-1]:.6e} AU")
    print(f"  最大距离偏差: {np.max(distances):.6e} AU")
    print(f"  平均距离偏差: {np.mean(distances):.6e} AU")
    
    # 可视化
    print("\n生成可视化图形...")
    
    # 创建图形
    fig = plt.figure(figsize=(16, 10))
    
    # 1. XY平面轨迹
    ax1 = plt.subplot(2, 3, 1)
    
    # 绘制所有质点的轨迹
    trajectories = solver.get_trajectories()
    colors = ['gold', 'blue', 'gray']
    
    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}"
        
        # 只绘制最后一部分轨迹（更清晰）
        if len(traj) > 1000:
            traj_to_plot = traj[-1000:]
        else:
            traj_to_plot = traj
            
        ax1.plot(traj_to_plot[:, 0], traj_to_plot[:, 1], 
                color=color, alpha=0.7, linewidth=1.5, label=label)
        
        # 绘制最终位置
        ax1.scatter(traj[-1, 0], traj[-1, 1], 
                   color=color, s=100, edgecolors='black', linewidth=1.5, zorder=5)
    
    # 绘制拉格朗日点位置
    ax1.scatter(lagrange_position[0], lagrange_position[1], 
               color='red', marker='*', s=300, label=f'L{lagrange_point}点', zorder=10)
    
    # 绘制等边三角形
    triangle_points = [
        [0, 0],  # 太阳
        [1, 0],  # 地球
        lagrange_position[:2]  # L4或L5点
    ]
    triangle_points.append(triangle_points[0])  # 闭合三角形
    triangle_points = np.array(triangle_points)
    ax1.plot(triangle_points[:, 0], triangle_points[:, 1], 
            'k--', alpha=0.5, linewidth=1, label='等边三角形')
    
    ax1.set_xlabel('X (AU)', fontsize=12)
    ax1.set_ylabel('Y (AU)', fontsize=12)
    ax1.set_title(f'拉格朗日点 {point_name} - XY平面', fontsize=14, fontweight='bold')
    ax1.legend(fontsize=10)
    ax1.grid(True, alpha=0.3)
    ax1.set_aspect('equal', adjustable='box')
    
    # 2. 距离随时间变化
    ax2 = plt.subplot(2, 3, 2)
    ax2.plot(time_points, distances, 'b-', linewidth=2, alpha=0.8)
    ax2.set_xlabel('时间 (年)', fontsize=12)
    ax2.set_ylabel(f'距离L{lagrange_point}点 (AU)', fontsize=12)
    ax2.set_title('测试质点轨道稳定性', fontsize=14, fontweight='bold')
    ax2.grid(True, alpha=0.3)
    
    # 添加平均距离线
    mean_distance = np.mean(distances)
    ax2.axhline(y=mean_distance, color='r', linestyle='--', alpha=0.7, 
               label=f'平均距离: {mean_distance:.3e}')
    ax2.legend(fontsize=10)
    
    # 3. 相空间图 (x vs vx)
    ax3 = plt.subplot(2, 3, 3)
    
    # 计算速度（使用位置差分）
    if len(trajectory) > 1:
        dt = solver.dt
        velocities = np.gradient(trajectory, dt, axis=0)
        x_positions = trajectory[:, 0]
        x_velocities = velocities[:, 0]
        
        # 使用颜色表示时间
        scatter = ax3.scatter(x_positions, x_velocities, c=time_points, 
                            cmap='viridis', alpha=0.7, s=20)
        plt.colorbar(scatter, ax=ax3, label='时间 (年)')
        
        ax3.set_xlabel('X 位置 (AU)', fontsize=12)
        ax3.set_ylabel('X 速度 (AU/年)', fontsize=12)
        ax3.set_title('测试质点相空间 (X维度)', fontsize=14, fontweight='bold')
        ax3.grid(True, alpha=0.3)
    
    # 4. 相对位置图（以地球为参考系）
    ax4 = plt.subplot(2, 3, 4)
    
    # 计算相对于地球的位置
    earth_trajectory = trajectories[1]  # 地球是第二个质点
    sun_trajectory = trajectories[0]    # 太阳是第一个质点
    test_trajectory = trajectories[2]   # 测试质点是第三个
    
    # 转换为以地球为中心的坐标系
    earth_centered_sun = sun_trajectory - earth_trajectory
    earth_centered_test = test_trajectory - earth_trajectory
    
    # 只绘制最后一部分
    if len(earth_centered_test) > 1000:
        earth_centered_test = earth_centered_test[-1000:]
    
    ax4.plot(earth_centered_test[:, 0], earth_centered_test[:, 1], 
            'gray', alpha=0.7, linewidth=1.5, label='测试质点')
    ax4.scatter(0, 0, color='blue', s=200, label='地球', edgecolors='black', linewidth=1.5)
    ax4.scatter(earth_centered_sun[-1, 0], earth_centered_sun[-1, 1], 
               color='gold', s=200, label='太阳', edgecolors='black', linewidth=1.5)
    
    # 绘制理论L4/L5点位置
    if lagrange_point == 4:
        l_point_relative = np.array([-0.5, np.sqrt(3)/2])
    else:  # L5
        l_point_relative = np.array([-0.5, -np.sqrt(3)/2])
    
    ax4.scatter(l_point_relative[0], l_point_relative[1], 
               color='red', marker='*', s=300, label=f'L{lagrange_point}点', zorder=10)
    
    ax4.set_xlabel('相对X位置 (AU)', fontsize=12)
    ax4.set_ylabel('相对Y位置 (AU)', fontsize=12)
    ax4.set_title('以地球为参考系', fontsize=14, fontweight='bold')
    ax4.legend(fontsize=10)
    ax4.grid(True, alpha=0.3)
    ax4.set_aspect('equal', adjustable='box')
    
    # 5. 能量随时间变化（简化）
    ax5 = plt.subplot(2, 3, 5)
    
    # 计算相对能量变化（简化）
    # 在实际实现中，需要记录能量历史
    time_array = np.linspace(0, total_time, len(distances))
    # 使用距离变化作为能量变化的代理
    energy_proxy = distances / distances[0]
    
    ax5.plot(time_array, energy_proxy, 'g-', linewidth=2, alpha=0.8)
    ax5.set_xlabel('时间 (年)', fontsize=12)
    ax5.set_ylabel('相对能量变化', fontsize=12)
    ax5.set_title('轨道能量变化', fontsize=14, fontweight='bold')
    ax5.grid(True, alpha=0.3)
    ax5.axhline(y=1.0, color='r', linestyle='--', alpha=0.5, label='初始能量')
    ax5.legend(fontsize=10)
    
    # 6. 3D视图
    ax6 = plt.subplot(2, 3, 6, projection='3d')
    
    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}"
        
        # 只绘制最后一部分轨迹
        if len(traj) > 1000:
            traj_to_plot = traj[-1000:]
        else:
            traj_to_plot = traj
            
        ax6.plot(traj_to_plot[:, 0], traj_to_plot[:, 1], traj_to_plot[:, 2], 
                color=color, alpha=0.7, linewidth=1.5, label=label)
        
        # 绘制最终位置
        ax6.scatter(traj[-1, 0], traj[-1, 1], traj[-1, 2], 
                   color=color, s=100, edgecolors='black', linewidth=1.5, zorder=5)
    
    ax6.scatter(lagrange_position[0], lagrange_position[1], lagrange_position[2], 
               color='red', marker='*', s=300, label=f'L{lagrange_point}点', zorder=10)
    
    ax6.set_xlabel('X (AU)', fontsize=10)
    ax6.set_ylabel('Y (AU)', fontsize=10)
    ax6.set_zlabel('Z (AU)', fontsize=10)
    ax6.set_title('3D视图', fontsize=14, fontweight='bold')
    ax6.legend(fontsize=9, loc='upper left')
    ax6.grid(True, alpha=0.3)
    
    plt.suptitle(f'拉格朗日点 {point_name} 稳定性分析', fontsize=16, fontweight='bold')
    plt.tight_layout()
    
    # 保存图形
    output_file = f"lagrange_point_{point_name}.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"\n图形已保存到: {output_file}")
    
    # 显示图形
    plt.show()
    
    return solver


def compare_lagrange_points():
    """比较L4和L5点的稳定性"""
    print("\n" + "=" * 60)
    print("拉格朗日点L4和L5稳定性比较")
    print("=" * 60)
    
    total_time = 50.0
    dt = 0.001
    
    results = []
    
    for lagrange_point in [4, 5]:
        point_name = f"L{lagrange_point}"
        print(f"\n模拟 {point_name} 点...")
        
        particles = ThreeBodyConfig.create_lagrange_point_config(lagrange_point=lagrange_point)
        solver = ThreeBodySolver([p.copy() for p in particles], dt=dt)
        solver.simulate(total_time=total_time, progress_interval=25000)
        
        # 分析测试质点的轨道稳定性
        test_particle = solver.particles[2]
        trajectory = test_particle.get_trajectory()
        
        if lagrange_point == 4:
            lagrange_position = np.array([0.5, np.sqrt(3)/2, 0.0])
        else:  # L5
            lagrange_position = np.array([0.5, -np.sqrt(3)/2, 0.0])
        
        distances = np.linalg.norm(trajectory - lagrange_position, axis=1)
        
        results.append({
            'point': point_name,
            'max_distance': np.max(distances),
            'mean_distance': np.mean(distances),
            'std_distance': np.std(distances),
            'final_distance': distances[-1]
        })
        
        print(f"  {point_name} 最大距离偏差: {np.max(distances):.6e} AU")
        print(f"  {point_name} 平均距离偏差: {np.mean(distances):.6e} AU")
    
    # 绘制比较图
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    
    points = [r['point'] for r in results]
    max_distances = [r['max_distance'] for r in results]
    mean_distances = [r['mean_distance'] for r in results]
    
    x = np.arange(len(points))
    width = 0.35
    
    axes[0].bar(x - width/2, max_distances, width, label='最大偏差', color='lightcoral')
    axes[0].bar(x + width/2, mean_distances, width, label='平均偏差', color='lightblue')
    axes[0].set_xlabel('拉格朗日点', fontsize=12)
    axes[0].set_ylabel('距离偏差 (AU)', fontsize=12)
    axes[0].set_title('L4和L5点稳定性比较', fontsize=14, fontweight='bold')
    axes[0].set_xticks(x)
    axes[0].set_xticklabels(points)
    axes[0].legend()
    axes[0].grid(True, alpha=0.3, axis='y')
    
    # 最终位置偏差
    final_distances = [r['final_distance'] for r in results]
    axes[1].bar(points, final_distances, color=['lightgreen', 'lightblue'])
    axes[1].set_xlabel('拉格朗日点', fontsize=12)
    axes[1].set_ylabel('最终距离偏差 (AU)', fontsize=12)
    axes[1].set_title('最终位置稳定性', fontsize=14, fontweight='bold')
    axes[1].grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    
    output_file = "lagrange_points_comparison.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"\n比较图形已保存到: {output_file}")
    plt.show()


if __name__ == "__main__":
    # 运行L4点示例
    print("运行拉格朗日点L4示例...")
    solver_l4 = run_lagrange_example(lagrange_point=4, total_time=50.0)
    
    # 运行L5点示例
    print("\n" + "="*60)
    print("运行拉格朗日点L5示例...")
    solver_l5 = run_lagrange_example(lagrange_point=5, total_time=50.0)
    
    # 比较L4和L5
    compare_lagrange_points()