from __future__ import division
from scipy import *
from scipy import integrate
from scipy import linalg

G = 6.67e-11 # m**3 * (kg)**-1 * s**-2
mass_sun = 1.99e30 # kg

initial_state = array([
    1.5e11, # m
    0,      # m
    0,      # m / s
    2.98e4  # m / s
])

t = linspace(0, 60*60*24*250, 500) # s

def f(state, time):
    """
    state is a vector
    state = [
        position in x
        position in y
        velocity in x
        velocity in y
    ]
    """
    return array([
        state[2],
        state[3],
        -G * mass_sun * state[0] / linalg.norm(state[:2])**3,
        -G * mass_sun * state[1] / linalg.norm(state[:2])**3
    ])

result = integrate.odeint(f, initial_state, t)

import pylab
pylab.plot(result[:,0], result[:,1])
pylab.plot([0], [0], 'r*', markersize=25)
pylab.axis('equal') # force Pylab to use equal units on the x and y axes
pylab.show()


