Satellite orbit change model

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Circle as mpl_Circle

class RK4:
    def __init__(self, ode, step_size):
        self.ode = ode
        self.step_size = step_size
        self.k1 = np.zeros_like(ode.state)
        self.k2 = np.zeros_like(ode.state)
        self.k3 = np.zeros_like(ode.state)
        self.k4 = np.zeros_like(ode.state)
        self.temp_state = np.zeros_like(ode.state)

    def step(self):
        h = self.step_size
        state = self.ode.state

        self.ode.get_rate(state, self.k1)

        self.temp_state[:] = state + h/2 * self.k1
        self.ode.get_rate(self.temp_state, self.k2)

        self.temp_state[:] = state + h/2 * self.k2
        self.ode.get_rate(self.temp_state, self.k3)

        self.temp_state[:] = state + h * self.k3
        self.ode.get_rate(self.temp_state, self.k4)

        self.ode.state += h/6 * (self.k1 + 2*self.k2 + 2*self.k3 + self.k4)

class Planet:
    GM = 4 * np.pi**2

    def __init__(self):
        self.state = np.zeros(5)
        self.trail_x = []
        self.trail_y = []
        self.circle = mpl_Circle((0, 0), 0.1, color='blue', alpha=0.5)
        self.ode_solver = None
        self.energy_data = []
        self.last_y = None
        self.last_period_time = None
        self.periods = []
        self.area_rate = []
        self.impulse_strength = 0.1  # 默认冲量强度
        self.speed_multiplier = 1.1  # 速度变化倍率


        # 新增极值检测相关变量
        self.max_speed_points = []  # 存储速度最大点 (x, y, t, speed)
        self.min_speed_points = []  # 存储速度最小点
        self.prev_state = None      # 前一步状态
        self.prev_r_dot = None      # 前一步的r_dot值

    def initialize(self, init_state, dt):
        self.state = np.array(init_state)
        self.trail_x = [self.state[0]]
        self.trail_y = [self.state[2]]
        self.circle.center = (self.state[0], self.state[2])
        self.ode_solver = RK4(self, dt)
        self.last_y = self.state[2]
        self.last_period_time = None
        self.periods = []
        self.max_speed_points = []
        self.min_speed_points = []
        self.prev_state = None
        self.prev_r_dot = None
        self.area_rate = []


    def get_rate(self, state, rate):
        x, vx, y, vy, t = state
        r = np.sqrt(x**2 + y**2)
        r3 = r ** 3
        ax = (-self.GM * x) / r3
        ay = (-self.GM * y) / r3
        rate[0] = vx
        rate[1] = ax
        rate[2] = vy
        rate[3] = ay
        rate[4] = 1

    def do_step(self):
        if self.ode_solver:
            self.ode_solver.step()
            self.energy_data.append((self.state[4], self.compute_energy()))

        self.trail_x.append(self.state[0])
        self.trail_y.append(self.state[2])
        self.circle.center = (self.state[0], self.state[2])
        self.period()
        self.detect_speed_extremes()
        self.compute_area()

    def apply_tangential_boost(self, multiplier):
        """应用切向速度冲量"""
        x, vx, y, vy, t = self.state
        r = np.sqrt(x**2 + y**2)
        if r < 1e-3:
            return

        # 计算切向方向单位向量
        tangential_dir = np.array([-y/r, x/r])

        # 计算当前切向速度分量
        current_tangent_speed = vx * tangential_dir[0] + vy * tangential_dir[1]

        # 施加冲量后的新速度
        new_vx = vx + (multiplier-1)*current_tangent_speed*tangential_dir[0]
        new_vy = vy + (multiplier-1)*current_tangent_speed*tangential_dir[1]

        self.state[1] = new_vx
        self.state[3] = new_vy
        self.last_period_time = None
        self.periods = []
        self.last_y = y

    def detect_speed_extremes(self):
        x, vx, y, vy, t = self.state
        current_r_dot = x * vx + y * vy  # 计算r_dot（位置矢量与速度矢量的点积）
        if self.prev_r_dot is not None:
            # 检测r_dot符号变化，判断是否经过极值点
            if self.prev_r_dot * current_r_dot < 0:
                # 比较绝对值确定更接近极值点的状态
                if abs(self.prev_r_dot) < abs(current_r_dot):
                    x_prev, vx_prev, y_prev, vy_prev, t_prev = self.prev_state
                    speed_prev = np.sqrt(vx_prev**2 + vy_prev**2)
                    # 根据r_dot的正负判断极大或极小
                    if self.prev_r_dot > 0:
                        self.min_speed_points.append((x_prev, y_prev, t_prev, speed_prev))
                    else:
                        self.max_speed_points.append((x_prev, y_prev, t_prev, speed_prev))
                else:
                    speed = np.sqrt(vx**2 + vy**2)
                    if current_r_dot < 0:
                        self.min_speed_points.append((x, y, t, speed))
                    else:
                        self.max_speed_points.append((x, y, t, speed))
        # 保存当前状态用于下一步比较
        self.prev_state = self.state.copy()
        self.prev_r_dot = current_r_dot

        # 保存当前状态用于下一步比较
        self.prev_state = self.state.copy()
        self.prev_r_dot = current_r_dot

    def compute_energy(self):
        x, vx, y, vy, t = self.state
        r = np.sqrt(x**2 + y**2)
        return 0.5*(vx**2 + vy**2) - self.GM/r

    def compute_area(self):
        # 计算面积
        x, vx, y, vy, t = self.state
        current_area_dAdt = 0.5 * abs(x * vy - y * vx)
        self.area_rate.append((t, current_area_dAdt))
        return self.area_rate

    def period(self):
        # 周期检测
        x, _, y, _, t = self.state
        if self.last_y is not None and y * self.last_y < 0:
            if x > 0:
                if self.last_period_time is not None:
                    period = t - self.last_period_time
                    self.periods.append(period)
                    print(f"检测到周期: {period:.5f} y")  # 立即输出周期
                self.last_period_time = t
        self.last_y = y
        return self.period

    def compute_semi_major_axis(self):
        if len(self.max_speed_points) > 0 and len(self.min_speed_points) > 0:
            # 获取最近的速度极值点
            max_point = self.max_speed_points[-1]
            min_point = self.min_speed_points[-1]

            # 计算距离
            r_max = np.sqrt(max_point[0]**2 + max_point[1]**2)
            r_min = np.sqrt(min_point[0]**2 + min_point[1]**2)

            # 计算半长轴
            semi_major_axis = (r_max + r_min) / 2
            return semi_major_axis
        return None

