first

pc-1

three_body_problem/visualizer.py

@@ -0,0 +1,321 @@
+"""
+三体问题可视化模块
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+from typing import List, Optional, Tuple
+from .particle import Particle
+from .solver import ThreeBodySolver
+
+
+class ThreeBodyVisualizer:
+    """三体问题可视化类"""
+    
+    def __init__(self, figsize: Tuple[int, int] = (12, 10)):
+        """
+        初始化可视化器
+        
+        参数:
+            figsize: 图形大小
+        """
+        self.figsize = figsize
+        self.fig = None
+        self.ax = None
+        
+    def create_3d_plot(self):
+        """创建3D图形"""
+        self.fig = plt.figure(figsize=self.figsize)
+        self.ax = self.fig.add_subplot(111, projection='3d')
+        
+    def plot_trajectories(self, solver: ThreeBodySolver, 
+                         show_current_positions: bool = True,
+                         show_com: bool = True,
+                         title: str = "三体问题轨迹"):
+        """
+        绘制三体运动轨迹
+        
+        参数:
+            solver: 三体问题求解器
+            show_current_positions: 是否显示当前位置
+            show_com: 是否显示质心
+            title: 图形标题
+        """
+        if self.ax is None:
+            self.create_3d_plot()
+        
+        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}"
+            
+            # 绘制轨迹线
+            self.ax.plot(traj[:, 0], traj[:, 1], traj[:, 2], 
+                        color=color, alpha=0.7, linewidth=1.5, label=label)
+            
+            # 绘制当前位置
+            if show_current_positions and len(traj) > 0:
+                current_pos = traj[-1]
+                self.ax.scatter(current_pos[0], current_pos[1], current_pos[2], 
+                              color=color, s=100, edgecolors='black', linewidth=1.5)
+        
+        # 绘制质心
+        if show_com:
+            com = solver.get_center_of_mass()
+            self.ax.scatter(com[0], com[1], com[2], 
+                          color='black', marker='x', s=200, label='质心', linewidth=2)
+        
+        # 设置图形属性
+        self.ax.set_xlabel('X (AU)', fontsize=12)
+        self.ax.set_ylabel('Y (AU)', fontsize=12)
+        self.ax.set_zlabel('Z (AU)', fontsize=12)
+        self.ax.set_title(title, fontsize=14, fontweight='bold')
+        self.ax.legend(fontsize=10)
+        self.ax.grid(True, alpha=0.3)
+        
+        # 设置等比例坐标轴
+        self.ax.set_aspect('auto')
+        
+    def plot_2d_projection(self, solver: ThreeBodySolver, 
+                          projection: str = 'xy',
+                          show_current_positions: bool = True,
+                          show_com: bool = True,
+                          title: str = "三体问题轨迹 (2D投影)"):
+        """
+        绘制2D投影
+        
+        参数:
+            solver: 三体问题求解器
+            projection: 投影平面 ('xy', 'xz', 'yz')
+            show_current_positions: 是否显示当前位置
+            show_com: 是否显示质心
+            title: 图形标题
+        """
+        if projection not in ['xy', 'xz', 'yz']:
+            raise ValueError("projection 必须是 'xy', 'xz' 或 'yz'")
+        
+        fig, ax = plt.subplots(figsize=(10, 8))
+        trajectories = solver.get_trajectories()
+        
+        # 确定坐标轴索引
+        if projection == 'xy':
+            x_idx, y_idx = 0, 1
+            x_label, y_label = 'X (AU)', 'Y (AU)'
+        elif projection == 'xz':
+            x_idx, y_idx = 0, 2
+            x_label, y_label = 'X (AU)', 'Z (AU)'
+        else:  # 'yz'
+            x_idx, y_idx = 1, 2
+            x_label, y_label = 'Y (AU)', 'Z (AU)'
+        
+        # 绘制每个质点的轨迹
+        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}"
+            
+            # 绘制轨迹线
+            ax.plot(traj[:, x_idx], traj[:, y_idx], 
+                   color=color, alpha=0.7, linewidth=1.5, label=label)
+            
+            # 绘制当前位置
+            if show_current_positions and len(traj) > 0:
+                current_pos = traj[-1]
+                ax.scatter(current_pos[x_idx], current_pos[y_idx], 
+                          color=color, s=100, edgecolors='black', linewidth=1.5)
+        
+        # 绘制质心
+        if show_com:
+            com = solver.get_center_of_mass()
+            ax.scatter(com[x_idx], com[y_idx], 
+                      color='black', marker='x', s=200, label='质心', linewidth=2)
+        
+        # 设置图形属性
+        ax.set_xlabel(x_label, fontsize=12)
+        ax.set_ylabel(y_label, fontsize=12)
+        ax.set_title(title, fontsize=14, fontweight='bold')
+        ax.legend(fontsize=10)
+        ax.grid(True, alpha=0.3)
+        ax.set_aspect('equal', adjustable='box')
+        
+        return fig, ax
+    
+    def plot_energy_conservation(self, solver: ThreeBodySolver, 
+                                time_points: Optional[np.ndarray] = None):
+        """
+        绘制能量守恒图
+        
+        参数:
+            solver: 三体问题求解器
+            time_points: 时间点数组（如果为None，则使用模拟时间）
+        """
+        # 从历史记录中提取能量信息
+        # 注意：这里需要修改求解器以记录能量历史
+        # 暂时返回一个简单的图形
+        fig, ax = plt.subplots(figsize=(10, 6))
+        
+        # 这里可以添加能量历史记录功能
+        ax.set_xlabel('时间 (年)', fontsize=12)
+        ax.set_ylabel('总能量', fontsize=12)
+        ax.set_title('能量守恒', fontsize=14, fontweight='bold')
+        ax.grid(True, alpha=0.3)
+        
+        return fig, ax
+    
+    def plot_phase_space(self, solver: ThreeBodySolver, 
+                        particle_index: int = 0,
+                        dimension: str = 'x'):
+        """
+        绘制相空间图（位置 vs 速度）
+        
+        参数:
+            solver: 三体问题求解器
+            particle_index: 质点索引 (0, 1, 2)
+            dimension: 维度 ('x', 'y', 'z')
+        """
+        if particle_index not in [0, 1, 2]:
+            raise ValueError("particle_index 必须是 0, 1, 或 2")
+        
+        if dimension not in ['x', 'y', 'z']:
+            raise ValueError("dimension 必须是 'x', 'y', 或 'z'")
+        
+        particle = solver.particles[particle_index]
+        trajectories = particle.get_trajectory()
+        
+        if len(trajectories) < 2:
+            print("警告：轨迹数据不足，无法绘制相空间图")
+            return None, None
+        
+        # 获取位置和速度历史
+        positions = trajectories[:, {'x': 0, 'y': 1, 'z': 2}[dimension]]
+        
+        # 注意：速度历史需要从求解器中获取
+        # 暂时使用位置差分近似速度
+        if len(positions) > 1:
+            dt = solver.dt
+            velocities = np.gradient(positions, dt)
+        else:
+            velocities = np.zeros_like(positions)
+        
+        fig, ax = plt.subplots(figsize=(10, 6))
+        
+        # 绘制相空间轨迹
+        scatter = ax.scatter(positions, velocities, c=np.arange(len(positions)), 
+                           cmap='viridis', alpha=0.7, s=20)
+        
+        # 添加颜色条表示时间
+        plt.colorbar(scatter, ax=ax, label='时间步')
+        
+        ax.set_xlabel(f'{dimension.upper()} 位置 (AU)', fontsize=12)
+        ax.set_ylabel(f'{dimension.upper()} 速度 (AU/年)', fontsize=12)
+        ax.set_title(f'质点 {particle_index+1} 的相空间图 ({dimension.upper()}维度)', 
+                    fontsize=14, fontweight='bold')
+        ax.grid(True, alpha=0.3)
+        
+        return fig, ax
+    
+    def animate_trajectories(self, solver: ThreeBodySolver, 
+                           save_path: Optional[str] = None,
+                           fps: int = 30,
+                           dpi: int = 100):
+        """
+        创建轨迹动画（需要额外安装 matplotlib.animation）
+        
+        参数:
+            solver: 三体问题求解器
+            save_path: 保存路径（如果为None则显示动画）
+            fps: 帧率
+            dpi: 分辨率
+        """
+        try:
+            from matplotlib.animation import FuncAnimation
+        except ImportError:
+            print("错误：需要安装 matplotlib.animation 模块")
+            return
+        
+        trajectories = solver.get_trajectories()
+        if any(len(traj) < 2 for traj in trajectories):
+            print("错误：轨迹数据不足，无法创建动画")
+            return
+        
+        fig = plt.figure(figsize=(10, 8))
+        ax = fig.add_subplot(111, projection='3d')
+        
+        # 初始化图形元素
+        lines = []
+        points = []
+        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}"
+            
+            # 创建轨迹线
+            line, = ax.plot([], [], [], color=color, alpha=0.7, linewidth=1.5, label=label)
+            lines.append(line)
+            
+            # 创建当前位置点
+            point, = ax.plot([], [], [], 'o', color=color, markersize=8, 
+                           markeredgecolor='black', markeredgewidth=1)
+            points.append(point)
+        
+        # 设置坐标轴范围
+        all_points = np.vstack(trajectories)
+        max_range = np.max([all_points[:, i].max() - all_points[:, i].min() for i in range(3)])
+        mid_x = (all_points[:, 0].max() + all_points[:, 0].min()) * 0.5
+        mid_y = (all_points[:, 1].max() + all_points[:, 1].min()) * 0.5
+        mid_z = (all_points[:, 2].max() + all_points[:, 2].min()) * 0.5
+        
+        ax.set_xlim(mid_x - max_range/2, mid_x + max_range/2)
+        ax.set_ylim(mid_y - max_range/2, mid_y + max_range/2)
+        ax.set_zlim(mid_z - max_range/2, mid_z + max_range/2)
+        
+        ax.set_xlabel('X (AU)', fontsize=12)
+        ax.set_ylabel('Y (AU)', fontsize=12)
+        ax.set_zlabel('Z (AU)', fontsize=12)
+        ax.set_title('三体问题动画', fontsize=14, fontweight='bold')
+        ax.legend(fontsize=10)
+        ax.grid(True, alpha=0.3)
+        
+        # 动画更新函数
+        def update(frame):
+            for i, (traj, line, point) in enumerate(zip(trajectories, lines, points)):
+                # 更新轨迹线（显示到当前帧）
+                line.set_data(traj[:frame+1, 0], traj[:frame+1, 1])
+                line.set_3d_properties(traj[:frame+1, 2])
+                
+                # 更新当前位置点
+                if frame < len(traj):
+                    point.set_data([traj[frame, 0]], [traj[frame, 1]])
+                    point.set_3d_properties([traj[frame, 2]])
+            
+            return lines + points
+        
+        # 创建动画
+        anim = FuncAnimation(fig, update, frames=min(len(traj) for traj in trajectories), 
+                           interval=1000/fps, blit=True)
+        
+        if save_path:
+            anim.save(save_path, writer='ffmpeg', fps=fps, dpi=dpi)
+            print(f"动画已保存到: {save_path}")
+        else:
+            plt.show()
+        
+        return anim
+    
+    def show(self):
+        """显示所有图形"""
+        plt.tight_layout()
+        plt.show()
+        
+    def save_figure(self, filename: str, dpi: int = 300):
+        """保存当前图形"""
+        if self.fig is not None:
+            self.fig.savefig(filename, dpi=dpi, bbox_inches='tight')
+            print(f"图形已保存到: {filename}")
+        else:
+            print("警告：没有活动的图形可以保存")