#!/usr/bin/env python

import pylab
import numpy as np
import time

def steepest_decent( f, x, tol, sf = 1.0, maxiter = 1000, verbose = False ):
    """Find minimum of multivarite function f(x).
    
    USAGE:
        y, err, iter = steepest_decent(f, x, tol, sf, maxiter, verbose)

    INPUT:
        f       - The objective function to minimize; must accept a single
                  argument which is an n-element array and return both the
                  function value at x and the gradient value at x.
        x       - Initial estimate of the minimum point; should be supplied
                  as an n-element list or array.
        tol     - Maximum allowable relative difference between consecutive
                  estimates of the location of the minimum.  This may be
                  either a scalar value or an n-element vector.
        sf      - Gradient step size scale factor.
        maxiter - Maximum number of iterations to perform.  Procedure will
                  terminate after this many iterations are done regardless
                  of the error.
        verbose - if True then information about each step is displayed.
  
    RETURNS:
        x       - Array containing coordinates of point which minimizes the
                  objective function.
        v       - value of function at x (the minimum value)
        iter    - Number of iterations performed.

    NOTES:
        This function uses the method of steepest decent to search for the
        minimum of a function f(x) where x is an n-element row vector that
        contains the values of the n independent variables used by f.

    AUTHOR:
        Jonathan R. Senning <jonathan.senning@gordon.edu>
        Gordon College
        December 4, 2008
    """

    # make sure x is a NumPy array and find the number of elements it contains
    n = len( x )
    x = np.array( x, dtype = float )

    # make sure the tolerance is a n-element NumPy array
    try:
        size = len( tol )
    except TypeError:
        tol = [tol] * n
    tol = np.array( tol, dtype = float )

    # if the user requests, we'll print out some information at each iteration
    if ( verbose ):
        print ' iter dif.norm         location'
        print '----- -------- -------------------------'
        print '%5d' % 0, '%8s' % ' ', x

    # main loop
    k = 0
    x0 = x + 2 * tol
    while k < maxiter and ( abs( x - x0 ) > tol ).any():
        k += 1
        x0 = x.copy()
        v, h = f( x )
        x = x - sf * h
        #sf *= 0.9
        if ( verbose ):
            print '%5d' % k, '%8.2e' % np.linalg.norm( x - x0 ), x
            ux = [x0[0],x[0]]
            vx = [x0[1],x[1]]
            pylab.plot( ux, vx, 'r-o' )
            pylab.draw()
            time.sleep( 0.5 )

    return ( x, v, k )

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

if __name__ == "__main__":

    def f( x ):
        a = 0.8
        y = ( x[0] - 1.0 )**2 + ( x[1] - 2.0 )**2 + a * np.cos( x[0] * x[1] )
        d1 = 2.0 * ( x[0] - 1.0 ) - a * np.sin( x[0] * x[1] ) * x[1]
        d2 = 2.0 * ( x[1] - 2.0 ) - a * np.sin( x[0] * x[1] ) * x[0]
        return ( y, np.array( [d1, d2] ) )


    n = 101
    x = pylab.linspace( -2.0, 4.0, n )
    y = pylab.linspace( -1.0, 5.0, n )
    xx = ( pylab.reshape( x, (n,1) ) * pylab.ones(n) ).transpose()
    yy = pylab.reshape( y, (n,1) ) * pylab.ones(n)
    z, g = f( [xx, yy] )

#    x0 = [3.8, 0.0]
#    x0 = [1, 4.8]
    x0 = [-0.35, 4.8]

    pylab.ion()
    pylab.contour( x, y, z, 41 )
    pylab.plot( x0[0], x0[1], 'r-o' )
    pylab.draw()

    raw_input( 'Press Enter to start...' )

    # search for minimum of ssqr() function

    sf1 = 0.5
    sf2 = 0.4
    sf3 = 0.3
    sf4 = 0.2
    sf5 = 0.1
    a, v, iter = steepest_decent( f, x0, 1e-2, sf4, 100, True )

    # report results

    print "a     = ", a
    print "value = ", v
    print "iter  = ", iter

    raw_input( 'Press Enter to quit...' )