#Python script coriolis_on_beta.py
#Integrate equations of motion for particle subject to Coriolis acceleration on beta plane.
#Inspired by the const_ang_mom_traj.m scripts in Holton's materials.
#To use this script,
#*Install the EPD as per the python learning materials.
#*Start a Pylab (ipython) session  
#*Type "import coriolis_on_beta" (make sure you are in the directory where the script is located).
#*To run the script again after you have made changes, type "reload coriolis_on_beta".

#The following "import"s all the necessary packages needed to run this script.
#I have, for this script, chosen to use the "from" method.
from matplotlib.pyplot import *
from numpy import *
#all quantities in MKS units
#reference value of coriolis parameter
f0 = 1e-4
#beta parameter
beta = 0
#initial x position
x=0
#initial y position
y=0
#initial x velocity
u=1.0
#initial y velocity
v=0.0
#state vector. The array command is part of the numpy module.
state = array([x, y, u, v])
#number of days to run program
days = 10.0
#runtime in seconds (86400 = 24*3600)
runtime = days*86400
#time step in seconds
dt = 300

#functions to be used below
#calculate derivative with respect to time of x, y, u, v
def tendency(f0, beta, y, u, v):
    #x velocity
    dxdt = u
    #y velocity
    dydt = v
    #x acceleration
    dudt =  (f0 + beta*y ) * v
    #y acceleration
    dvdt = -(f0 + beta*y ) * u
    return array([dxdt, dydt, dudt, dvdt])

#calculate update to state vector s using Adams-Bashforth formula:
#s(n) = s(n-1) + dt*(a*dsdt(n)+b*dsdt(n-1)+c*dsdt(n-2))
#a=23/12, b=-16/12, c = 5/12, note that a + b + c = 1
def adams_bashforth(dstate_dt_n,dstate_dt_1,dstate_dt_2,dt):
    output = dt*(23.0*dstate_dt_n-16.0*dstate_dt_1+5.0*dstate_dt_2)/12.0
    return output

#calculate tendency and initialize tendency arrays for Adams-Bashforth calculation
dstate_dt_n = tendency(f0,beta,y,u,v)
dstate_dt_1 = dstate_dt_n
dstate_dt_2 = dstate_dt_1

#create an array of time points
time = arange(0,runtime,dt)
#create a solution array for plotting
x_solution = zeros(len(time))
y_solution = zeros(len(time))

#integrate equations of motion
#range(len(time)) creates an array [0,...,len(time)-1]
for j in range(len(time)):
    #update state
    state = state + adams_bashforth(dstate_dt_n,dstate_dt_1,dstate_dt_2,dt)
    #save tendencies to previous timesteps
    dstate_dt_2 = dstate_dt_1
    dstate_dt_1 = dstate_dt_n
    #new state
    [x,y,u,v] = state
    #get new tendency
    dstate_dt_n = tendency(f0,beta,y,u,v)
    #save current state for plotting
    x_solution[j] = x/1000
    y_solution[j] = y/1000
    #print diagnostics every 12 hours
    if time[j]%(3600*12)==0:
        print 'time(d): %5.2f, x(km): %8.2f, y(km): %8.2f, u: %8.2f, v: %8.2f'%(time[j]/86400,x/1000,y/1000,u,v)


figure(1)
clf()
xlabel('x (km)')
ylabel('y (km)')
title('Coriolis force with beta')
#Here's a fancier title string. Uncomment it if you want to use it.
#The "r" before the string indicates that latex-style formatting will be used. The %5.1f are format strings that format the values (f0,beta) at the end of the string.
#title_string_fancy = r'Coriolis force with $f_0$ = %5.1f and $\beta$ = %5.1f'%(f0,beta)
#title(title_string_fancy)

#plot solution
plot(x_solution,y_solution)
#plot starting point with a big green circle. For an array A, A[0] is first entry.
plot([x_solution[0]],[y_solution[0]],'go',markersize=10)
#plot ending point with a big red circle. For an array A, A[-1] is last entry.
plot([x_solution[-1]],[y_solution[-1]],'ro',markersize=10)
#set aspect ratio of plot to 1
axis('scaled')
#uncomment to save figure in pdf format
#savefig('coriolis_on_beta.pdf')