#! /usr/bin/env python
"""Always comment your scripts!
This is done in three quotes, as shown here.

This script calculates the behavior of a harmonic oscillator.

Date: November 9, 2008

Author: Bob Wimmer

"""

from math import pi   #This imports all math functions and constants


#-------- Declarations ----------------------------

#Define parameters
m = 1.0   #mass
k = 1.0   #spring constant
g = 0.1   #damping constant
A = 1.0   #amplitude of driving force
w = 10.*pi #pi is imported from math

#Define initial conditions
x0 = 0.0  #initial position
xp0 = 1.0 #initial speed

#Define computational parameters
nt = 101     #number of timesteps for output
t_stop = 30. #time to stop

#Define filenames
outfilename = "oscil_v1.dat"


def func(y,t):
    x,v = hsplit(y,2)
    xp = v
    vp = A/m*sin(w*t) - g/m*v - k/m*x
    #print(vp)
    z = hstack([xp, vp])
    return z



#--------- End of Declarations --------------------

#--------- Begin calculations ---------------------
"""Use the well-known trick to reduce a second-order ODE to a
system of two first-order ODEs"""

from scipy.integrate import odeint
from numpy import *

x = x0                        #initial position
v = xp0                       #initial speed
y = hstack([x,v])             #array contains initial conditions
t = linspace(0., t_stop, nt)    #array contains timesat which to evaluate
z = odeint(func,y,t)
x_arr, v_arr = hsplit(z,2)

#open output file
with open(outfilename, 'w') as f:
    i = 0
    while i < nt:
        f.write('%g \t %g \n' % (t[i], x_arr[i]))
        i = i + 1