from scipy import *
from scipy import linalg, integrate
n = 3
#
#  Linear eigenvalue problem : K x = lambda M x
#
K = array(([1.0,-1.0,0.0],
[-1.0,2.0,-1.0],
[0.0,-1.0,2.0]))
#
M = array(([2.0,1.0,0.0],
[1.0,4.0,1.0],
[0.0,1.0,4.0]))/6.0
#
w, vr = linalg.eig(K,M)

def P(t,z):

    return array(([z[n]/tan(z[n]),-z[n]/sin(z[n]),0.0],
    [-z[n]/sin(z[n]),2.0*z[n]/tan(z[n]),-z[n]/sin(z[n])],
    [0.0,-z[n]/sin(z[n]),2.0*z[n]/tan(z[n])]))

def Pl(t,z):

    return array(([cos(z[n])-z[n],z[n]*cos(z[n])-sin(z[n]),0.0],
    [z[n]*cos(z[n])-sin(z[n]),2.0*(cos(z[n])-z[n]),z[n]*cos(z[n])-sin(z[n])],
    [0.0,z[n]*cos(z[n])-sin(z[n]),2.0*(cos(z[n])-z[n])]))/sin(z[n])**2
 
def A(P, Pl, M, K, t,z):
    tmp = zeros((n,n),Float)
    tmp.flat[0:n:n+1] = (1.0-t)*(K.flat[0:n:n+1]-z[n]*M.flat[0:n:n+1])+t*P.flat[0:n:n+1]
    tmp[0:n,n-1]   = -(1-t)*dot(M,z[:n])+t*dot(Pl,z[:n])
    tmp[n-1,0:n]   = z[:n]
    return tmp

def f(P, M, K, t, z):
    tmp = dot(K,z[:n])-z[n]*dot(M,z[:n])-dot(P,z[:n])
    tmp[n-1] = 0.0
    return tmp

def F(z, t):
    p = P(t,z)
    pl = Pl(t,z)
    x = linalg.solve(A(p, pl, M, K, t, z), f(p, M, K, t, z))
    return x

t = arange(0, 0.1, 0.05)  # or whatever
z0=zeros(n+1,Float)
z0[0:n] = vr[:,0]               
z0[n] = abs(w[0])
print z0

z = integrate.odeint(F, z0, t)
print z