"""

  Solving the second order differential orbit equations
    u'' + u - 1/r0 = 0
  and
    u'' + u - \gamma/r0 = 0
  using odsolve which solves two coupled first order equations for
  the variables u and u' which are named u[0] and u[1] respectively.

"""

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

# Constants
# Values chosen is radius and eccentricity for Mercury's orbit
r0   = 57.9    # Radius of circular orbit [10^6 km]
phi0 =  0.0    # Starting angle
ecc  =  0.2056 # Eccentricity
eps  =  0.0266 # Small relativistic parameter (in reality eps=2*10^-8)

# Non-relativistic equation
def f(u,x):
    return (u[1],-u[0]+1/r0)

# Relativistic equation (first order in v^2/c^2 approximation)
def f_rel(u,x):
    return (u[1],-u[0]+1/r0+0.5*eps*r0*u[0]**2)

# Initial condition in terms of u(0) and u'(0)
# Set by using u[0]=1/r0+ecc/r0*\cos\phi_0 and u'[0]=ecc/r0*\sin\phi_0
# from exact solution of non-relativistic equation
u0 = [1/r0+ecc/r0*np.cos(phi0),ecc/r0*np.sin(phi0)]

# Range of angles to plot
norbits = 10
phi = np.linspace(0,norbits*2*np.pi,1000)

# Numerical solutions
us = odeint(f, u0, phi)
us_rel = odeint(f_rel, u0, phi)  # For comparison use same intial conditions

# Pick out solutions for u (not u')
u = us[:,0]
u_rel = us_rel[:,0]

# Plot in a polar plot
plt.figure()
plt.polar(phi, 1./u , linewidth=2)
plt.polar(phi, 1./u_rel , linewidth=2)
plt.xlabel(r'$\phi$ ')
plt.ylabel(r'$\rho$ [$10^6$ km]')
plt.savefig('orbits.png')