The two-body problem

import numpy as np
from scipy.integrate import odeint
import math
import matplotlib.pyplot as plt

def xvyplot(t,sol,e):
    plt.plot(sol[:, 0],sol[:, 1] , 'p', label='x(t) vs y(t)')
    plt.title("x v y with eccentricity of " + str(e))
    plt.legend(loc='best')
    plt.xlabel('x')
    plt.ylabel('y')
    plt.grid()
    plt.show()

def xvtplot(t,sol,e):
    plt.plot(sol[:, 0],t , 'p', label='x(t) vs t')
    plt.title("x v t with eccentricity of " + str(e))
    plt.legend(loc='best')
    plt.xlabel('x')
    plt.ylabel('t')
    plt.grid()
    plt.show()

def yvtplot(t,sol,e):
    plt.plot(sol[:, 1],t , 'p', label='y(t) vs t')
    plt.title("y v t with eccentricity of " + str(e))
    plt.legend(loc='best')
    plt.xlabel('y')
    plt.ylabel('t')
    plt.grid()
    plt.show()

e1 = 0.9 #eccentricity
#also try 0 and 0.5
e2 = 0  # circular orbit
e3 = 0.5
NUM_POINTS = 90
t = np.linspace(0, np.pi*6, NUM_POINTS)

def twobody(v,t):
    x, y, xPrime, yPrime = v
    r = math.sqrt(math.pow(x,2) + math.pow(y,2))

    dydt = [
        xPrime, # xprime formula
        yPrime, # yprime formula
        -x/(r*r*r), # xdprime
        -y/(r*r*r) # y dprime
    ]
    return dydt

def conservationOfEnergy(x, y, xPrime, yPrime):
    r = np.sqrt(x*x + y*y)
    return (1/2)*(xPrime*xPrime + yPrime*yPrime) - 1 / r

def conservationOfAngularMomentum(x, y, xPrime, yPrime):
    return x*yPrime-y*xPrime



e = e1
sol = odeint(
    twobody, 
    t=t,
    y0=np.array([
        1 - e,
        0,
        0,
        math.sqrt((1 + e) / (1 - e))
    ])
)

xvyplot(t,sol,e)
xvtplot(t,sol,e)
yvtplot(t,sol,e)

x = sol[:,0]
y = sol[:,1]
energy = conservationOfEnergy(x, y, sol[:,2], sol[:,3])
momentum = conservationOfAngularMomentum(x, y, sol[:,2], sol[:,3])
'''
1. Convert the 2nd order equations into first order equations
2. Plug in initial values
'''

e = e2
sol = odeint(
    twobody, 
    t=t,
    y0=np.array([
        1 - e,
        0,
        0,
        math.sqrt((1 + e) / (1 - e))
    ])
)

xvyplot(t,sol,e)
xvtplot(t,sol,e)
yvtplot(t,sol,e)

e = e3
sol = odeint(
    twobody, 
    t=t,
    y0=np.array([
        1 - e,
        0,
        0,
        math.sqrt((1 + e) / (1 - e))
    ])
)

xvyplot(t,sol,e)
xvtplot(t,sol,e)
yvtplot(t,sol,e)

e = e1
plt.plot(t , energy, 'p',)
plt.title("t v energy with eccentricity of " + str(e))
plt.legend(loc='best')
plt.xlabel('t')
plt.ylabel('energy')
plt.grid()
plt.show()

plt.plot(t , momentum, 'p',)
plt.title("t v momentum with eccentricity of " + str(e))
plt.legend(loc='best')
plt.xlabel('t')
plt.ylabel('momentum')
plt.grid()
plt.show()