polynomial.py

#!/usr/stsci/pyssg/Python-2.7/bin/python

"""
polynomial.py

Created by Colin Cox on 2013-04-29.
Based on surfit.pro
"""

import scipy
from scipy import linalg
import matplotlib.pyplot as P
from math import *

def triangle(A, order):
    '''Print coefficients in triangular layout'''
    k = 0
    for i in range(order+1):
        for j in range(i+1):
            #print('%12.5e' %A[k], end=' ')
            print('{:f.3}'.format(A[k]))
            k += 1
        print()    

def triangulate(A, order):
    '''Convert linear array to 2-D array with triangular coefficient layout'''
    AT = scipy.zeros((order+1, order+1))
    k = 0
    for i in range(order+1):
        for j in range(i+1):
            AT[i,j] = A[k]
            k += 1
    return AT
    
def flatten(A, order):
    '''Convert triangular layout to linear array'''
    terms = (order+1)*(order+2)//2
    AF = scipy.zeros(terms)
    k = 0
    for i in range(order+1):
        for j in range(i+1):
            AF[k] = A[i,j]
            k += 1
    return AF        
                

def reorder(A,B, verbose=False) :
    '''Reorder Sabatke coefficients to my convention'''
    order = 5
    terms = (order+1)*(order+2)//2
    Aarray = scipy.zeros((order+1,order+1))
    Barray = scipy.zeros((order+1,order+1))
    #Carray = scipy.zeros((order+1,order+1))
    #Darray = scipy.zeros((order+1,order+1))
    k1 = 0
    for i in range(order+1):
        for j in range(order+1-i):
            Aarray[j,i] = A[k1]
            Barray[j,i] = B[k1]
            #Carray[j,i] = C[k1]
            #Darray[j,i] = D[k1]
            k1 += 1

    A2 = scipy.zeros((terms))
    B2 = scipy.zeros((terms))
    #C2 = scipy.zeros((terms))
    #D2 = scipy.zeros((terms))
    k2 = 0
    for i in range(order+1):
        for j in range(i+1):
            A2[k2] = Aarray[j,i-j]
            B2[k2] = Barray[j,i-j]
            #C2[k2] = Carray[j,i-j]
            #D2[k2] = Darray[j,i-j]
            k2 += 1

    if verbose:
        print('A')
        triangle(A2, order)        
        print('\nB')
        triangle(B2, order)        
        #print '\nC'
        #polyfit(C2, order)        
        #print '\nD'
        #triangle(D2, order)
                    
    return (A2, B2) #, C2, D2)       


def poly(a,x,y,order):
    pol = 0.0
    k = 0 # index for coefficients
    for i in range(order+1):
        for j in range(i+1):
            pol = pol + a[k]*x**(i-j)*y**j
            k+=1    
    return pol

def dpdx(a,x,y,order): # Differential wrt x
    dpdx = 0.0
    k = 1 # index for coefficients
    for i in range(1,order+1):
        for j in range(i+1):
            if i-j > 0: dpdx = dpdx + (i-j)*a[k]*x**(i-j-1)*y**j
            k+=1 
    return dpdx

def dpdy(a,x,y,order): # Differential wrt y
    dpdy = 0.0
    k = 1 # index for coefficients
    for i in range(1,order+1):
        for j in range(i+1):
            if j > 0: dpdy = dpdy + j*a[k]*x**(i-j)*y**(j-1)
            k+=1
    return dpdy
    
def jacob(a,b,x,y,order):
    '''Calculation of Jacobean, or relative area'''
    j = dpdx(a,x,y,order)*dpdy(b,x,y,order) - dpdx(b,x,y,order)*dpdy(a,x,y,order)
    j = scipy.fabs(j)
    return j


