帮我写一个三体运动的代码

以下是一个简单的三体运动代码：

import numpy as np
import matplotlib.pyplot as plt

# 定义常数
G = 6.67430e-11 # 万有引力常数
M = np.array([5.97e24, 3.3e23, 1.9885e30]) # 地球、月球和太阳的质量
R = np.array([6.371e6, 1.738e6, 6.96342e8]) # 地球、月球和太阳的半径
r0 = np.array([0, 384e6, 1.5e11]) # 质心坐标
v0 = np.array([0, 1023, 0]) # 初始速度

# 定义时间步长和总时间
dt = 60 # 时间步长为1分钟
t = np.arange(0, 365*24*60*60, dt) # 总时间为一年

# 初始化坐标和速度数组
r = np.zeros((len(M), len(t), 3)) # 每个星体在每个时间步长的坐标
v = np.zeros((len(M), len(t), 3)) # 每个星体在每个时间步长的速度
r[:, 0] = np.array([0, 384e6, 1.5e11]) # 初始坐标

# 定义三体运动函数
def three_body_motion(r, v, M, R, G, dt):
    a = np.zeros((len(M), 3)) # 每个星体在当前时间步长的加速度
    for i in range(len(M)):
        for j in range(len(M)):
            if i == j:
                continue
            r_ij = r[i] - r[j] # 计算星体i和星体j之间的距离向量
            a_ij = G * M[j] / np.linalg.norm(r_ij)**3 * r_ij # 计算星体i在星体j的引力加速度
            a[i] -= a_ij # 计算星体i的加速度
    v += a * dt # 更新速度
    r += v * dt # 更新坐标
    return r, v

# 通过时间循环计算三体运动
for i in range(1, len(t)):
    r[:, i], v[:, i] = three_body_motion(r[:, i-1], v[:, i-1], M, R, G, dt)

# 绘制轨迹图
plt.figure(figsize=(8, 8))
plt.plot(r[0, :, 0], r[0, :, 1], label='Earth')
plt.plot(r[1, :, 0], r[1, :, 1], label='Moon')
plt.plot(r[2, :, 0], r[2, :, 1], label='Sun')
plt.legend()
plt.xlabel('X [m]')
plt.ylabel('Y [m]')
plt.title('Three-Body Motion')
plt.show()

该代码将计算地球、月球和太阳的三体运动，将它们的轨迹显示在同一张图中。你可以根据需求修改和扩展该代码，例如添加更多星体，改变初始条件等。