#!/net/python/bin/python

import pylab as p
import numpy as num
from numpy import random

#---------------------------------------------------------------

def lfit(x, y, sig, func, ma):
    """
    Given a set of data points x[1..ndat], y[1..ndat]
    with individual standard deviations sig[1..ndat],
    use chi^2 minimization to fit for
    some or all of the coefficients a[1..ma]
    of a function that depends linearly on a,
    y = Sum_i[ ai * afunci(x)].

    The input array ia[1..ma] indicates by nonzero entries
    those components of a that should be fitted for, and by
    zero entries those components that should be held fixed
    at their input values.

    func supplies the basis functions: assumed to be called
    in the form func(x,i) = F_i(x), where i=0,1,2...

    The program returns values for a[1..ma], chisq,
    and the covariance matrix covar[1..ma][1..ma]. (Parameters
    held fixed will return zero covariances.)

    """

    ndat = len(x)
    assert (len(y) == ndat) and (len(sig) == ndat)

    #create some arrays needed for the routine
    a = num.zeros(ma)
    beta = num.zeros([ma,1])
    afunc = num.zeros(ma)
    #covar will hold the covariance matrix
    covar = num.zeros([ma,ma])

    for i in range(ndat): #Loop over data to accumulate coefficients of
        afunc = [func(x[i],mi) for mi in range(ma)] #the normal equations.
        ym=y[i]

        sig2i=1.0/(sig[i]**2)
        j=0
        for l in range(ma):
            wt=afunc[l]*sig2i
            for k in range(l+1):
                covar[j][k] += wt*afunc[k]
            beta[j][0] += ym*wt
            j+=1

    for j in range(1,ma): # Fill in above the diagonal from symmetry.
        for k in range(j-1):
            covar[k][j]=covar[j][k]
    gaussj(covar,beta) #Matrix solution.
    
    j=0
    for l in range(ma):
        a[l]=beta[j][0] #Partition solution to 
        j+=1            # appropriate coefficients a.
    
    chisq=0.0
    for i in range(ndat): # Evaluate chi2 of the fit.
        afunc = [func(x[i],mi) for mi in range(ma)]
        sum = 0.0
        for j in range(ma):
            sum += a[j]*afunc[j]
        chisq += ((y[i]-sum)/sig[i])**2

    return a,covar,chisq

#---------------------------------------------------------------

def gaussj(a,b):
    """
    Linear equation solution by Gauss-Jordan elimination,
    equation (2.1.1) in Press NRC.
     a[1..n][1..n] is the input matrix.
     b[1..n][1..m] is input containing the m right-hand side vectors.
    On output, a is replaced
    by its matrix inverse, and b is replaced by the corresponding
    set of solution vectors.
    """
    n,m = b.shape
    assert a.shape == (n,n)
    
    indxc=num.zeros(n) #The integer arrays ipiv, indxr, 
    indxr=num.zeros(n) #and indxc are used for
    ipiv=num.zeros(n)  #bookkeeping on the pivoting.
    
    for i in range(n): #This is the main loop over the 
        big=0.0        #  columns to be reduced.
        for j in range(n):     # This is the outer loop of the search 
            if ipiv[j] != 1: # for a pivot element.
                for k in range(n):
                    if ipiv[k] == 0:
                        if num.abs(a[j][k]) >= big:
                            big=num.abs(a[j][k])
                            irow=j
                            icol=k
        ipiv[icol] += 1
        
        #We now have the pivot element, so we interchange rows,
        #if needed, to put the pivot
        #element on the diagonal. The columns are not
        #physically interchanged, only relabeled:
        #indxc[i], the column of the ith pivot element,
        #is the ith column that is reduced, while
        #indxr[i] is the row in which that pivot element
        #was originally located. If indxr[i] =
        #indxc[i] there is an implied column interchange.
        #With this form of bookkeeping, the
        #solution b's will end up in the correct order,
        #and the inverse matrix will be scrambled by columns.
        if irow != icol:
            for l in range(n):
                a[irow][l],a[icol][l] = a[icol][l],a[irow][l]
            for l in range(m):
                b[irow][l],b[icol][l] = b[icol][l],b[irow][l]
        indxr[i]=irow #We are now ready to divide the pivot row by the
        indxc[i]=icol # pivot element, located at irow and icol.
        if (a[icol][icol] == 0.0):
            raise ValueError, "Error: gaussj: Singular Matrix"
        pivinv=1.0/a[icol][icol]
        a[icol][icol]=1.0
        for l in range(n): a[icol][l] *= pivinv
        for l in range(m): b[icol][l] *= pivinv
        

        for ll in range(n): #Next, we reduce the rows...
            if (ll != icol): #...except for the pivot one, of course.
                dum=a[ll][icol]
                a[ll][icol]=0.0
                for l in range(n): a[ll][l] -= a[icol][l]*dum
                for l in range(m): b[ll][l] -= b[icol][l]*dum
    #This is the end of the main loop over columns of the reduction.
    #It only remains to unscramble the solution in view of the
    #column interchanges. We do this by interchanging pairs of
    #columns in the reverse order that the permutation was built up.
    for l_m in range(1,n+1):
        l = n-l_m
        if indxr[l] != indxc[l]:
            for k in range(n):
                a[k][indxr[l]],a[k][indxc[l]]=a[k][indxc[l]],a[k][indxr[l]]
    #And we are done.
    
#-------------------------------------------------------------

if __name__ == '__main__':

    def value(a,basis,x):
        sum = 0
        for i in range(len(a)):
            sum += a[i]*basis(x,i)
        return sum

    def sin_basis(x,i,xmin=0.0,xmax=100.0):
        return num.sin( (i+1) * num.pi*(x-xmin)/(xmax-xmin) )

    def poly_basis(x,i):
        return x**i

    basis = sin_basis
    
    x = num.sort([random.random()*100 for i in range(20)])

    y = [(num.sin(xi*num.pi/100)+0.3*(random.random()-0.5))**2 for xi in x]

    dy = [0.1*(random.random()-0.5) for xi in x]

    a,covar,chi2 = lfit(x, y, dy, func=basis, ma=20)

    y_comp = [value(a,basis,xi) for xi in x]

    print a
    print chi2
    
    p.errorbar(x,y,dy,fmt='.k')
    p.plot(x,y_comp)
    p.show()