#This scripts illustrates how to use
#the integrator in ClimateUtilities to
#numerically solve a second ordinary differential equation
# of the form 
#
#                d^2Y/dt^2 = G(t,Y)
#
#for the function Y(t). Y is a scalar.  One example of this
#would be the harmonic oscillator, defined by G = -kY
#
#The method used is to rewrite this as a pair of first order
#equations:
#                dZ1/dy = Z2
#                dZ2/dt = G(t,Z2) 

import math
from ClimateUtilities import *

#Define the function G.  The following yields the
#linear harmonic oscillator. For generality, the independent variable
#is included in the argument list even though it isn't
#used in this example. 
def G(t,Y,k):
    return -k*Y

#Define the function that computes the derivatives.
#the third argument allows us to pass parameters to
#the function.  In this case, the dependent variable Z is
#a numpy array, and the function returns an array of slopes.
#
#The parameter that G expects is a real number, for the
#harmonic oscillator example, but if using a function G
#that requires more information to evaluate, param can be
#a list of values, or indeed any Python object at all, so long
#as G is set up to handle it.
def f(t,Z,param):
    Z0prime = Z[1]
    Z1prime = G(t,Z[0],param)
    #It's not really necessary to define the above intermediate
    #variables, since the right hand sides could be put directly
    #in the list passed to numpy.array below, but defining the
    #intermediate variables makes it a bit easier to keep track
    #of what is being done
    return numpy.array([Z0prime,Z1prime])
#------------Done with function definitions----------
#
#Set the spring constant
k = 4.
#Set some constants defining the integration
dt = .1      #The increment to use. 
             #Smaller values give more accurate results
Zstart = numpy.array([0.,1.])#The initial value of Z. Must be a numpy array
tstart = 0.  #The initial value of x
tfinish = 2.*math.pi#Where you want to stop the integration
#
#Tip:  In this example you start from a smaller value of
#      t and integrate toward a large value.  It's also
#      OK to do an integration backwards starting from
#      a larger value and finishing at the smaller value.
#      To do this, you deed to redefine tstart and tfinish
#      and you also need to make dt negative.  You will also
#      have to change the condition in the "while" loop below
#      to use a > instead of a < . 
#
#-----------------------------------------------------
#Now create an instance of the integrator class. This
#class implements Runge-Kutta numerical integration. It
#is defined in ClimateUtilities. We call the instance "m"
#but of course you can use any name you want, and also create
#as many different instances as you want
#
m = integrator(f,tstart,Zstart,dt)
#
#Set the parameter value (optional if you aren't using a parameter)
m.setParams(k) #Sets the parameter  to the value of k
#
#Make lists to save the results for plotting
tList = [tstart]
YList = [Zstart[1]]
#
#Step through until x exceeds the desired target
t = tstart
while t <= tfinish:
    #Now we call the "next" method of the integrator object
    #m.next() does all the work of carrying out the integration
    #It takes the current value of t and Z (conveniently stored in
    #the object, so you don't have to tell the next method about what
    #these values are!), computes slope information, and figures out
    #what the new value of Z is at t+dt.  It returns a pair of values,
    #consisting of the new value of t and the new value of Z . Z is
    #an array, but the variable we want here is just Z[1]. Z[0] is
    #dY/dt, i.e. the velocity, and we could save that too if we wanted
    #
    t,Z = m.next() # next() can take a new value of dx as optional argument
                   # e.g. m.next(.01) if you want to change to a smaller
                   # increment.  The new value you use will continue to be
                   #used in subsequent calls to next() even if you call
                   #next without an argument
    #
    #Put the values in the results list
    tList.append(t)
    YList.append(Z[1])
#
#----------Integration loop complete----------------------
#
#Now we make some plots to see what the result looks like
#
#Put results in a Curve object and plot
c = Curve()
c.addCurve(tList,'t')
c.addCurve(YList,'Y')
w = plot(c)
#