gravity.py

# This is code written by Nathaniel Virgo, for the sake of replying to a Twitter thread.
#
# The code is in the public domain and you can do what you want with it.
#
# It simulates N masses moving in 2D with Newtonian gravity, with symmetric initial
# conditions, and outputs an animated gif.
#
# This code makes no attempt to be accurate, because the whole point was to show
# the system breaking down as a result of numerical inaccuracies. In particular, it
# doesn't use a conservative ODE solver, so even before it breaks down a little bit of
# energy gets dissipated, which it shouldn't do. Just don't rely on this code for any
# actual research, is what I'm saying.
#
# It also makes no particular effort to be fast, and it relies on Python loops.
#
# You need Python 3, scipy and numpy, and I guess you need ImageMagick to be installed
# if you want to save the animation. (If you don't, see the comment at the bottom of
# the file.) Please don't ask me for help installing prerequisites, because I did that
# a long time ago and have no knowledge of the process.
#
# just run
#
# python3 gravity.py
#
# It will create a very large gif file in the same folder as the code. You will need
# some kind of gif editing software to reduce the file size to something you can post online.
# I use "GIF Brewery 3" on the Mac, which is free.

from scipy.integrate import solve_ivp
import numpy as np


N = 7 # no. of bodies
inv_m = 1 # 1/mass (all the same for now)
init_v = 1/2 # initial radial velocity of the masses

T = 20 # amount of simulated time
steps = 200 # integration steps per unit of integrated time.
# (So there are T*steps integration steps in total)
skip = 10 # when creating the animation, only show one frame every 'skip' steps

output_file = "./animation.gif"

x_s = np.s_[:N]
y_s = np.s_[N:2*N]
vx_s = np.s_[2*N:3*N]
vy_s = np.s_[3*N:4*N]

z = np.zeros(N*4)
# this is the x's, y's, vx's, vy's, all in one vector so we can pass it to
# the ode solver

def newton(t, z):
	x = z[x_s]
	y = z[y_s]
	vx = z[vx_s]
	vy = z[vy_s]
	dz = np.zeros_like(z)
	dz[x_s] = vx
	dz[y_s] = vy
	for i in range(N):
		for j in range(i):
			Dx = x[j]-x[i]
			Dy = y[j]-y[i]
			r2 = Dx*Dx + Dy*Dy
			r  = np.sqrt(r2)
			Fx = Dx/r/r2
			Fy = Dy/r/r2
			dz[vx_s][i] += Fx*inv_m
			dz[vy_s][i] += Fy*inv_m
			dz[vx_s][j] -= Fx*inv_m
			dz[vy_s][j] -= Fy*inv_m
	return dz

for i in range(N):
	theta = 2*np.pi*i/N
	z[x_s][i] = np.sin(theta)
	z[y_s][i] = np.cos(theta)
	z[vx_s][i] = np.sin(theta+np.pi/2)*init_v
	z[vy_s][i] = np.cos(theta+np.pi/2)*init_v

t = np.linspace(0,T,T*steps)

sol = solve_ivp(newton, [0,T], z, t_eval=t, method='RK23')

x = sol.y[x_s,:]
y = sol.y[y_s,:]

from matplotlib import pyplot as plt
import matplotlib.animation as animation

fig, ax = plt.subplots()

lines = [ax.plot(x[i,:],y[i,:])[0] for i in range(N)]
points = [ax.plot(x[i,-1],y[i,-1],'.',color = lines[i].get_color())[0] for i in range(N)]

plt.xlim([-1.1,1.1])
plt.ylim([-1.1,1.1])
ax.set_aspect('equal')

plt.axis('off')

def animate(t):
	for i in range(N):
		lines[i].set_data(x[i,:t],y[i,:t])
		points[i].set_data(x[i,t],y[i,t])
	return lines, points

lines[i].set_data([],[])
points[i].set_data([],[])

anim = animation.FuncAnimation(fig, animate, frames=range(0,len(sol.t),skip),interval=1)

# if you just want to watch the animation, comment out the next line and uncomment
# the one after it

anim.save(output_file, writer='imagemagick', fps=60)
# plt.show()