#!/net/python/bin/python
import math

#Contains:
#   - smooth: a physics based algorithm to smooth noisy data
#   - splrep & splev: cubic spline interpolation algorithms
#   - integrate: Romberg integration algorithm

def smooth(x,y,dy=None,s=0.1,MaxSteps=200,tol=0.05):
    """
    given an x, y and dy, smooths the data using a physics
    based method:
     Each uncertainty is treated as a spring with restoring
     force F_s, with a rod connected between them.  The rod
     is given a restoring force F_r proportional to its
     curvature.  The system is iterated over until an equilibrium
     is reached, with F_s+F_r<tol at each point.
    if dy == None, then dyi's are set to 1
    """

    #make sure data is in correct format
    if dy is None:
        dy = [1 for yi in y]
    n = len(x)
    if (len(y) != n) or (len(dy) != n):
        raise ValueError, 'x,y,dy must be same length!'

    #now add points so that
    #  x[-1] = x[0]
    #  x[n] = x[n-1]
    #this simplifies the boundary conditions in loops
    x = [xi for xi in x] + [2*x[-1]-x[-2]] + [2*x[0]-x[1]]
    y = [yi for yi in y] + [y[-1]] + [y[0]]
    dy = [dyi for dyi in dy] + [dy[-1]] + [dy[0]]

    #y2 initialization
    y2 = [yi for yi in y]

    #holder for spring force & line restoring force
    F_s = [0.0 for i in x]
    F_r = [0.0 for i in x]

    for steps in range(MaxSteps):
        #first calculate all forces involved
        for i in range(n):
            F_s[i] = ( y[i]-y2[i] )/dy[i]
            F_r[i] = s * ( (y2[i+1]-y2[i])/(x[i+1]-x[i])\
                           - (y2[i]-y2[i-1])/(x[i]-x[i-1]) )*\
                           2/(x[i+1]-x[i-1])

        F_t = [ F_s[i]+F_r[i]\
                -0.5*( F_r[i+1]+F_r[i-1] )\
                for i in range(n) ]

        #factor effectively sets the time-step small enough
        # so our solution doesn't blow up
        factor = math.sqrt(steps+1) * max( [ abs(F_t[i]/dy[i] )\
                                          for i in range(n)] )
        if factor > 1.0:
            F_t = [val/factor for val in F_t]

        #add the new displacement dy = 0.5 dt^2 * F/m
        # other physical constants are set by factor, above
        for i in range(n):
            y2[i]+=F_t[i]

        #if net force at each point becomes small,
        # we are finished iterating
        if factor * max(abs(fi) for fi in F_t) < tol:
            break

    print 'smooth: stopped after',steps+1,'steps.'

    return y2[:n]

def splrep(x, y, yp1 = 1E30, ypn = 1E30):
    """
    Adapted from Press et al, Numerical Recipies in C

    Given arrays x[1..n] and y[1..n] containing a tabulated function,
    i.e., yi = f(xi), with x1 < x2 < .. . < xN, and given
    values yp1 and ypn for the first derivative of the interpolating
    function at points 1 and n, respectively, this routine returns a
    tuple (x,y,y2) where y2[1...n] contains the second derivatives of
    the interpolating function at the points xi.
    
    If yp1 and/or ypn are equal to 1 * 10^30 or larger, the routine is
    signaled to set the corresponding boundary condition for a natural
    spline, with zero second derivative on that boundary.
    """
    n = len(x)
    if len(y) != n:
        raise ValueError, "x and y must be equal length arrays"

    y2 = [0.0 for i in y]
    u = [0.0 for i in y]
    
    if (yp1 > 0.99E30): #lower boundary condition is y''=0
        y2[0] = 0.0
        u[0] = 0.0
    else:               #lower boundary has a specified first derivative.
        y2[0] = -0.5
        u[0]=(3.0/(x[1]-x[0]))*((y[1]-y[0])/(x[1]-x[0])-yp1)

    for i in range(1,n-1): #decomposition loop of the tridiagonal algorithm.
        sig=(x[i]-x[i-1])/(x[i+1]-x[i-1])
        p=sig*y2[i-1]+2.0
        y2[i]=(sig-1.0)/p 
        u[i]=(y[i+1]-y[i])/(x[i+1]-x[i]) - (y[i]-y[i-1])/(x[i]-x[i-1])
        u[i]=(6.0*u[i]/(x[i+1]-x[i-1])-sig*u[i-1])/p

    if (ypn > 0.99E30):#upper boundary condition is y''=0
        qn = 0.0 
        un = 0.0
    else:              #upper boundary has specified first derivative. 
        qn=0.5
        un=(3.0/(x[n-1]-x[n-2]))*(ypn-(y[n-1]-y[n-2])/(x[n-1]-x[n-2]))

    y2[n-1]=(un-qn*u[n-2])/(qn*y2[n-2]+1.0);
    for i in range(2,n): #backsubstitution loop of the tridiagonal
        k = n-i          # algorithm.
        y2[k]=y2[k]*y2[k+1]+u[k]

    return x,y,y2


