simulation_movie.py

import rebound 
import numpy as np
import matplotlib.pyplot as plt
import rebound
import argparse


parser=argparse.ArgumentParser(description='Make a movie from rebound SimlationArchive')
parser.add_argument('loc',help='Data location')
parser.add_argument('--xlim', type=float, default=1)
parser.add_argument('--ylim', type=float, default=1)
args=parser.parse_args()

##Location of simulation data...
loc=args.loc
xlim=args.xlim
ylim=args.ylim

##You will have to modify the names of the data files (the archive_out part) here. 
sa=rebound.SimulationArchive('{0}/archive.bin'.format(loc))
frames=len(sa)


pts=200
##Arrays to store data from stellar orbits
x=np.empty([frames, len(sa[0].particles)-1, pts+1])
y=np.empty([frames, len(sa[0].particles)-1, pts+1])
z=np.empty([frames, len(sa[0].particles)-1, pts+1])

fig,ax=plt.subplots(figsize=(10,10))
col='0.5'

for idx in range(frames):
	sim=sa[idx]
	orbs=sim.calculate_orbits(primary=sim.particles[0])
	for i in range(len(orbs)):
		p=sim.particles[i+1]
		##Plot orbits of all bound particles 
		if orbs[i].a>0:
			pos=np.array(p.sample_orbit(pts, primary=sim.particles[0]))
			x[idx,i,:-1] = pos[:,0]
			x[idx,i][-1]=np.nan
			y[idx,i,:-1] = pos[:,1]
			y[idx,i][-1]=np.nan
			z[idx,i,:-1] = pos[:,2]
			z[idx,i][-1]=np.nan

			ax.plot(x[idx,i,:-1], y[idx,i,:-1], '-', color=col)
	##Label each plot with the time
	ax.annotate('t={0:.2f}'.format(sim.t), (0.01,0.99), va='top', xycoords='axes fraction')

	ax.set_xlabel('x [pc]')
	ax.set_ylabel('y [pc]')
	ax.set_xlim(-xlim, xlim)
	ax.set_ylim(-ylim, ylim)
	##Save figure data
	fig.savefig(loc+'/tmp_movie_{0:03d}.png'.format(idx), ppi=72)
	plt.cla()