#Assumes numpy has been imported
from ClimateUtilities import *

def TriDiag(a,b,c,r):

#Converted to Python May 16, 2007
#/*    Dec 12, 2002. Adapted from numpyal recipes, to run standalone without
#   the utility routines, and to use globally declared workspace storage. 
#   Returns solution if successful, 'Error' if not
#   Global Workspace *gam must be allocated as float with dimensin NMAX */
    n = len(a)
    gam = numpy.zeros(n,numpy.Float)
    u = numpy.zeros(n,numpy.Float) #for returning result
    if b[0]== 0.0:
        return 'Error' #Not invertible
    beta = b[0]
    u[0]=r[0]/beta;
    for j in range(1,n):
       gam[j]=c[j-1]/beta
       beta=b[j]-a[j]*gam[j]
       if beta == 0.0:
            return 'Error' 
       u[j]=(r[j]-a[j]*u[j-1])/beta
    for j in range(n-2,-1,-1):
       u[j] -= gam[j+1]*u[j+1]
    return u

#Now test the tridiagonal solver
N = 1000
a = numpy.zeros(N,numpy.Float)
b = numpy.zeros(N,numpy.Float)
c = numpy.zeros(N,numpy.Float)
r = numpy.zeros(N,numpy.Float)


lam = -.001
b[:] = -2. + lam
a[:] = 1.
c[:] = 1.
a[0]=c[0] = 0.
a[-1]=c[-1] = 0.
r[N/4] = 1.
b[0] = b[-1] = 1.

u = TriDiag(a,b,c,r)
c = Curve()
c.addCurve(range(N))
c.addCurve(u)
plot(c)