Now we've done the 2D case, it shouldn't be too hard to make a 3D $n$-body simulation, and animate the result in Julia. I can't resist showing you the final product:
# Function for creating an embedded video given a filename
function html_video(filename)
base64_video = base64(open(readbytes, filename))
"""<video controls src="data:video/x-m4v;base64,$base64_video">"""
end
display("text/html", html_video("4Body3D.mp4"))
I find it mesmirizing. But we must get back to work! We start by calling the necessary packages.
using PyPlot
using PyCall
using ODE
@pyimport matplotlib.animation as anim
@pyimport mpl_toolkits
INFO: Loading help data...
The ODE's we'll be solving are of the form:
$$\frac{d^2\mathbf{r}_i}{dt^2} = \sum_j G\frac{m_j}{||\mathbf{r}_i-\mathbf{r}_j||^3}(\mathbf{r}_i-\mathbf{r}_j)$$where of course each of the $\mathbf{r}_i$ are position vectors for each planet.
We'll make these first order by considering the equivalent system:
$$\frac{d\mathbf{r}_i}{dt} = \mathbf{v}_i$$with $$\frac{d\mathbf{v}_i}{dt} = \sum_j G\frac{m_j}{||\mathbf{r}_i-\mathbf{r}_j||^3}(\mathbf{r}_i-\mathbf{r}_j)$$
We'll encode all of the position and velocity vectors in one vector $y$, in the format:
$$y = [\text{First Planet's Position}, \text{First Planet's Velocity}, \text{Second Planet's Position}, \text{Second Planet's Velocity},\ldots]$$To start, we must define the right hand side of the sytem above.
function F(t,y)
#Number of Planets (note we're solving the n-body problem, rather than just a 3-body problem)
n = int(length(y)/6)
#Extract current position and velocity
r = zeros(n,3)
v = zeros(n,3)
for i=1:n
r[i,:] = y[(i-1)*6+1:(i-1)*6+3]
v[i,:] = y[(i-1)*6+4:(i-1)*6+6]
end
#Calculate spatial derivatives
drdt = v
#Work out velocity derivatives (ie accelerations)
dvdt = zeros(n,3)
for i = 1:n
for j = 1:n
if i != j
dvdt[i,:] += -G*m[j]*(r[i,:]-r[j,:])/(norm(r[i,:]-r[j,:])^3)
end
end
end
#Combine back together into dydt vector
dydt = zeros(6*n)
for i = 1:n
dydt[6(i-1)+1] = drdt[i,1]
dydt[6(i-1)+2] = drdt[i,2]
dydt[6(i-1)+3] = drdt[i,3]
dydt[6(i-1)+4] = dvdt[i,1]
dydt[6(i-1)+5] = dvdt[i,2]
dydt[6(i-1)+6] = dvdt[i,3]
end
return dydt
end
F (generic function with 1 method)
#Set the planetary masses
m = [5,4,3,5]
n = length(m)
#Set the gravitational field strength
G = 1
#Set initial positions and velocities
r0 = zeros(n,3)
r0[1,:] = [1.0,-1.0,1.0]
r0[2,:] = [1,3,0.0]
r0[3,:] = [-1,-2,0.0]
r0[4,:] = [0,0,0.0]
v0 = zeros(n,3)
#Define y0
y0 = zeros(6*n)
for i = 1:n
y0[6(i-1)+1] = r0[i,1]
y0[6(i-1)+2] = r0[i,2]
y0[6(i-1)+3] = r0[i,3]
y0[6(i-1)+4] = v0[i,1]
y0[6(i-1)+5] = v0[i,2]
y0[6(i-1)+6] = v0[i,3]
end
#Solve the system
tf = 50
stepsPerUnitTime = 5000
tspan = linspace(0,tf,tf*stepsPerUnitTime)
t,y = ode23s(F, y0, tspan; points=:specified);
#Extract the data into a useful form
ymat= hcat(y...)
rcoords = sort([[6(i-1)+1 for i = 1:n],[6(i-1)+2 for i = 1:n],[6(i-1)+3 for i = 1:n]])
rcoords = convert(Array{Int64,1}, rcoords)
r = ymat[rcoords,:];
#Find Axis Limits
xmin = minimum(r[1,:])
for i = 2:n
xmin = min(xmin, minimum(r[3(i-1)+1,:]))
end
xmax = maximum(r[1,:])
for i = 2:n
xmax = max(xmax, maximum(r[3(i-1)+1,:]))
end
xmin, xmax = xmin - 0.1(xmax-xmin), xmax+ 0.1*(xmax-xmin)
ymin = minimum(r[2,:])
for i = 2:n
ymin = min(ymin, minimum(r[3(i-1)+2,:]))
end
ymax = maximum(r[2,:])
for i = 2:n
ymax = max(ymax, maximum(r[3(i-1)+2,:]))
end
ymin, ymax = ymin - 0.1(ymax-ymin), ymax+ 0.1*(ymax-ymin)
zmin = minimum(r[3,:])
for i = 2:n
zmin = min(zmin, minimum(r[3(i-1)+3,:]))
end
zmax = maximum(r[3,:])
for i = 2:n
zmax = max(zmax, maximum(r[3(i-1)+3,:]))
end
zmin, zmax = zmin - 0.1(zmax-zmin), zmax+ 0.1*(zmax-zmin)
#Define some random colours
colors = ["r", "g", "b", "k", "y", "c","m","0.75","0.5","0.25","#eeefff","#123456","#a125f6","#ff5868","#ef555e","#98745e","#12ff56","#789463","#8dd56d","#ddee77","#698851","#5544bb"]
#Construct Figure and Plot Data
fig = figure()
ax = plt.axes(xlim = (xmin,xmax),ylim=(ymin,ymax),zlim=(zmin,zmax), projection="3d")
for i = 1:n
line = ax[:plot]([],[],[],colors[i])[1]
x = r[3(i-1)+1,:]
x = reshape(x,length(x))
y = r[3(i-1)+2,:]
y =reshape(y,length(y))
z = r[3(i-1)+3,:]
z =reshape(z,length(z))
line[:set_data](x,y)
line[:set_3d_properties](z)
end
plt.show()
This mimics the standard Matplotlib syntax, with the same idiosyncracies we saw in the 2D case. I personally based this on the rather pretty example here:
https://jakevdp.github.io/blog/2013/02/16/animating-the-lorentz-system-in-3d/
#Use external backend
pygui(true)
#Construct Figure and Plot Data
fig = figure()
#Optional: Redefine plotting range so it is centered at the origin
xmin = -3; xmax = 3; ymin = -3; ymax = 3; zmin = -3; zmax = 3
#Create 3d axes
ax = plt.axes(xlim = (xmin,xmax),ylim=(ymin,ymax),zlim = (zmin,zmax),projection="3d")
#Make cartesian axes invisible
ax[:axis]("off")
#Add lines and planet points.
global line = Array(Any,n)
global planets = Array(Any,n)
for i = 1:n
line[i] = ax[:plot]([],[],[],color = colors[i])[1]
planets[i]= ax[:plot]([],[],[],"o", color = colors[i])[1]
end
function init()
global line
global planets
for i = 1:n
line[i][:set_data]([],[])
line[i][:set_3d_properties]([])
planets[i][:set_data]([],[])
planets[i][:set_3d_properties]([])
end
result = tuple(tuple([line[i] for i = 1:n]...)...,tuple([planets[i] for i = 1:n]...)...,None)
return result
end
#To speed up the animation, we will only plot every 'step'th data point
step = 125
function animate(i)
global line
global planets
for k = 1:n
#Get Planetary Data
x = r[3(k-1)+1,:]
x = reshape(x,length(x))
y = r[3(k-1)+2,:]
y = reshape(y,length(y))
z = r[3(k-1)+3,:]
z =reshape(z,length(z))
#Plot Data in 3D - both the planet position and the past trajectory
line[k][:set_data](x[max(1,step*(i-200)):(step*i+1)],y[max(1,step*(i-200)):(step*i+1)])
line[k][:set_3d_properties](z[max(1,step*(i-200)):(step*i+1)])
planets[k][:set_data](x[step*i+1],y[step*i+1])
planets[k][:set_3d_properties](z[step*i+1])
#Rotate axes
ax[:view_init](30, 0.3 * i)
end
result = tuple(tuple([line[i] for i = 1:n]...)...,tuple([planets[i] for i = 1:n]...)...,None)
return result
end
# call the animator.
myanim = anim.FuncAnimation(fig, animate, init_func=init, frames=ifloor(length(t)/step), interval=20)
#ffmpeg or some equialent will be required to save this animation
#myanim[:save]("4Body3D.mp4", bitrate=-1, extra_args=["-vcodec", "libx264", "-pix_fmt", "yuv420p"])
PyObject <matplotlib.animation.FuncAnimation object at 0x7f37ca3fb9d0>
# Function for creating an embedded video given a filename
function html_video(filename)
base64_video = base64(open(readbytes, filename))
"""<video controls src="data:video/x-m4v;base64,$base64_video">"""
end
display("text/html", html_video("4Body3D.mp4"))