from pylab import *
from numpy import *

h = 0.1
x = array([-0.24, -.10, 0.00, 0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.75, 0.97])
t = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]

def F(t):
    return 3*cos(3*t) - 2*sin(3*t)

def g(x, t, h, i):
    x2 = (x[i-1]-2.0*x[i]+x[i+1])/(h**2)
    x1 = (x[i+1] - x[i-1])/(2*h)
    x0 = x[i]
    return (2.0*x2) - (3.0*(x1**2)/4.0) + 3.0*x0 - F(t[i])

def par1(h=h):
    return -4.0/h**2 + 3.0

def par2(p1, p2, h=h):
    return 2.0/h**2 - (6.0/(16.0*h**2))*(p2-p1)

def Jac(x, h=h):
    Jacob = array([[par1(),par2(x[0],x[2]),0,0,0,0,0,0,0],
    [par2(x[1],x[3]),par1(),par2(x[1],x[3]),0,0,0,0,0,0],
    [0,par2(x[2],x[4]),par1(),par2(x[2],x[4]),0,0,0,0,0],
    [0,0,par2(x[3],x[5]),par1(),par2(x[3],x[5]),0,0,0,0],
    [0,0,0,par2(x[4],x[6]),par1(),par2(x[4],x[6]),0,0,0],
    [0,0,0,0,par2(x[5],x[7]),par1(),par2(x[5],x[7]),0,0],
    [0,0,0,0,0,par2(x[6],x[8]),par1(),par2(x[6],x[8]),0],
    [0,0,0,0,0,0,par2(x[7],x[9]),par1(),par2(x[7],x[9])],
    [0,0,0,0,0,0,0,par2(x[8],x[10]),par1()]])
    return Jacob

def corr(x, t=t):
    xnew = zeros((9,))
    for i in range(1,10):
        xnew[i-1] = -g(x, t, h, i)
    return xnew

def corr_solver(x, t=t):
    Jacob = Jac(x)
    xnew = corr(x)
    return solve(Jacob, xnew)

def damp(x,d):
    P = x
    P[1:-1] = (1/(2**(1.0*d)))*corr_solver(x) + array(x[1:-1])
    return P

def fdam(x, t, h, d):
    X = norm(corr(damp(x,d)))
    return X

def guess(x, t, h):
    for d in range(0,100):
        P = damp(x,d)
        if fdam(x, t, h, d)<(1-1/(2**(1.0*d+2)))*norm(array(corr(P))):
            break
    return damp(x,d)

def Newton(P, tol=10**(-9)):
    iterations = 0
    while norm(corr_solver(P))>tol:
        P[1:-1] += corr_solver(P)
        iterations += 1
    return P

def Damped_Newton(x, t=t, h=h):
    W = guess(x, t, h)
    while norm(W - x)>10**(-9):
        x = W
    return x

print Damped_Newton(x)
print "done"
plot(t,Damped_Newton(x),'o')
title('Position vs Time for spring damped system')
hold('true')
xlabel('Time (seconds)')
ylabel('Position (meters)')
show()
print "Done"