def invert(a,b,u,v,n,verbose=False):
    '''Given that order n polynomials of (x,y) have the result (u,v), find (x,y)
    Newton Raphson method in two dimensions'''
    tol = 1.0e-10
    err = 1.0
    #Initial guesses - Linear approximation
    det = a[1]*b[2] - a[2]*b[1]
    x0 = (b[2]*u-a[2]*v)/det
    y0 = (-b[1]*u + a[1]*v)/det
    if verbose: print('Initial guesses', x0, y0)
    x = x0
    y = y0
    X = scipy.array([x,y])
    iter = 0
    while err > tol:
        f1 = scipy.array([poly(a,x,y,n)-u, poly(b,x,y,n)-v])
        j = scipy.array([[dpdx(a,x,y,n), dpdy(a,x,y,n)],[dpdx(b,x,y,n), dpdy(b,x,y,n)]])
        invj = scipy.linalg.inv(j)
        X = X - scipy.dot(invj,f1)
        if verbose: print('[X1,Y1]', X)
        x1 = X[0]
        y1 = X[1]
        err = hypot(x-x1,y-y1)
        if verbose: print('Error %10.2e' % err)
        [x,y] = [x1,y1]
        iter += 1
    
    return  (x, y, err, iter)  
            
    
def choose(n,r):
    '''The number of ways of choosing r items from n'''
    if n<0 or r <0:
        print('Negative values not allowed')
        return 0
    if r > n:
        print('r must not be greater than n')
        return 0
            
    combin = 1
    if r > n/2: r1 = n-r
    else: r1 = r 
    for k in range(r1): combin = combin*(n-k)//(k+1) 
    return combin
        
def ShiftCoeffs(a,xshift, yshift, order, verbose=False):
    '''Calculate coefficients of polynomial when shifted to new origin'''

    # First place in triangular layout
    at = scipy.zeros([order+1,order+1])
    atshift = scipy.zeros([order+1,order+1])
    ashift = scipy.zeros([len(a)]) # Copy shape of a 
    k = 0
    for p in range(order+1):
        for q in range(p+1):
            at[p,q] = a[k]
            k+=1
            
    # Apply shift
    for p in range(order+1):
        for q in range(p+1):
            if verbose: print("A'%1d%1d" %(p,q))
            for i in range(p,order+1):
                for j in range(q,i+1-(p-q)):
                    f = choose(j,q)*choose(i-j,p-q) 
                    atshift[p,q] = atshift[p,q] + f*xshift**((i-j)-(p-q))*yshift**(j-q)*at[i,j]
                    if verbose: print('%2d A(%1d,%1d) x^%1d y^%1d' %(f, i , j, i-j-(p-q), (j-q)))
            if verbose: print()
                    
    # Put back in linear layout
    k = 0
    for p in range(order+1):
        for q in range(p+1):
            ashift[k] = atshift[p,q]
            k+=1
                               
    return ashift    







def polyfit(u,x,y,order):
    '''Fit polynomial to a set of u values on an x,y grid
    u is a function u(x,y) being a polynomial of the form
    u = a[i,j] x**(i-j) y**j. x and y can be on a grid or be arbitrary values'''
    
    
    # First set up x and y powers for each coefficient
    px = []
    py = []
    for i in range(order+1):
        for j in range(i+1):
            px.append(i-j)
            py.append(j)
    terms = len(px)        
    #print terms, ' terms for order ', order
    #print px
    #print py 
           
    # Make up matrix and vector
    vector = scipy.zeros((terms))
    mat = scipy.zeros((terms,terms))
    for i in range(terms):
        vector[i] = (u*x**px[i]*y**py[i]).sum()
        for j in range(terms):
            mat[i,j] =  (x**px[i]*y**py[i]*x**px[j]*y**py[j]).sum()
            
    #print 'Vector', vector
    #print 'Matrix'
    #print mat
    imat = linalg.inv(mat)
    #print 'Inverse'
    #print imat
    # Check that inversion worked
    #print scipy.dot(mat,imat)
    coeffs = scipy.dot(imat,vector)        
    return coeffs
                        


def polyfit2(u,x,y,order):
    '''Fit polynomial to a set of u values on an x,y grid
    u is a function u(x,y) being a polynomial of the form
    u = a[i,j]x**(i-j)y**j. x and y can be on a grid or be arbitrary values
    This version uses solve instead of matrix inversion'''


    # First set up x and y powers for each coefficient
    px = []
    py = []
    for i in range(order+1):
        for j in range(i+1):
            px.append(i-j)
            py.append(j)
    terms = len(px)        
    #print terms, ' terms for order ', order
    #print px
    #print py 

    # Make up matrix and vector
    vector = scipy.zeros((terms))
    mat = scipy.zeros((terms,terms))
    for i in range(terms):
        vector[i] = (u*x**px[i]*y**py[i]).sum() # Summing over all x,y
        for j in range(terms):
            mat[i,j] =  (x**px[i]*y**py[i]*x**px[j]*y**py[j]).sum()

    coeffs = linalg.solve(mat,vector)        
    return coeffs

                
