import numpy as N

########################################################
def tridag(a, b, c, r):
    """Solve linear system with tridiagonal coefficient matrix.
    a is the lower band, b is the diagonal, c is the upper band, and r
    is the right hand side. The solution is returned in u.
    [b1 c1 0 ...           ] [u1]   [r1] 
    [a1 b2 c2 0 ...        ] [ :]   [ :] 
    [ 0 a2 b3 c3 0 ...     ] [  ] = [  ] 
    [                      ] [  ]   [  ] 
    [ ...  0 an-2 bn-1 cn-1] [ :]   [ :] 
    [        ... 0 an-1 bn ] [un]   [rn] """ 

    n = len(b)
    tmp = N.zeros(n-1) # necessary temporary array 
    if b[0] == 0: 
        raise RuntimeError, 'System is effectively order N-1' 

    beta = b[0]
    u=N.zeros_like(r)
    u[0] = r[0] / beta 
    for i in range(1, n): # Decompose and forward substitution 
        tmp[i-1] = c[i-1] / beta 
        beta = b[i] - a[i-1] * tmp[i-1] 
        if beta == 0: 
            raise RuntimeError, 'Method failure'
        u[i] = (r[i] - a[i-1] * u[i-1]) / beta 

    for i in range(n-1, 0, -1): # Backward substitution 
        u[i-1] -= tmp[i-1] * u[i] 

    return u

########################################################
def cyclic(a, b, c, alpha, beta, r):
    """Solve linear system with tridiagonal coefficient matrix with
    corners alpha and beta (i.e. for a cyclic matrix).
    a is the lower band, b is the diagonal, c is the upper band, and r
    is the right hand side. The solution is returned in u.
    [b1 c1 0 ...       beta] [u1]   [r1] 
    [a1 b2 c2 0 ...        ] [ :]   [ :] 
    [ 0 a2 b3 c3 0 ...     ] [  ] = [  ] 
    [                      ] [  ]   [  ] 
    [ ...  0 an-2 bn-1 cn-1] [ :]   [ :] 
    [alpha   ... 0 an-1 bn ] [un]   [rn] """ 
#
    n=len(b)
    bb=N.zeros_like(b)
    u=N.zeros_like(b)
#
    gamma=-b[0]
    bb[0]=b[0]-gamma
    bb[-1]=b[-1]-alpha*beta/gamma
    for i in range(1,n-1):
    	bb[i]=b[i]
    x=tridag(a,bb,c,r)
    u[0]=gamma
    u[n-1]=alpha
    for i in range(1,n-1):
    	u[i]=0.
    z=tridag(a,bb,c,u)
    fact=(x[0]+beta*x[n-1]/gamma)/(1.+z[0]+beta*z[n-1]/gamma)
    for i in range(n):
    	x[i]=x[i]-fact*z[i]
    return x