Python测试粒子模拟长宽比计算维度不匹配问题求助

问题描述 #

在Python粒子模拟项目中，需要计算每个时间步所有粒子的平均长度与宽度，使长度、宽度数组形状为(19,)，但当前代码计算出的是对应单个粒子的(1000,)形状数组，执行绘图时触发以下错误：

plt.plot(t_output, ratio)
ValueError: x and y must have same first dimension, but have shapes (19,) and (1000,)

已完成质心系转换、轴旋转等步骤，但仍无法解决问题，需修正代码实现需求。

错误原因分析 #

1. 质心系转换不完整：原代码仅转换了速度到质心系，未对位置做质心系转换，导致后续旋转和长宽计算基于全局坐标系而非云团自身的质心系。
2. 旋转逻辑维度错误：原代码针对每个粒子计算旋转角，而非整个云团的整体运动方向，导致旋转后的位置维度混乱。
3. 长宽计算轴参数错误：原代码计算length和width时，错误地对粒子维度取最大最小，而非针对每个时间步的所有粒子范围计算。

修正后的代码 #

import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation
from mpl_toolkits.mplot3d import Axes3D


# Constants
M_Sun = 1.989e33  # Solar Mass in g
M_bh = 4.3e6 * M_Sun # Mass of BH
G = 6.67430e-8  # cm^3 g^(-1) s^(-2)
yr = 365 * 24 * 60 * 60  # 1 year in seconds
pc = 3.086e18 # Parsec in cm
R = .225 * 0.04 * pc / 2 # Radius of cloud

# Number of particles
num_particles = 1000

# Uniformly distributed particles in Sphere of radius R
phi = np.random.uniform(0, 2*np.pi, size=num_particles)
costheta = np.random.uniform(0, 1, size=num_particles)
u = np.random.uniform(0, 1, size=num_particles)

theta = np.arccos(costheta)
r = R * (u**(1/3))

x = r * np.sin(theta) * np.cos(phi)
y = r * np.sin(theta) * np.sin(phi)
z = r * np.cos(theta)

# Initial Conditions

RA = .425 # Right Ascension
Decl = .523 # Declination

vx = 550e5 # x-component of velocity
vy = 730e5 # y-component of velocity
vz = 0 # z-component of velocity

# Transform RA and Decl to x,y,z components (1 arcsec = 0.04 pc in the galactic center)
x_0 = RA * 0.04 * pc
y_0 = Decl * 0.04 * pc
z_0 = 0

# Empty lists for position and velocity
initial_pos = []
initial_vel = []

for i in range(num_particles):
    initial_position = (x_0 + x[i], y_0 + y[i], z_0 + z[i]) # x_0 + x (x is already multiplied by R), same for y and z axes
    initial_velocity = (vx, vy, vz) # Fixed velocity for all particles
    initial_pos.append(initial_position)
    initial_vel.append(initial_velocity)

initial_pos = np.array(initial_pos) #Shape (num_particles, coordinates)
initial_vel = np.array(initial_vel) #Shape (num_particles, coordinates)

# Time
t_end = 100*yr # Total time of integration
dt_constant = 0.1 # Magic number for time stepping scheme
intervals = 1000000 # seconds after which to store pos & vel
next_interval_time = intervals

# Initial conditions
time = np.zeros(1)
_pos_t = initial_pos.copy()
_vel_t = initial_vel.copy()

# Lists to store outputs
pos_output = []
vel_output = []
t_output = []

# Number of Stored outputs
output = 20