class PlanetSimulation:
    def __init__(self):
        self.fig, (self.ax1, self.ax2, self.ax3) = plt.subplots(1, 3, figsize=(24, 12))
        self.ax1.set(xlim=(-3, 10), ylim=(-5, 5), aspect='equal',
                    xlabel="x (AU)", ylabel="y (AU)", title="Planet Simulation")
        self.ax2.set(xlabel="Time (yr)", ylabel="Energy (AU²/yr²)",
                    title="Energy vs Time", xlim=(0, 10), ylim=(-40, 0))
        self.sun = mpl_Circle((0, 0), 0.2, color='yellow', alpha=1.0)
        self.ax1.add_patch(self.sun)
        self.planet = Planet()
        self.ax1.add_patch(self.planet.circle)
        self.trail, = self.ax1.plot([], [], 'b-', lw=0.5)
        self.energy_line, = self.ax2.plot([], [], 'r-', lw=1)

        self.ax3.set(xlabel="Time (yr)", ylabel="dA/dt (AU²/yr)",
        title="Area Rate Verification",
        xlim=(0, 10), ylim=(0, 5))
        self.area_line, = self.ax3.plot([], [], 'b-', lw=1)
        self.connect_events()

        # 添加半长轴显示文本
        self.semi_major_axis_text = self.ax1.text(0.95, 0.85, '',
                                                 transform=self.ax1.transAxes,
                                                 ha='right', va='top',
                                                 fontsize=12,
                                                 bbox=dict(facecolor='white', alpha=0.5))

        # 添加周期显示文本
        self.period_text = self.ax1.text(0.95, 0.95, '',
                                        transform=self.ax1.transAxes,
                                        ha='right', va='top',
                                        fontsize=12,
                                        bbox=dict(facecolor='white', alpha=0.5))
        self.last_period_count = 0  # 跟踪已显示的周期数量

        self.max_scatter = self.ax1.scatter([], [], color='red', marker='*',
                                           s=100, label='Max Speed')
        self.min_scatter = self.ax1.scatter([], [], color='green', marker='*',
                                           s=100, label='Min Speed')
        self.ax1.legend()
        # 在AX1中添加文本显示
        self.area_text = self.ax1.text(0.95, 0.65, '',
                                    transform=self.ax1.transAxes,
                                    ha='right', va='top',
                                    fontsize=10,
                                    bbox=dict(facecolor='white', alpha=0.5))

        self.dt = 0.01
        self.steps_per_frame = 5

    def initialize(self, x, vx, y, vy, dt):
        self.planet.initialize([x, vx, y, vy, 0], dt)
        self.trail.set_data(self.planet.trail_x, self.planet.trail_y)
        self.energy_line.set_data([], [])

    def connect_events(self):
        """连接所有交互事件"""
        self.fig.canvas.mpl_connect('button_press_event', self.on_mouse_click)

    def on_mouse_click(self, event):
        """鼠标点击事件处理"""
        click_x, click_y = event.xdata, event.ydata
        self.planet.apply_tangential_boost(self.planet.speed_multiplier)

        # 更新冲量显示
        current_speed = np.sqrt(self.planet.state[1]**2 + self.planet.state[3]**2)
        self.impulse_text.set_text(f'Speed: {current_speed:.2f} AU/yr\n'+
                                      f'Last boost: {self.planet.speed_multiplier}x')

        # 显示视觉反馈
        self.ax1.plot(click_x, click_y, 'r*', markersize=15, alpha=0.5)
        self.fig.canvas.draw_idle()

    def update(self, frame):
        for _ in range(self.steps_per_frame):
            self.planet.do_step()

        self.trail.set_data(self.planet.trail_x, self.planet.trail_y)

        # 更新能量曲线
        if self.planet.energy_data:
            times, energies = zip(*self.planet.energy_data)
            self.energy_line.set_data(times, energies)
            self.ax2.set_xlim(0, max(times))
            self.ax2.set_ylim(min(energies)-1, max(energies)+1)

        # 更新极值点标记
        max_points = np.array([[p[0], p[1]] for p in self.planet.max_speed_points])
        min_points = np.array([[p[0], p[1]] for p in self.planet.min_speed_points])
        if len(max_points) > 0:
            self.max_scatter.set_offsets(max_points)
        if len(min_points) > 0:
            self.min_scatter.set_offsets(min_points)

        # 更新周期显示
        current_count = len(self.planet.periods)
        if current_count > self.last_period_count:
            latest_period = self.planet.periods[-1]
            self.period_text.set_text(f'Period: {latest_period:.2f} yr')
            self.last_period_count = current_count

        # 更新半长轴显示
        semi_major_axis = self.planet.compute_semi_major_axis()
        if semi_major_axis is not None:
            self.semi_major_axis_text.set_text(f'Semi-major Axis: {semi_major_axis:.2f} AU')

        # 更新面积率图
        if self.planet.area_rate:
            times, dAdt = zip(*self.planet.area_rate)
            self.area_line.set_data(times, dAdt)
            self.ax3.set_xlim(0, max(times))

        return [self.trail, self.planet.circle, self.energy_line,
                self.max_scatter, self.min_scatter, self.period_text,
                self.semi_major_axis_text, self.area_line]


    def run(self):
        ani = FuncAnimation(self.fig, self.update,frames=200, interval=50, blit=True)
        plt.show()


if __name__ == "__main__":
    sim = PlanetSimulation()
    sim.initialize(x=1.0, vx=0.0, y=0.0, vy=6.28, dt=0.01)
    sim.run()