Python绘制散点图：两体运动计算代码相关技术咨询

我来帮你搞定这段两体运动代码的问题，同时实现散点图绘制。先梳理下你当前代码里的几个关键问题，再给出完整的可运行版本：

代码里的核心问题 #

• 参数被覆盖：你定义的twomasses函数接收x1,x2,p参数，但函数内部直接赋值了x0_1 = [1,2]、x0_2 = [4,5]、p=3，导致传入的参数完全无效，建议把初始值改为函数参数或默认参数。
• 状态变量不匹配：你的ruku_step函数期望状态更新函数是单参数的，但外部的F(x1,x2)需要两个位置参数。两体问题的状态应该是合并后的向量（包含两个质点的位置和速度），这样才能用Runge-Kutta统一更新。
• 未完成的lambda函数：末尾的f = lambda y: ...没有实现，这是连接受力计算和状态更新的核心部分，必须补全状态导数的逻辑。
• 全局变量冗余：使用global定义gamma,m1,m2没必要，直接在函数内部处理更规范，避免污染全局命名空间。

修正后的完整代码 #

import numpy as np
import matplotlib.pyplot as plt

def twomasses(M1, M2, x0_1, x0_2, p, h, n):
    gamma = 1  # 引力常数，可按需调整
    m1 = M1
    m2 = M2
    
    # 初始化速度：根据给定动量p分配，这里假设初始动量沿y轴
    v1 = [0, p / m1]
    v2 = [0, -p / m2]
    
    # 合并状态变量：[x1x, x1y, v1x, v1y, x2x, x2y, v2x, v2y]
    y = np.array(x0_1 + v1 + x0_2 + v2, dtype=np.float64)
    
    def compute_forces(y):
        # 从状态向量中提取两个质点的位置
        x1 = y[:2]
        x2 = y[4:6]
        r = x2 - x1
        r_norm = np.linalg.norm(r, 2)
        # 计算万有引力：F = γ*m1*m2*(r)/|r|³
        force = (gamma * m1 * m2 / (r_norm ** 3)) * r
        return force
    
    def state_derivative(y):
        # 提取两个质点的速度
        v1 = y[2:4]
        v2 = y[6:8]
        force = compute_forces(y)
        # 计算状态导数：位置的导数是速度，速度的导数是加速度
        dx1_dt = v1
        dv1_dt = force / m1
        dx2_dt = v2
        dv2_dt = -force / m2  # 第二个质点受力与第一个大小相等、方向相反
        # 合并所有导数为一个向量
        return np.concatenate([dx1_dt, dv1_dt, dx2_dt, dv2_dt])
    
    def ruku_step(F, y, h):
        # 四阶Runge-Kutta迭代步骤
        k1 = F(y)
        k2 = F(y + (h/2)*k1)
        k3 = F(y + (h/2)*k2)
        k4 = F(y + h*k3)
        y_new = y + (h/6)*(k1 + 2*k2 + 2*k3 + k4)
        return y_new
    
    # 初始化轨迹存储数组
    traj1 = np.zeros((n+1, 2))
    traj2 = np.zeros((n+1, 2))
    traj1[0] = y[:2]
    traj2[0] = y[4:6]
    
    # 迭代计算每一步的位置
    for i in range(n):
        y = ruku_step(state_derivative, y, h)
        traj1[i+1] = y[:2]
        traj2[i+1] = y[4:6]
    
    return traj1, traj2

# 示例参数配置
m1 = 1.0
m2 = 0.5
initial_pos1 = [1.0, 2.0]
initial_pos2 = [4.0, 5.0]
initial_momentum = 3.0
step_size = 0.01
total_steps = 1000

# 计算两个质点的运动轨迹
traj1, traj2 = twomasses(m1, m2, initial_pos1, initial_pos2, initial_momentum, step_size, total_steps)

# 绘制散点图
plt.figure(figsize=(8, 8))
plt.scatter(traj1[:, 0], traj1[:, 1], s=5, label=f'质点1 (质量={m1})', alpha=0.6)
plt.scatter(traj2[:, 0], traj2[:, 1], s=5, label=f'质点2 (质量={m2})', alpha=0.6)
plt.xlabel('X坐标')
plt.ylabel('Y坐标')
plt.title('两体运动轨迹散点图')
plt.legend()
plt.grid(True)
plt.axis('equal')  # 保持坐标轴比例一致，避免轨迹变形
plt.show()

关键说明 #

• 我们将两个质点的位置和速度合并成一个8维状态向量，让Runge-Kutta方法可以统一处理整个系统的更新。
• state_derivative函数负责计算状态的变化率：位置的导数就是速度，速度的导数由引力加速度决定（引力除以质量）。
• 用traj1和traj2数组存储每一步的位置数据，最后通过plt.scatter绘制散点图展示轨迹。
• 添加了plt.axis('equal')来保证坐标轴比例一致，避免天体运动轨迹因为坐标轴拉伸而变形。

内容的提问来源于stack exchange，提问作者undergrad