first

pc-1

three_body_problem/particle.py

@@ -0,0 +1,64 @@
+"""
+质点类，表示三体问题中的一个天体
+"""
+
+import numpy as np
+from typing import Tuple, List
+
+
+class Particle:
+    """表示一个质点的类"""
+    
+    def __init__(self, mass: float, position: np.ndarray, velocity: np.ndarray, 
+                 name: str = "", color: str = None):
+        """
+        初始化质点
+        
+        参数:
+            mass: 质量（太阳质量）
+            position: 位置向量 (x, y, z) [AU]
+            velocity: 速度向量 (vx, vy, vz) [AU/年]
+            name: 质点名称
+            color: 可视化颜色
+        """
+        self.mass = mass
+        self.position = np.array(position, dtype=np.float64)
+        self.velocity = np.array(velocity, dtype=np.float64)
+        self.name = name
+        self.color = color
+        
+        # 存储轨迹历史
+        self.position_history = []
+        self.velocity_history = []
+        
+    def update(self, new_position: np.ndarray, new_velocity: np.ndarray):
+        """更新质点的状态并记录历史"""
+        self.position_history.append(self.position.copy())
+        self.velocity_history.append(self.velocity.copy())
+        self.position = new_position
+        self.velocity = new_velocity
+    
+    def get_trajectory(self) -> np.ndarray:
+        """获取轨迹历史"""
+        if not self.position_history:
+            return np.array([self.position])
+        return np.array(self.position_history + [self.position])
+    
+    def get_energy(self) -> float:
+        """计算质点的动能"""
+        v_squared = np.sum(self.velocity**2)
+        return 0.5 * self.mass * v_squared
+    
+    def __repr__(self) -> str:
+        return (f"Particle(name='{self.name}', mass={self.mass:.3f}, "
+                f"position={self.position}, velocity={self.velocity})")
+    
+    def copy(self) -> 'Particle':
+        """创建质点的副本"""
+        return Particle(
+            mass=self.mass,
+            position=self.position.copy(),
+            velocity=self.velocity.copy(),
+            name=self.name,
+            color=self.color
+        )