while time[-1] <= t_end: #We end up doing one more timestep after t_end
    r = np.linalg.norm(_pos_t, axis=1)
    acc = -G * M_bh / r[:, np.newaxis]**3 * _pos_t #np.newaxis for broadcasting with _pos_t

    # Calculate the time step for the current particle
    current_dt = dt_constant * np.sqrt(np.linalg.norm(_pos_t, axis=1)**3 / (G * M_bh))
    min_dt = np.min(current_dt) # Use the minimum time step for all particles

    # leap frog integration

    # calculate velocity at half timestep
    half_vel = _vel_t + 0.5 * acc * min_dt
    # current position only used in this timestep
    _pos_t = _pos_t + half_vel * min_dt # shape: (#particles, #coordinates)

    # Recalculate acceleration with the new position
    r = np.linalg.norm(_pos_t, axis=1) 
    # Acceleration at timestep
    acc = -G * M_bh / r[:, np.newaxis]**3 * _pos_t #np.newaxis for broadcasting with _pos_t
    # current velocity only used in this timestep
    _vel_t = half_vel + 0.5 * acc * min_dt # shape: (#particles, #coordinates)

    # time at timestep t
    _time_t = time[-1] + min_dt

    # Check if the elapsed time has surpassed the next interval
    if _time_t >= next_interval_time: # >= because time steps aren't constant and it may not be an exact multiple
        # Store data at intervals
        pos_output.append(_pos_t.copy())
        vel_output.append(_vel_t.copy())
        t_output.append(_time_t.copy())
        next_interval_time += intervals

    time = np.append(time, _time_t)

    # show current status by printing timestep number (-1 because initial conditions)
    print(f'timestep: {time.size - 1} [progress: {_time_t/t_end*100:.3f}%]')


pos_output = np.array(pos_output)  # shape: (#stored timesteps, #particles, #coordinates)
vel_output = np.array(vel_output)  # shape: (#stored timesteps, #particles, #coordinates)
t_output = np.array(t_output)

# Plotting the ratio of length to breadth as a function of time for the entire cloud 
# find the velocity vector at each timestep 
# move to COM frame
# find the angle of x or y position to the vel vector #dot product
# rotate one axis so it aligns with vel vector # use rotation matrix
# then x and y are perpendicular and one is aligned with vel vector
# then do the ratio calculation

# Move to the Center of Mass (COM) frame
# 计算每个时间步的质心位置和质心速度
com_pos = np.mean(pos_output, axis=1)  # shape: (#timesteps, 3)
com_vel = np.mean(vel_output, axis=1)  # shape: (#timesteps, 3)

# 转换位置到质心系
pos_rel_com = pos_output - com_pos[:, np.newaxis, :]  # shape: (#timesteps, #particles, 3)

# 针对每个时间步，计算质心速度的方向，生成旋转矩阵
# 我们需要将质心速度方向旋转到x轴方向
rotations = []
for vel in com_vel:
    # 处理零速度情况（避免除以零）
    if np.linalg.norm(vel) < 1e-10:
        rot = Rotation.from_euler('y', 0)
    else:
        # 计算速度向量与x轴的夹角
        vel_unit = vel / np.linalg.norm(vel)
        # 计算绕y轴的旋转角（将速度向量转到x轴）
        angle = -np.arccos(vel_unit[0])
        # 检查z分量符号，调整旋转方向
        if vel_unit[2] < 0:
            angle = -angle
        rot = Rotation.from_euler('y', angle)
    rotations.append(rot.as_matrix())

rotation_matrices = np.array(rotations)  # shape: (#timesteps, 3, 3)

# 旋转每个时间步的所有粒子位置
pos_rotated = np.einsum('tij,tpj->tpj', rotation_matrices, pos_rel_com)

# Calculate the length and width for each timestep
# 长度：x轴方向所有粒子的范围
length = np.max(pos_rotated[:, :, 0], axis=1) - np.min(pos_rotated[:, :, 0], axis=1)  # shape: (#timesteps,)
# 宽度：y轴方向所有粒子的范围
width = np.max(pos_rotated[:, :, 1], axis=1) - np.min(pos_rotated[:, :, 1], axis=1)  # shape: (#timesteps,)

# 避免除以零
width[width < 1e-10] = 1e-10
# Calculate the ratio for each timestep
ratio = length / width  # shape: (#timesteps,)

# Plot the ratio as a function of time
plt.plot(t_output, ratio)
plt.xlabel('Time (s)')
plt.ylabel('Ratio')
plt.title('Ratio of Length to Width of Cloud')
plt.show()

关键修改说明 #

1. 调整输出数组维度：将pos_output和vel_output的维度调整为(#timesteps, #particles, 3)，更符合时间步优先的逻辑，便于后续处理。
2. 完整质心系转换：同时转换位置和速度到质心系，确保基于云团自身的运动状态计算。
3. 全局旋转逻辑：针对每个时间步的质心速度方向计算旋转矩阵，将整个云团的位置转到沿质心速度的坐标系，而非单个粒子的旋转。
4. 正确计算长宽：针对每个时间步，取所有粒子在旋转后坐标轴上的最大最小值之差，得到每个时间步的云团长宽，维度与t_output匹配。

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