N体模拟器开发遇ValueError：数组广播维度不匹配问题求助

问题描述 #

开发N体模拟器，输入初始位置与速度，输出随时间变化的位置。编写了微分方程函数拟传入solve_ivp求解，但计算dvelocity_dt（加速度）时触发广播错误：

ValueError: could not broadcast input array from shape (10,3) into shape (3,)

已尝试将G和mass转为[10,1]数组，问题仍存在。

背景代码 #

import scipy.integrate as integrate
from scipy.integrate import solve_ivp
import numpy as np
import random

# 控制模拟模式
mode = 'galaxy'

if mode == 'solar':
    G = 6.6e-11  # N m^2/kg^2
    N_bodies = 4
    # 四大天体初始位置、速度、质量
    sun_r0 = np.array([0.0, 0.0, 0.0])
    sun_v0 = np.array([0.0, 0.0, 0.0])
    sun_M = 2e30  # kg
    
    earth_r0 = np.array([149.6e9, 0.0, 0.0])  # m
    earth_v0 = np.array([0.0, 29.8e3, 0.0])  # m/s
    earth_M = 5.97e24  # kg
    
    venus_r0 = np.array([108e9, 0.0, 0.0])
    venus_v0 = np.array([0.0, 35e3, 0.0])
    venus_M = 5e24
    
    mars_r0 = np.array([230e9, 0.0, 0.0])
    mars_v0 = np.array([0.0, 24e3, 0.0])
    mars_M = 6.4e23

    a = 1e-20
    t_max = 31557600.0*10  # 模拟时长：10年
    vis_radius = 300e9  # 可视化范围

    # 初始状态数组：[位置数组, 速度数组]
    y_initial = np.array([
        [sun_r0, earth_r0, venus_r0, mars_r0],
        [sun_v0, earth_v0, venus_v0, mars_v0]
    ])
    mass = np.array([sun_M, earth_M, venus_M, mars_M])

elif mode == 'galaxy':
    G = 4.49e-18  # kpc^3/M_solar/kyear^2
    R = 50  # kpc
    z_scale_factor = 0.25
    N_bodies = 10
    M_bodies = 1e8  # M_solar
    rc = 1  # kpc
    a = rc/(2**(2/5)-1)**(1/2)
    t_max = 1000000
    mass = np.full(N_bodies, M_bodies)
    vis_radius = 50
    
    # 生成初始位置
    r_initial = np.zeros((N_bodies,3))
    v_initial = np.zeros((N_bodies,3))
    
    for i in range(N_bodies):
        ran_thi = random.random()
        ran_theta = random.random()
        ran_r = random.random()
        thi = 2*np.pi*ran_thi
        theta = np.arccos(1-(2*ran_theta))
        r = R*ran_r
        r_1 = r*np.sin(theta)*np.cos(thi)
        r_2 = r*np.sin(theta)*np.sin(thi)
        r_3 = 0.25*r*np.cos(theta)
        r_initial[i] = [r_1, r_2, r_3]

    # 初始状态数组
    y_initial = np.array([
        r_initial,
        v_initial
    ])
else:
    print("Invalid mode")
print(y_initial)

出错的微分方程函数代码 #

# 扁平化初始值
y_initial_flat = y_initial.flatten()
t_eval_point = np.linspace(0, t_max, 200)
t_span = (0, t_max)
rs, vs = y_initial_flat.reshape(y_initial.shape)

# 状态解包函数
def get_rs_vs(y_flat):
    rs, vs = y_flat.reshape(y_initial.shape)
    return rs, vs

