from numpy import *
from numpy.linalg import *
from myarray import *

## plot, show, axes, clf
from pylab import *

def condition_number( A ):
    '''
    Returns the condition number of a matrix.
    The condition number is the ratio of absolute largest to absolute smallest eigenvalues (for symmetric matrices).
    '''
    eigenvalues = eig( A )[0]
    abseigenvalues = abs( eigenvalues )
    return max( abseigenvalues ) / min( abseigenvalues )

def lesson1():
    '''
    Prints the condition numbers of an identity matrix and an identity matrix with scaled rows.
    '''
    
    A = identityf(2)
    print 'condition_number( [1 0; 0 1] ):', condition_number( A )
    
    A[0,:] *= 1e20
    A[1,:] *= 1e-20
    print 'condition_number( [1e20 0; 0 1e-20] ):', condition_number( A )


def lesson2a():
    '''
    Plots all eigenvalues of our curve editing system.  Note the 'S' shape.  Keep in mind
    that the condition number is the largest divided by the smallest.
    '''
    
    import direct_linear
    
    steps = 101
    undeformed = zerosf(steps+1)
    h = 1
    
    A,b = direct_linear.curve_editing_1d_system( undeformed, h )
    #A,b = direct_linear.curve_editing_1d_system_fixed( {}, undeformed, h )
    
    ## Plot the n eigenvalues
    eigenvalues = eig(A)[0]
    eigenvalues.sort()
    plot( eigenvalues )
    plot( eigenvalues, 'go' )
    ylabel('eigenvalue')
    title('eigenvalues of the curve editing system')


def lesson2b():
    '''
    Plots the condition number of our curve editing system as a function of 'h', the resolution
    of our discrete curve approximation.
    '''
    
    import direct_linear
    
    steps = 21
    undeformed = zerosf(steps+1)
    fixed = {}
    
    ## Calculate the condition number as a function of h
    hs = arange( .01, .05, .001 )
    cs = [ condition_number( direct_linear.curve_editing_1d_system_fixed( fixed, undeformed, h )[0] ) for h in hs ]
    
    plot( hs, cs )
    plot( hs, cs, 'go' )
    xlabel('h'), ylabel('condition number')
    title('h versus condition number, curve editing system')



def improve_condition_number( A, b ):
    '''
    Improves the condition number of a system by scaling all coefficients
    in an equation so that the diagonal entry is 1.
    Returns the improved system A,b.
    NOTE: Systems where the diagonal entry is not the largest coefficient in a row (equation)
          should be pivoted to make this true.
    '''
    A = A.copy()
    b = b.copy()
    assert( A.shape[0] == A.shape[1] )
    assert( A.shape[0] == b.shape[0] )
    for row in range( A.shape[0] ):
        diag = A[ row, row ]
        oodiag = 1./diag
        A[ row, : ] *= oodiag
        b[ row ] *= oodiag
    return A,b

def residual( A, x, b ):
    '''
    Returns the residual of a system of equations Ax = b.
    The residual of a system of equations Ax = b is defined as |Ax-b|.
    If there were no roundoff error, this would be 0.
    '''
    r = dot( A, x ) - b
    return sqrt( sum( r ** 2 ) )


def lesson3():
    '''
    Print the original and improved condition numbers of our curve editing system.
    Print the residual of the solutions obtained from the original system and the improved system.
    '''
    import direct_linear
    
    steps = 101
    undeformed = zeros(steps+1)
    fixed = {steps//2: 1.0}
    h = .01
    
    A,b = direct_linear.curve_editing_1d_system_fixed( fixed, undeformed, h )
    x = solve( A, b )
    #x = dot( inv(A), b )
    c = condition_number( A )
    print 'condition number:', c
    
    print 'residual:', residual( A, x, b )
    
    Ap,bp = improve_condition_number( A, b )
    xp = solve( Ap, bp )
    #xp = dot( inv(Ap), bp )
    cp = condition_number( Ap )
    print 'condition number, improved:', cp
    
    print 'residual, improved:', residual( Ap, xp, bp )


def main():
    import matplotlib
    ion()
    
    lesson1()
    
    figure(1)
    lesson2a()
    
    figure(2)
    lesson2b()
    
    lesson3()
    
    show()


if __name__ == '__main__': main()