#import for rosenbrock plot:
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.colors import LogNorm
import matplotlib.pyplot as plt
import numpy as np
# other
from pylab import *


#import for rosenbrock function
from scipy.optimize import fmin
from scipy.optimize import fmin_powell
from scipy.optimize import fmin_ncg
from scipy.optimize import leastsq


#rosenbrock function code using fmin
def rosen(x):
	return sum(100.0*(x[1:]-x[:-1]**2.0)**2.0 + (1-x[:-1])**2.0)
	
def r_hess(x):
        x = asarray(x)
        H = diag(-400*x[:-1],1) - diag(400*x[:-1],-1)
        diagonal = zeros(len(x),x.typecode())
        diagonal[0] = 1200*x[0]-400*x[1]+2
        diagonal[-1] = 200
        diagonal[1:-1] = 202 + 1200*x[1:-1]**2 - 400*x[2:]
        H = H + diag(diagonal)
        return H

def r_der(x):
        xm = x[1:-1]
        xm_m1 = x[:-2]
        xm_p1 = x[2:]
        der = zeros(x.shape,x.typecode())
        der[1:-1] = 200*(xm-xm_m1**2) - 400*(xm_p1 - xm**2)*xm - 2*(1-xm)
        der[0] = -400*x[0]*(x[1]-x[0]**2) - 2*(1-x[0])
        der[-1] = 200*(x[-1]-x[-2]**2)
        return der
        
#x0 = [1.3,0.7,0.8,1.9,1.2] #Start location: an array that represents an initial guess for the x parameters which minimizes the scalar function.
x0 = np.arange(-2,2.)

#Methods to analyze rosenbrock func: 
#xopt = fmin(rosen,x0)       #nelder-mead
#xopt = fmin_powell(rosen,x0) #direction-set
xopt = fmin_ncg(rosen,x0,r_der,fhess=r_hess)    #conjugate gradient method
#xopt = leastsq(rosen,x0)   #Levenberg-Marquardt


#rosenbrock plot code:
fig = plt.figure()
#ax = Axes3D(fig, azim = -128, elev = 43)
#ax = Axes3D(fig, azim = 115, elev =45) good
ax = Axes3D(fig, azim = 115, elev =90)

s = .05
X = np.arange(-2, 2.+s, s) #arange(start,finish,increment), stores resulting vector in X
Y = np.arange(-1, 3.+s, s)
X, Y = np.meshgrid(X, Y)   #create the mesh grid
Z = (1.-X)**2 + 100.*(Y-X*X)**2 # rosenbrock function
ax.plot_surface(X, Y, Z, rstride = 1, cstride = 1, norm = LogNorm(), cmap = cm.jet)

#add
CS = plt.contour(X,Y,Z) #plot contour

plt.clabel(CS,inline=1,fontsize=10)
CB = plt.colorbar(CS, shrink=0.8, extend='both')  #colorbar


#ax.plot(xs=X,ys=[0]*len(X),zs=Y,zdir='z')
 #Notes: 
 #rstride = row stride (step size)
 #cstride = column stride (step size)
plt.xlabel("x")
plt.ylabel("y")
 
#plt.savefig("Rosenbrock function.svg")

plt.show()