python - python中N体问题的矢量化耦合ODE

Question

我正在尝试使用汉密尔顿运动方程解决 n 体问题；这是我所做的：

1.) 首先我为动量和向量位置定义随机初始值：

import numpy as np
import matplotlib.pyplot as plt
from scipy import integrate

n= 3 # of bodys
G = 1 # constant
m = np.random.random_sample((1,n)) # random values for masses

#----- Random Values for momentums and positions
P_0 = np.random.random_sample(size=(3,n)) # each column of the matrix represents one vector
R_0= np.random.random_sample(size=(3,n))
Y_0 = np.array([P_0,R_0])

2.) 然后我写出第 k 个粒子的汉密尔顿方程：

# Hamilton equations:

 p_kDot = lambda t,r_k,r_j,m_k,m_j: G*(m_k*m_j/np.linalg.norm(r_k-r_j)**3)*np.subtract(r_k-r_j) # Equation for the momentum

 r_kDot = lambda t,p_k,m_k: (1/m_k)*p_k # Equation for the vectors positions.

3.）然后我对所有粒子求和并定义函数 U，它将传递给 scipy.integrate.solve_ivp：

def U(t,Y,m):

  partial_Ham_r = np.array( [p_kDot(t,Y[1][k],Y[1][j],m[k],m[j]) for k in range(0,len(Y[1])) 
  for j in range(0,k)] )

  partial_Ham_p = np.array( [r_kDot(t,Y[0][i],m[i]) for i in range(0,len(Y[0]))] )                           




  return (partial_Ham_r,partial_Ham_p)

4.) 调用 scipy 积分器，但是，正如文档所说，Y_0 必须是一个向量，但我找不到将我的初始条件表示为 n 维向量的方法！：

Sol = scipy.integrate.solve_ivp(U,t_span= 
[0,10],y0=Y_0,args=m,dense_output=True,vectorized=True)

显然会导致错误说 Y_0 必须是一维的。

有没有办法表达 Y_0 以便代码工作？或者我应该在一维数组中表达位置和矩向量？

标签： pythonscipynumerical-methodsodeorbital-mechanics

Answer 1

我不知道你是否在乎，但就像你一样，我承担了编写“多体”动力学的矢量化版本的任务，并写了几个版本。如果您可能感兴趣，我将分享我写的内容。我在两个身体问题上进行了测试，甚至尝试制作一些基本的动画。

'''
vectorized many body
'''

import math
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

'''
vectorization of the gravitational acceleration field 
'''
def G_accel(positions, mass):
    q = positions.reshape((int(len(positions)/dim), dim))
    q = q.T    
    Dq = - q[:, :, np.newaxis] + q[:, np.newaxis, :]    
    D = np.linalg.norm(Dq, axis = 0) + np.identity(Dq.shape[1])
    D = 1/D
    D = D - np.identity(D.shape[1])     
    D = mass[:, np.newaxis].dot(mass[np.newaxis, :]) * D**3 # column times row times elemntwise
    Dq = D * Dq 
    Dq = Dq.sum(axis=-1) / mass[np.newaxis,:]        
    Dq = Dq.T
    return Dq.reshape((1, q.shape[0]*q.shape[1]))[0] 
    '''
    sum the matrix along the last dimension, i.e. a tensor contraction along the third index
    ''' 
'''
vectorized right-hand side of the Newtonian ODEs
'''   
def f(t, y, masses):
    n = int(len(y)/2)
    positions = y[0:n]
    velocities = y[n:]
    return np.concatenate( ( velocities,  G_accel(positions, masses) ))

'''
vectorized set up for the initial values, so that momentum is conserved and it is 0
'''
def initialize(positions, velocities, masses):
    momenta = masses[:, np.newaxis] * velocities
    velocities[-1,:] = - np.sum(momenta[0:-1, :], axis=0) / masses[-1]
    n3 = positions.shape[0]*positions.shape[1]
    pos = positions.reshape((1, n3))[0]
    vel = velocities.reshape((1, n3))[0]
    return np.concatenate((pos, vel))

'''
Test with a two-body version in the x,y plane and z = 0 component
'''
dim = 3
masses = np.array([2, 1])

q = np.array([[ -1, 0, 0],
              [1.3, 0, 0]]) 

v = np.array([[0, 0.5, 0],
              [0, 0, 0]])

n_particles = q.shape[0]


t_start = 0
t_step = 0.1 
n_steps = 5000
t_stop = n_steps * t_step


t_nodes = np.linspace(t_start, t_stop, n_steps)
y_initial = initialize(q, v, masses)

solution  = solve_ivp(fun = lambda t, y :  f(t, y, masses), 
                  t_span=[t_start, t_stop], 
                  y0=y_initial, 
                  t_eval=t_nodes, 
                  method='Radau')

pos = solution.y
pos = pos[0: int(pos.shape[0]/2)]

'''
animation plot of time-evolution:
'''
plt.style.use('seaborn-whitegrid')

fig = plt.figure()
ax = plt.axes()

ax.set_xlim(-10, 10)
ax.set_ylim(-10, 10)
line = np.empty(n_particles, dtype=object)
line[0], = ax.plot(pos[0, 0], pos[1, 0])
line[1], = ax.plot(pos[3, 0], pos[4, 0])

def animate(i):
    '''
    update plot
    '''
    line[0].set_xdata(pos[0, 0:i])
    line[0].set_ydata(pos[1, 0:i])
    line[1].set_xdata(pos[3, 0:i])
    line[1].set_ydata(pos[4, 0:i])
    return line

anim = FuncAnimation(fig, animate, frames=2000, interval=10)
ax.set_aspect( 1 )
plt.show()