Vectorized coupled ODE’s for a N body Problem in python

Question

I’m trying to solve a n-Body problem using hamilton equations of motion; here is what I’v done:

1.) First I define random initial values for momentums and the vectors positions:

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.) then I write the hamilton equations for a k-th particle:

# 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.)Then I sum over all the particles and define the function U, which is gonna be pass to 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.) Call the scipy integrator, but, as documentation say, Y_0 must be a vector, but i dont find a way to express my initial conditions as a n-dimensional vector!:

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

obviously outcomes the an error saying that Y_0 must be 1-Dimensional.

Is there a way to express Y_0 so that the code works? or should I express the position and moment vectors in 1 dimensional arrays?

Answer 1

I do not know if you care at all, but just like you I took upon the task to write vectorized version of a "many-body" dynamics and wrote a couple of versions. In case you might be interested, I will share what I wrote. I tested it on the two body problem and tried to even make some rudimentary animation.

'''
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()

Answer 2

I think I solved it, the code works when the initial conditions are such that there are no collisions between the particles (in which case the algorithm diverges):

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

n= 3 # of bodys
G = 1 # constant
#m = np.random.randint(low=1,high=100,size=(n,)) # random values for masses
m = [1,1,1]

#----- 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))
R_0= np.array([[0,0,0],[0,1,0],[0,1,0.96]]
Y_0 = np.concatenate((P_0,R_0),axis=1)


#----------------------------------------------
# Hamilton equations:

r_kDot = lambda t,p_k,m_k: (1/m_k)*p_k

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)*(r_k- 
r_j)


def U(t,Y):

Y = Y.reshape(3,2*n)
Y = [Y[:,:n],Y[:,n:]]


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

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)] )
                      
return np.concatenate((partial_Ham_r,partial_Ham_p),axis=1)

I just reshaped the Y arg into the U function for have the sime shape as the Y_0 concatenated matrix, a minus sign was missing as someone commented

#%matplotlib notebook
Sol = integrate.solve_ivp(U,y0=Y_0.flatten(),t_span= 
[0,10],dense_output=True,vectorized=True,t_eval=np.linspace(0,10,9999))
print (np.shape(Sol.y))
print (Sol.y)
r_vectors = Sol.y[9:,:]



fig , axes = plt.subplots(subplot_kw={ "projection":"3d" })


axes.plot(r_vectors[6,:],r_vectors [7,:],r_vectors [8,:]) 
axes.plot(r_vectors[0,:],r_vectors [1,:],r_vectors [2,:])
axes.plot(r_vectors[3,:],r_vectors [4,:],r_vectors [5,:])

I also generalized it to n bodies:

m = np.random.randint(low=1,high=100,size=(n,)) # random values for masses
m = np.ones((n,))
m[0]=2
#----- 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.concatenate((P_0,R_0),axis=1)

Sol = integrate.solve_ivp(U,y0=Y_0.flatten(),t_span= 
[0,10],dense_output=True,vectorized=True,t_eval=np.linspace(0,10,9999))


TrayX = Sol.y[3*n::3,:]
TrayY = Sol.y[3*n+1::3,:]
TrayZ = Sol.y[3*n+2::3,:]

print(np.shape(TrayX))
print(np.shape(TrayY)) 
print(np.shape(TrayZ)) 

fig , axes = plt.subplots(subplot_kw={ "projection":"3d" })

for j in range(n):
   axes.plot(TrayX[j],TrayY[j],TrayZ[j])

perhaps some of you may know how to prevent the code from failing when there are particle collisions ...