# 微分方程函数
def dy_dt(t, y):
    position , velocity = get_rs_vs(y_initial)  # 错误：应该用传入的y，而非初始值
    dvelocity_dt = np.zeros_like(velocity)

    if mode == 'solar':
        dposition_dt = velocity

    if mode == 'galaxy':
        G_array = np.full(N_bodies, G)
        for i in range(N_bodies):
            dvelocity_dt[i] = (
                dvelocity_dt[i] +
                (G_array*mass).reshape((10,1))*(position - position[i])/(np.full(N_bodies,(((np.linalg.norm(position - position[i]))**2 + a**2)**(3/2)))).reshape((10,1))
            )
        
        # 计算速度大小
        velocity = np.sqrt(np.linalg.norm(dvelocity_dt)*np.linalg.norm(position))

        # 计算v_x和v_y
        v_x = np.zeros(N_bodies)
        v_y = np.zeros(N_bodies)
        v_n = np.zeros((N_bodies,3))
        for i in range(N_bodies):
            v_x[i] = np.sqrt(velocity/(1+(dvelocity_dt[i,0]**2/dvelocity_dt[i,1]**2)))
            v_y[i] = -v_x[i]*(dvelocity_dt[i,0]/dvelocity_dt[i,1])
            v_n[i] = [v_x[i], v_y[i], 0]

报错信息 #

ValueError                                Traceback (most recent call last)
Cell In[156], line 10
      8 G_array = np.full(N_bodies,G)
      9 for i in range(0,N_bodies):
---> 10     dvelocity_dt[i] = (
     11     dvelocity_dt[i]+
     12     (G_array*mass).reshape((10,1))*(position - position[i])/(np.full(N_bodies,(((np.linalg.norm(position - position[i]))**2 + a**2)**(3/2)))).reshape((10,1)))
     14 #Calculate the magnitude of velocity (Part B)
     15 velocity = np.sqrt(np.linalg.norm(dvelocity_dt)*np.linalg.norm(position))

ValueError: could not broadcast input array from shape (10,3) into shape (3,)

错误原因与修复方案 #

核心错误点 #

1. 状态参数误用：dy_dt函数里用了y_initial而非传入的y，导致始终用初始位置计算，而非当前时间步的位置。
2. 形状不匹配：计算第i个天体的加速度时，position - position[i]得到(10,3)的数组，直接赋值给(3,)的dvelocity_dt[i]，广播失败。正确逻辑是对所有其他天体的引力加速度求和，得到单个(3,)的加速度向量。
3. 未处理自引力项：position - position[i]包含j=i的情况（距离为0），会导致除零问题，需要排除或置零。

修复后的微分方程函数 #

# 扁平化初始值
y_initial_flat = y_initial.flatten()
t_eval_point = np.linspace(0, t_max, 200)
t_span = (0, t_max)

# 状态解包函数
def get_rs_vs(y_flat):
    rs, vs = y_flat.reshape(y_initial.shape)
    return rs, vs

# 微分方程函数
def dy_dt(t, y):
    position, velocity = get_rs_vs(y)  # 改用传入的当前状态y
    dposition_dt = velocity  # 位置导数就是速度，两种模式通用
    dvelocity_dt = np.zeros_like(velocity)

    if mode == 'galaxy' or mode == 'solar':
        G_val = G
        for i in range(N_bodies):
            # 计算所有天体到i的位置差
            r_ij = position - position[i]  # (N_bodies, 3)
            # 计算距离（带软参数a避免除零）
            dist_ij = np.sqrt(np.sum(r_ij**2, axis=1) + a**2)  # (N_bodies,)
            # 计算每个天体对i的引力加速度，j=i时r_ij为0，该项自动为0
            acc_ij = G_val * mass[:, np.newaxis] * r_ij / dist_ij[:, np.newaxis]**3  # (N_bodies, 3)
            # 求和得到i的总加速度
            dvelocity_dt[i] = np.sum(acc_ij, axis=0)  # (3,)
    
    # 把导数扁平化，符合solve_ivp的要求
    dy = np.concatenate([dposition_dt.flatten(), dvelocity_dt.flatten()])
    return dy

额外说明 #

• 修复了状态参数的问题，确保用当前时间步的y计算导数。
• 通过求和得到单个天体的加速度，解决形状不匹配问题。
• 用np.sum(r_ij**2, axis=1)计算距离，结合软参数a避免j=i时的除零错误。
• 最后将位置和速度的导数扁平化后返回，符合solve_ivp对输入输出的一维数组要求。
• 原代码中关于速度计算的部分（Part B）需要结合实际需求调整，当前修复先保证微分方程的正确性。

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