Three-body Problem simulation code

# Three body problem graphical simulation

import numpy as np
import matplotlib.pyplot as plt
from  scipy.integrate import odeint
from matplotlib.animation import FuncAnimation
from matplotlib.widgets import Button
from matplotlib.animation import writers

acc = lambda r, m, g : -r*m*g/np.linalg.norm(r)**3

m1 = 1
m2 = 1
G = 1
init = np.array([0.0, -0.00, 0.2, -0.2,
                  -1.0, 0.0, 0.0, 0.5,
                  1.0, 0.0, 0.0,-0.5,
                 0.0, 0.00, 0.2005, -0.2,])
init2 = np.array([0.0, -0.00, 0.5, -0.2,
                  -1.0, 0.0, 0.0, 0.2,
                  1.0, 0.0, 0.0,-0.5,
                 0.0001, 0.00, 0.5, -0.2,])


def f(x, t):
    # This describes gravitational pull of two planets at (-2,0) and (2,0)
    # Position of planets
    x1 = x[4:6]
    x2 = x[8:10]

    r1 = x[:2]-x1 # Get radial vectors
    r2 = x[:2]-x2
    r12 = x2-x1
    a = acc(r1, m1, G)+acc(r2, m2, G) # Calculate net acceleration
    a3 = acc(x[12:14]-x1, m1, G)+acc(x[12:14]-x2, m2, G)

    a1 = acc(-r12, m2, G)
    a2 = acc( r12, m1, G)
    return np.array([x[2], x[3], a[0], a[1],
                     x[6], x[7],a1[0],a1[1],
                    x[10],x[11],a2[0],a2[1],
                    x[14],x[15],a3[0],a3[1]])



# First two numbers position, second two velocity, for test mass, m1, m2
# Calculate different solutions
res1 = odeint(f, init, np.arange(0,100,0.05))

m2 = 0.4
res2 = odeint(f, init2, np.arange(0,100,0.05))
res = res1

log1 = np.log(np.linalg.norm(res1[:,0:2]-res1[:,12:14], axis=1))
log2 = np.log(np.linalg.norm(res2[:,0:2]-res2[:,12:14], axis=1))
log = log1

c1 = plt.Circle(res[0, 4:6], 0.08, fc='orange', ec='black')
c2 = plt.Circle(res[0, 8:10], 0.08, fc='g', ec='black')
h1 = plt.Circle(res[0,:2], 0.04, fc='r', ec='black')
h2 = plt.Circle(res[0,12:14], 0.04, fc='b', ec='black')

def setup():
    ax.set_xlim(-3,3)
    ax.set_ylim(-3,3)
    ax.plot(res[0,0], res[0,1], 'r--', linewidth=2)
    ax.plot(res[0,12], res[0,13], 'b--', linewidth=2)
    ax.add_artist(c1)
    ax.add_artist(c2)
    ax.add_artist(h1)
    ax.add_artist(h2)

    ax2.plot(0, log[0], '-', color='black', linewidth=0.5)

def update(i):
    ax.lines[0].set_data(res[max(i-100,0):i+1,0], res[max(i-100,0):i+1,1])
    ax.lines[1].set_data(res[max(i-100,0):i+1,12], res[max(i-100,0):i+1,13])

    h1.center = res[i,:2]
    h2.center = res[i,12:14]
    c1.center = res[i, 4:6]
    c2.center = res[i, 8:10]
    ax2.lines[0].set_data(np.arange(0,i,1)*0.05,log[:i:1])

state = 0
def callback(event):
    global res, state, log
    state += 1
    if(state%2):
        res = res2
        log = log2
        t.set_text('Simulation:\nUneven Masses')
    else:
        res = res1
        log = log1
        t.set_text('Simulation:\nEven Masses')


fig = plt.figure(figsize=(10,5))
ax = plt.axes([0.1, 0.2, 0.35, 0.7], aspect='equal')
ax.set_title('Restricted Three-Body Problem Simulation')
ax2 = plt.axes([0.55, 0.1, 0.35, 0.8])
ax2.set_title('Logarithm of distance between balls')
ax2.set_xlim(0,100)
ax2.set_ylabel(r'$\log(distance) \rightarrow$')
ax2.set_ylim(-10,2)
ax2.set_xlabel(r'time $\rightarrow$')
# Add Button
axbutton = plt.axes([0.1, 0.05, 0.15, 0.08])
button = Button(axbutton, 'Change simulation', color='red', hovercolor='green')
axsimname = plt.axes([0.30, 0.05, 0.15, 0.1])
t = axsimname.text(0.5,0.5,'Simulation:\nEven Masses', horizontalalignment='center',
               verticalalignment='center', transform=axsimname.transAxes)
axsimname.axis('off')
button.on_clicked(callback)

ani = FuncAnimation(fig, update, frames=range(len(res)),
                    init_func=setup, blit=False, interval=30, repeat=False)
#ani.save('Even.mp4')