def splev(x, tuple):
    """
    Adapted from Press et al, Numerical Recipies in C
    
    Given the tuple (xa,ya,y2a) which is the output of splrep()
    above, and given a values of array x, this routine returns an
    array of cubic-spline interpolated values y.
    """
    xa = tuple[0]
    ya = tuple[1]
    y2a = tuple[2]
    n = len(xa)
    if len(ya) != n or len(y2a) != n:
        raise ValueError, "xa,ya,y2a must be same length"
     #We will find the right place in the table by means of
     # bisection, using previous values of klo and khi to
     # optimize in case x's are in order (but this condition
     # is not necessary)
    y = []
    klo=0
    khi=n-1
    for xi in x:
        if xa[klo] > xi: klo = 0
        if xa[khi] < xi: khi = n-1
        while (khi-klo > 1):
            k= (khi+klo) >> 1 #floor of average value
            if (xa[k] > xi):
                khi=k
            else:
                klo=k      #klo and khi now bracket the input value of x.
        h=xa[khi]-xa[klo]

        #make sure the xa's are distinct
        if h == 0.0: raise ValueError, "Bad xa input to routine splint"

        a=(xa[khi]-xi)/h
        b=(xi-xa[klo])/h #Cubic spline polynomial is now evaluated.
        yi = a*ya[klo]+b*ya[khi]+\
            ((a*a*a-a)*y2a[klo]+(b*b*b-b)*y2a[khi])*(h*h)/6.0
        y.append(yi)

    return y



def integrate(f,xmin,xmax,acc = 1E-6,MaxIterations = 25):
    """
    Approximate the definite integral of f from a to b by Romberg's method.
    acc is the desired accuracy.
    """
    R = [[0.5 * (xmax - xmin) * (f(xmin) + f(xmax))]]  # R[0][0]
    times_within_limits = 0
    
    for i in range(1,MaxIterations):
        h = float(xmax-xmin)/2**i
        R.append((i+1)*[None])  # Add an empty row.
        R[i][0] = 0.5*R[i-1][0] + \
                  h*sum(f(xmin+(2*k-1)*h) for k in range(1, 2**(i-1)+1))
        for k in range(1, i+1):
            R[i][k] = R[i][k-1] + (R[i][k-1] - R[i-1][k-1]) / (4**k - 1)

        #check for convergence after i == 5
        # important for oscillating functions if the
        # N points are in resonance
        if i>=5 and abs(R[i][i]-last) <= acc:
            times_within_limits += 1
        else:
            times_within_limits = 0

        #if it's really converging, it will be within errors twice
        if times_within_limits == 2:
            return R[i][i]
        last = R[i][i]
    print "Warning: Romberg Integration did not converge"
    return last


if __name__ == '__main__':
    pass