for filename in glob.glob('tmp_errors*.pdf'):
    os.remove(filename)

def exact_S_solution(t):
    return g*(1 - np.cos(t))

## Default values
m = 1; b = 2; L = 10; k = 1; beta = 0; S0 = 0; 
dt = 2*np.pi/40; g = 9.81; w_formula='0'; N = 200

m, b, L, k, beta, S0, dt, g, w, N = init_prms(m, b, L, k, beta, S0, dt, g, w_formula, N)

def w(t):
    return 0

S = solve(m, k, beta, S0, dt, g, w, N)
S = np.array(S)

tcoor = np.linspace(0, N*dt, len(S))
exact = exact_S_solution(tcoor)
error = exact - S
fig1, (ax1, ax2) = plt.subplots(1,2)
ax1.plot(tcoor, S, 'r-', label='numeric solution')
ax1.plot(tcoor, exact, 'b-', label='exact solution')
ax1.set_xlabel('time'); ax1.set_ylabel('S'); ax1.legend(loc=1)
ax2.plot(tcoor, error, 'b-', label='error')
ax2.set_xlabel('time'); ax2.set_ylabel('error'); ax2.legend(loc=1)
fig1.show()
fig1.savefig('tmp_errors_1.pdf')

fig2 = plt.figure()
# -*- coding: utf-8 -*-
"""
Created on Thu Oct 08 23:14:33 2015

@author: Georg
"""

from ode_project import init_prms, solve
from math import cos
import matplotlib.pyplot as plt

def exact_S_solution(t):
    return g*(1-cos(t))
    
def check_S(S, t, step):
    error = exact_S_solution(t) - S[step]
    print 't = %5.2f   S[%04d] = %+g   error = %g' % (t, step, S[step], error)

## Fixed values for a test
from math import pi
m = 1; b = 2; L = 10; k = 1; beta = 0; S0 = 0; dt = 2*pi/40; g = 9.81; N = 200

def w(t):
    return 0

S = solve(m, k, beta, S0, dt, g, w, N, user_action=check_S)

plt.plot(S)
plt.xlabel('time')
plt.ylabel('Y')
m, b, L, k, beta, S0, dt, g, w, N = init_prms(m, b, L, k, beta, S0, dt, g, w_formula, N)

## Vectorize StringFunction w
w_formula = str(w)
if ' else ' in w_formula:
    w = np.vectorize(w)   ## General vectorization
else:
    w.vectorize(globals())

plt.ion()
lines = plt.plot(0, w(0)-L-S0-b/2.0)
plt.axis([0, dt*N, -35, -8])
plt.xlabel('time')
plt.ylabel('Y')

S = solve(m, k, beta, S0, dt, g, w, N, user_action=plot_S)

plt.savefig('tmp_Y.pdf')

tcoor = np.linspace(0, N*dt, N+1)
S     = np.array(S)

nrows = 2
if beta != 0:
    nrows += 1
if w_formula != '0':
    nrows += 1
    
fig, (axarr) = plt.subplots(nrows, sharex=True, sharey=False)

## Position Y(t)