def testpoly():
    [x,y] = scipy.mgrid[0:10,0:10]
    #print 'X'
    #print x
    #print 'Y'
    #print y
    u = scipy.zeros((10,10))
    v = scipy.zeros((10,10))
    # Random polynomials
    a0 = scipy.random.rand(1)
    a1 = 0.1*(scipy.random.rand(2)-0.5)
    a2 = 0.01*(scipy.random.rand(3)-0.5)
    a = scipy.concatenate((a0,a1))
    a = scipy.concatenate((a,a2))
    a[2] = 0.01*a[2]
    print('A coefficients')
    print(a)
    b0 = scipy.random.rand(1)
    b1 = 0.1*(scipy.random.rand(2)-0.5)
    b2 = 0.01*(scipy.random.rand(3)-0.5)
    b = scipy.concatenate((b0,b1))
    b = scipy.concatenate((b,b2))
    b[1] = 0.01*b[1]
    print('B coeffcicients')
    print(b)
    for i in range(10):
        for j in range(10):
            u[i,j] = poly(a,x[i,j], y[i,j],2) #+ scipy.random.normal(0.0, 0.01) 
            v[i,j] = poly(b,x[i,j], y[i,j],2)  #+ scipy.random.normal(0.0,0.01)      
    #print z
    s1 = polyFit2(u,x,y,2)
    s2 = polyFit2(v,x,y,2)
    print('S1', s1)
    print('S2', s2) 
    uc = poly(s1,x,y,2)
    vc = poly(s2,x,y,2)
    
    P.figure(1)
    P.clf()
    P.grid(True)
    P.plot(u,v,'gx')
    P.plot(uc,vc,'r+') 
    

def FlipX(A, order):
    """Change sign of all coefficients with odd x power"""
    terms = (order+1)*(order+2)//2
    AF = scipy.zeros(terms)
    k = 0
    for i in range(order+1):
        for j in range (i+1):
            AF[k] = (-1)**(i-j)*A[k]
            k += 1
    return  AF       


def FlipY(A, order):
    """Change sign of all coefficients with odd y power"""
    terms = (order+1)*(order+2)//2
    AF = scipy.zeros(terms)
    k = 0
    for i in range(order+1):
        for j in range (i+1):
            AF[k] = (-1)**(j)*A[k]
            k += 1
    return AF
    
    
def FlipXY(A, order):
    "Change sign for coeffs where sum of x and y powers is odd"
    terms = (order+1)*(order+2)//2
    AF = scipy.zeros(terms)
    k = 0
    for i in range(order+1):
        for j in range (i+1):
            AF[k] = (-1)**(i)*A[k]
            k += 1
    return AF
    
                

def RotateCoeffs(a, theta, order, verbose=False):
    '''Rotate axes of coefficients by theta degrees'''
    c = cos(radians(theta))
    s = sin(radians(theta))

    # First place in triangular layout
    at = scipy.zeros([order+1,order+1])
    k = 0
    for m in range(order+1):
        for n in range(m+1):
            at[m,n] = a[k]
            k+=1

    # Apply rotation
    atrotate = scipy.zeros([order+1,order+1])
    arotate = scipy.zeros([len(a)]) # Copy shape of a 
    for m in range(order+1):
        for n in range(m+1):
            for mu in range(0,m-n+1):
                for j in range(m-n-mu, m-mu+1):
                    factor = (-1)**(m-n-mu)*choose(m-j,mu)*choose(j,m-n-mu)
                    cosSin = c**(j+2*mu-m+n)*s**(2*m-2*mu-j-n)
                    atrotate[m,n] = atrotate[m,n] + factor*cosSin*at[m,j]
                    if verbose: print(m, n, j, factor, 'cos^', j+2*mu-m+n, 'sin^',2*m-2*mu-j-n, ' A',m,j)
    # Put back in linear layout
    k = 0
    for m in range(order+1):
        for n in range(m+1):
            arotate[k] = atrotate[m,n]
            k+=1

    return arotate        


