first

pc-1

three_body_problem/solver.py

@@ -0,0 +1,196 @@
+"""
+三体问题求解器主模块
+"""
+
+import numpy as np
+from typing import List, Tuple, Optional, Callable
+from .particle import Particle
+from .integrator import RK4Integrator
+
+
+class ThreeBodySolver:
+    """三体问题求解器"""
+    
+    # 万有引力常数 (AU^3 / (M_sun * year^2))
+    G = 4 * np.pi**2  # 在天文单位制中
+    
+    def __init__(self, particles: List[Particle], dt: float = 0.001):
+        """
+        初始化三体问题求解器
+        
+        参数:
+            particles: 三个质点的列表
+            dt: 时间步长（年）
+        """
+        if len(particles) != 3:
+            raise ValueError("三体问题需要恰好3个质点")
+        
+        self.particles = particles
+        self.dt = dt
+        self.integrator = RK4Integrator(dt)
+        self.time = 0.0
+        
+        # 验证总动量和角动量守恒（可选）
+        self.initial_total_momentum = self._calculate_total_momentum()
+        self.initial_angular_momentum = self._calculate_total_angular_momentum()
+    
+    def _calculate_accelerations(self, particles: List[Particle]) -> List[np.ndarray]:
+        """
+        计算所有质点的加速度
+        
+        参数:
+            particles: 质点列表
+            
+        返回:
+            加速度列表
+        """
+        accelerations = [np.zeros(3) for _ in range(3)]
+        
+        # 计算每对质点之间的引力
+        for i in range(3):
+            for j in range(3):
+                if i != j:
+                    # 计算相对位置向量
+                    r_vec = particles[j].position - particles[i].position
+                    r = np.linalg.norm(r_vec)
+                    
+                    # 避免除以零
+                    if r < 1e-10:
+                        r = 1e-10
+                    
+                    # 计算加速度 (F = ma -> a = F/m)
+                    acceleration = self.G * particles[j].mass * r_vec / (r**3)
+                    accelerations[i] += acceleration
+        
+        return accelerations
+    
+    def _calculate_total_momentum(self) -> np.ndarray:
+        """计算系统总动量"""
+        total_momentum = np.zeros(3)
+        for p in self.particles:
+            total_momentum += p.mass * p.velocity
+        return total_momentum
+    
+    def _calculate_total_angular_momentum(self) -> np.ndarray:
+        """计算系统总角动量"""
+        total_angular_momentum = np.zeros(3)
+        for p in self.particles:
+            # 角动量 L = r × p = r × (m*v)
+            angular_momentum = np.cross(p.position, p.mass * p.velocity)
+            total_angular_momentum += angular_momentum
+        return total_angular_momentum
+    
+    def _calculate_total_energy(self) -> float:
+        """计算系统总能量（动能+势能）"""
+        kinetic_energy = 0.0
+        potential_energy = 0.0
+        
+        # 计算动能
+        for p in self.particles:
+            kinetic_energy += 0.5 * p.mass * np.sum(p.velocity**2)
+        
+        # 计算势能
+        for i in range(3):
+            for j in range(i+1, 3):
+                r_vec = self.particles[j].position - self.particles[i].position
+                r = np.linalg.norm(r_vec)
+                if r > 1e-10:  # 避免除以零
+                    potential_energy -= self.G * self.particles[i].mass * self.particles[j].mass / r
+        
+        return kinetic_energy + potential_energy
+    
+    def step(self) -> List[Particle]:
+        """
+        执行一个时间步的积分
+        
+        返回:
+            更新后的质点列表
+        """
+        # 计算加速度
+        acceleration_func = lambda p: self._calculate_accelerations(p)
+        
+        # 执行积分
+        self.particles = self.integrator.step(self.particles, acceleration_func)
+        self.time += self.dt
+        
+        return self.particles
+    
+    def simulate(self, total_time: float, progress_interval: int = 1000) -> List[Particle]:
+        """
+        模拟三体运动
+        
+        参数:
+            total_time: 总模拟时间（年）
+            progress_interval: 进度打印间隔步数
+            
+        返回:
+            模拟结束后的质点列表
+        """
+        steps = int(total_time / self.dt)
+        
+        print(f"开始模拟: 总时间={total_time:.2f}年, 步数={steps}, 步长={self.dt:.6f}年")
+        print(f"初始总能量: {self._calculate_total_energy():.6e}")
+        
+        acceleration_func = lambda p: self._calculate_accelerations(p)
+        
+        for step in range(steps):
+            self.particles = self.integrator.step(self.particles, acceleration_func)
+            self.time += self.dt
+            
+            # 打印进度
+            if step % progress_interval == 0:
+                energy = self._calculate_total_energy()
+                print(f"进度: {step}/{steps}步, 时间={self.time:.3f}年, 能量={energy:.6e}")
+        
+        final_energy = self._calculate_total_energy()
+        print(f"模拟完成: 最终时间={self.time:.2f}年, 最终能量={final_energy:.6e}")
+        
+        return self.particles
+    
+    def get_trajectories(self) -> List[np.ndarray]:
+        """获取所有质点的轨迹"""
+        return [p.get_trajectory() for p in self.particles]
+    
+    def get_center_of_mass(self) -> np.ndarray:
+        """计算系统质心位置"""
+        total_mass = sum(p.mass for p in self.particles)
+        com = np.zeros(3)
+        for p in self.particles:
+            com += p.mass * p.position
+        return com / total_mass if total_mass > 0 else np.zeros(3)
+    
+    def get_conservation_errors(self) -> Tuple[float, float, float]:
+        """
+        计算守恒定律的误差
+        
+        返回:
+            (动量误差, 角动量误差, 能量误差)
+        """
+        # 动量误差
+        current_momentum = self._calculate_total_momentum()
+        momentum_error = np.linalg.norm(current_momentum - self.initial_total_momentum)
+        
+        # 角动量误差
+        current_angular_momentum = self._calculate_total_angular_momentum()
+        angular_momentum_error = np.linalg.norm(
+            current_angular_momentum - self.initial_angular_momentum
+        )
+        
+        # 能量误差（相对误差）
+        initial_energy = self._calculate_total_energy()  # 重新计算初始能量
+        current_energy = self._calculate_total_energy()
+        if abs(initial_energy) > 1e-10:
+            energy_error = abs((current_energy - initial_energy) / initial_energy)
+        else:
+            energy_error = abs(current_energy - initial_energy)
+        
+        return momentum_error, angular_momentum_error, energy_error
+    
+    def reset(self):
+        """重置求解器状态（清除历史记录）"""
+        for p in self.particles:
+            p.position_history.clear()
+            p.velocity_history.clear()
+        self.time = 0.0
+        self.initial_total_momentum = self._calculate_total_momentum()
+        self.initial_angular_momentum = self._calculate_total_angular_momentum()