N = 25
dts = np.linspace(-3.1, -0.1, N)
e0s = np.linspace(0, -8, N)
integrators = ["wh", "whfast-nocor"]

niter = []
energyerror = []
energyerror_sign = []
timing = []

for integrator in integrators:
    print("Running " + integrator)
    parameters = [(1.732, dts[i], e0s[j], integrator) for j in range(N)
                  for i in range(N)]

    pool = rebound.InterruptiblePool(12)
    res = pool.map(simulation, parameters)
    res = np.nan_to_num(res)
    niter.append(res[:, 0].reshape((N, N)))
    energyerror.append(res[:, 1].reshape((N, N)))
    timing.append(res[:, 2].reshape((N, N)))
    energyerror_sign.append(res[:, 3].reshape((N, N)))

import matplotlib
matplotlib.use("pdf")
import matplotlib.pyplot as plt
from matplotlib import ticker
from matplotlib.colors import LogNorm

f, axarr = plt.subplots(3, 2, figsize=(13, 14), sharex='col', sharey='row')
extent = [dts.min(), dts.max(), e0s.max(), e0s.min()]
        rebound.calculate_megno(),
        1. / (rebound.calculate_lyapunov() * 2. * np.pi)
    ]  # returns MEGNO and Lypunov timescale in years


### Setup grid and run many simulations in parallel
N = 100  # Grid size, increase this number to see more detail
a = np.linspace(7., 10., N)  # range of saturn semi-major axis in AU
e = np.linspace(0., 0.5, N)  # range of saturn eccentricity
parameters = []
for _e in e:
    for _a in a:
        parameters.append([_a, _e])

# Run simulations in parallel
pool = rebound.InterruptiblePool(
)  # Number of threads default to the number of CPUs on the system
print("Running %d simulations on %d threads..." %
      (len(parameters), pool._processes))
res = np.nan_to_num(np.array(pool.map(simulation, parameters)))
megno = np.clip(res[:, 0].reshape((N, N)), 1.8,
                4.)  # clip arrays to plot saturated
lyaptimescale = np.clip(np.absolute(res[:, 1].reshape((N, N))), 1e1, 1e5)

### Create plot and save as pdf
import matplotlib
matplotlib.use("pdf")
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

# Setup plots
f, axarr = plt.subplots(2, figsize=(10, 10))
    planets.append({'m': .37})
    planets.append({'m': 0.00082, 'rho': 1.04534352444679,'a':0.138,'e':0.01,'inc':0 })
    planets.append({'m':0.00256,'rho':1.04534352444679,'a': 0.25,'e': 0.01, 'inc':0 })
    planets.append({'m':5.0e-05,'rho':1.04534352444679,'a':0.45 ,'e':0.01,'inc':0 })
    fname='results_test/res_t%d_k%d_d%d'%(params['tau_e'],params['tau_a'],int(params['direct']))
    res,res1,sim = mg.evolve(planets,**params)
    res.save_state(fname=fname+'_a.dat')
    res1.save_state(fname=fname+'_b.dat')

    fig,axes=plt.subplots(7,2,figsize=(20,15))
    res.plot(fig=fig,axes=axes[:,0])
    res1.plot(fig=fig,axes=axes[:,1])
    axes[0,0].set_title('$\\tau_e = %d$, $\\tau_a=%d$'%(res.tau_e,res.tau_a),fontsize=15)
    fig.savefig(fname+'.png')
    return

ta_vals = np.array([3e1,1e2,3e2,1e3,3e3,1e4,3e4,1e5,3e5])
te_vals = np.array([3e1,1e2,3e2,1e3,3e3,1e4,3e4,1e5,3e5])
direct = np.array([True,False])
args=[]
for ta in ta_vals:
    for te in te_vals:
        for d in direct:
            args.append((ta,te,d))

pool = rb.InterruptiblePool()
pool.map(pevolve,args)
print 'Finished'