def TransCoeffs(A, a,b,c,d, order, verbose=False):
    '''Transform polynomial coefficients to allow for
    xp = a*x + b*y
    yp = c*x + d*y'''
    
    A1 = scipy.zeros((order+1, order+1))
    A2 = scipy.zeros((order+1, order+1))
    ncoeffs = (order+1)*(order+2)//2 
    if verbose: print(ncoeffs, 'coefficients for order', order)
    AT = scipy.zeros((ncoeffs)) 
    
    # First place A in triangular layout
    k = 0
    for i in range(order+1):
        for j in range(i+1):
            A1[i,j] = A[k]
            k+=1

    for m in range(order+1):
        for n in range(m+1):
            if verbose: print('\nM,N',m,n)
            for mu in range(m-n+1):
                for j in range(m-n-mu,m-mu+1):
                    if verbose: print('J, MU', j,mu)
                    if verbose: print('Choose', m-j, mu, 'and', j, m-n-mu)
                    factor = choose(m-j,mu)*choose(j,m-n-mu)
                    A2[m,n] += factor*a**mu*b**(m-j-mu)*c**(m-n-mu)*d**(mu+j-m+n)*A1[m,j]
                    if verbose: print(m, j, ' Factor', factor)
    # Restore A2 to flat layout in AT
    k = 0
    for m in range(order+1):
        for n in range(m+1):
            AT[k] = A2[m,n]
            k += 1       
    return AT                       
                    
                    
def TwoStep(A, B, a, b, order):
    """Change coefficients when
    xp = a[0] + a[1].x + a[2].y
    yp = b[0] + b[1].x + b[2].y"""
    terms = (order+1)*(order+2)//2
    A2 = scipy.zeros((order+1,order+1))
    B2 = scipy.zeros((order+1,order+1))
    
    k=0
    for i in range(order+1):
        for j in range(i+1):
            for alpha in range(i-j+1):
                for beta in range(i-j-alpha+1):
                    f1 = choose(i-j,alpha)*choose(i-j-alpha,beta)*a[0]**(i-j-alpha-beta)*a[1]**alpha*a[2]**beta
                    for gamma in range(j+1):
                        for delta in range(j-gamma+1):
                            f2 = choose(j,gamma)*choose(j-gamma,delta)*b[0]**(j-gamma-delta)*b[1]**gamma*b[2]**delta
                            A2[alpha+beta+gamma+delta,beta+delta] += A[k]*f1*f2
                            B2[alpha+beta+gamma+delta,beta+delta] += B[k]*f1*f2
            k += 1
    
    # Flatten A@ and B2                    
    k = 0
    Aflat = scipy.zeros(terms)
    Bflat = scipy.zeros(terms)
    for i in range(order+1):
        for j in range(i+1):
            Aflat[k] = A2[i,j]
            Bflat[k] = B2[i,j]
            k += 1                      
    return (Aflat, Bflat)     
    
def TestTwoStep():
    A = scipy.array([10.0, 2.0, 0.1, 0.01, -0.02, 0.03])
    B = scipy.array([4.0, 1.8, 0.2, 0.02, 0.03, -0.02])
    a = scipy.array([1.0, 0.5, 0.1])
    b = scipy.array([2.0, 0.2, 0.6])
    print('\nA')
    triangle(A,2)
    print('B')
    triangle(B,2)
    print('a\n',a)
    print('b\n', b)
    (A2, B2) = TwoStep(A,B,a,b,2)
    print('\nA2')
    triangle(A2,2)
    print('B2')
    triangle(B2,2)
    
    # Now do a test calculation
    (x,y) = (10,5)
    xp = a[0] + a[1]*x + a[2]*y
    yp = b[0] + b[1]*x + b[2]*y
    print('x,y', x,y)
    print('xp,yp', xp,yp)
    
    u = poly(A, xp, yp, 2)
    v = poly(B, xp, yp, 2)
    up = poly(A2,x,y,2)
    vp = poly(B2,x,y,2)
    print('Two step', u, v)
    print('One step', up, vp)
    return