#!/usr/bin/env python
# coding: utf-8

# # Cosmological Constant Correction
# The cosmological constant is the simplest model for dark energy, which provides a mechanism for the accelerating expansion of the universe. In the vicinity of a spherically symmetric point mass, $M$, the gravitational potential is approximately
# 
# $$\Phi = -\frac{GM}{r} - \frac{c^2 \Lambda r^2}{6},$$
# 
# where $G$ is Newton's gravitational constant and $\Lambda$ is the cosmological constant. Consequently, objects straying too far from the central mass $M$ can be blown away by cosmic inflation if $\Lambda > 0$.
# 
# ### Problem
# 1. Find the effective potential, $V_{eff}$, for a particle orbiting the mass $M$, explaining why angular momentum is conserved.
# 2. Sketch the effective potential for a range of angular momenta and cosmological constants, identifying where gravity becomes repulsive.
# 3. Determine a condition on this crossover radius. Test this by using `CosmologicalConstant.ipynb` to visualise the paths around $M$ for various starting positions, velocities and values of $\Lambda$.
# 4. By considering the effective potential, explain the relationship between an orbit's energy and its precession.
# 5. Using the above crossover radius as an order of magnitude estimate for the attractive-repulsive crossover radius, discuss the implications for the solar system if the cosmological constant has the value $6 \times 10^{-34}$m$^{-2}$.

# In[12]:


#NAME: Cosmological Constant
#DESCRIPTION: Investigating the cosmological constant correction term for the gravitational potential around a point mass.

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from matplotlib import animation
#HTML is used to run the animation inline
from IPython.display import HTML

#gravitational constant
G = 1.0
#mass of star
M = 100.0
#cosmological constant * c^2
C = 0.1

#starting co-ordinates (x, y, v_x, v_y)
a0 = [5.0, 0.0, -3.0, 3.0]

#time steps
dt = 0.05
Nsteps = 500
t = np.linspace(0, Nsteps*dt, Nsteps)


# The equation of motion for a particle in gravitational field $\vec{g}$ is
# $$\frac{d^2 \vec{r}}{dt^2} = \vec{g}.$$
# This can be separated into two sets of first-order differential equations which, setting 
# $$\vec{a} = \left( \begin{matrix}
#                 x\\
#                 y\\
#                 \frac{dx}{dt}\\
#                 \frac{dy}{dt}\\
#                 \end{matrix} \right)$$
# can be written in 2D as
# $$\frac{d \vec{a}}{dt} = \left( \begin{matrix}
#                              a_2\\
#                              a_3\\
#                              g_0\\
#                              g_1\\
#                              \end{matrix} \right)$$

# In[13]:


#calculate gravitational field
def g(x,y):
    #x and y are scalar floats
    r = np.sqrt(x**2 + y**2)
    x_component = (-G*M/(r**3) + C/3)*x
    y_component = (-G*M/(r**3) + C/3)*y
    return x_component, y_component

#derivates of coordinates
def derivative(a, t):
    gx, gy = g(a[0], a[1])
    return [a[2], a[3], gx, gy]

#solve the differential equation above
trajectory = odeint(derivative, a0, t)


# In[14]:


#create the animation
fig = plt.figure(figsize = (8,6))
plt.xlim(-10,10)
plt.ylim(-10,10)
plt.title("Cosmological Constant Correction \n to Newton's Law of Gravitation")
plt.xlabel('x')
plt.ylabel('y')
#central mass
mass, = plt.plot([0.0],[0.0], 'ko')
#orbiting body
point, = plt.plot([trajectory[0][0]],[trajectory[0][1]], 'bo')
path, = plt.plot([trajectory[0][0]],[trajectory[0][1]], 'b-')

def animator(i):
    point.set_data([trajectory[i][0]], [trajectory[i][1]])
    path.set_data(trajectory[:i,0], trajectory[:i,1])
    return point, path

anim = animation.FuncAnimation(fig, animator, Nsteps, interval = 20, repeat = False, blit = True)
html_video = anim.to_html5_video()
HTML(html_video)