from numpy import *
	
class NelderMead:
	"""Implementation of Nelder Mead Algorithm"""	
		
	def __init__(self, f):
		self.function = f
		
	def Run(self):				
		function = self.function		
		ndim = 2 # 2D problem
		npts = ndim + 1 # a simplex
		# Define the starting simplex and evaluate the function at its vertices
		x = zeros((npts, ndim))
		x[0,:] = 2 # Starting point
		for i in range(1,npts,1):
			x[i,:] = x[i-1,:] + array([0.1,0.2])*i
		f = apply_along_axis(function,1,x)
		# Sort the simplex and the vector of function values
		index = argsort(f, axis = 0) # best - middle - worst
		x = x[index,:] 
		f = f[index]
		# set up storage for simplex path plotting
		self.path = copy(x)
		fpath = copy(f)
		converged = False
		while (converged==False):
			# 1. Find the center of the face of the simplex (except the worst point)
			x_mean = mean(x[0:(npts-1),:], axis = 0)    
			# 2. Reflect the worst point across the face
			xr = 2*x_mean-x[(npts-1),:]     
			if function(xr) <= f[0]: # is the new point the new best?
				x2r = 3*x_mean-2*x[(npts-1),:] # grow further in that direction        
				if function(x2r) < function(xr): # did the growing improve?
					x[(npts-1),:] = x2r # accept the grown point and go to the next iteration
					f[(npts-1)] = function(x2r)
					print 'accepted grown point'
				else: # the growing step did not help
					x[(npts-1),:] = xr # accept the original reflected point and go to the next iteration
					f[(npts-1)] = function(xr)            
					print 'accepted reflected point (it was the best)'
			elif function(xr) >= f[(npts-2)]: # is the new point still the worst?
				xs = 1.5*x_mean - 0.5*x[(npts-1),:] # reflect + shrink
				if function(xs) < f[(npts-2)]: # better than the second worst?
					x[(npts-1),:] = xs # accept and move on
					f[(npts-1)] = function(xs)
					print 'accepted reflected & shrunk point'
				else: # not better than the second worst
					xs = 0.5*x_mean+0.5*x[(npts-1),:] # just shrink (no reflection)
					if function(xs) > f[(npts-2)]:
						x[1:npts,:] = 0.5*(x[1:npts,:] + x[0,:]) # shrink everything (except the best, 0) towards the best
						f = apply_along_axis(function, 1, x)
						print 'shrank everything towards the best'
					else:
						x[(npts-1),:] = xs # accept and move on
						f[(npts-1)] = function(xs)
						print 'accepted shrunk only point'                          
			else: # the new reflected point is neither the best nor the worst. keep it
				x[(npts-1),:] = xr
				f[(npts-1)] = function(xr)
				print 'accepted reflected point (though it was not a new best)'   
			# Sort the simplex and the vector of function values - think of a better way that doesn't need sorting every time
			index = argsort(f, axis = 0) # best - middle - worst
			x = x[index,:] 
			f = f[index]    
			# Record the history of the simplex:
			self.path = concatenate((self.path,x),axis = 0)			
			if fabs(f[0]- f[(npts-1)]) < 1e-14:
				converged = True
		print '\nconverged at: ' + str(x[0,:]) + ', with a function value of ' + str(f[0]) + '\n'
		  
	def plot(self):
		path = self.path
		function = self.function
		#Plot the Rosenbrock function
		from mpl_toolkits.mplot3d import Axes3D
		from mpl_toolkits.mplot3d.art3d import Line3D
		from matplotlib import cm
		from matplotlib.colors import LogNorm
		import matplotlib.pyplot as plt
		import numpy as np
		fig = plt.figure()
		ax = Axes3D(fig, azim = -128, elev = 43)
		s = .1
		X = np.arange(-2, 2.5+s, s)
		Y = np.arange(-1, 3.5+s, s)
		X, Y = np.meshgrid(X, Y)
		Z = (1.-X)**2 + 100.*(Y-X*X)**2
		ax.plot_surface(X, Y, Z, rstride = 1, cstride = 1, norm = LogNorm(), cmap = cm.gray)
		plt.xlabel("x1")        
		plt.ylabel("x2")
		ax.plot(ravel(path[:,0]), ravel(path[:,1]), ravel(apply_along_axis(function, 1, path[:,:])), color = 'g', marker = '^')
        #       # Plot the simplex path
        #	    for i in range(0,path.shape[0],3):
        #		    x_path = path[i:i+3,0]
        #		    y_path = path[i:i+3,1]
        #		    f_path = apply_along_axis(function, 1, path[i:i+3,:])
        #			lstLines = [ [ ( x_path[j], y_path[j], f_path[j] ), ( x_path[k], y_path[k], f_path[k] )] for (j, k) in ((0,1),(0,2),(1,2))]
        #			ax.plot_wireframe(ravel(x_path), ravel(y_path), ravel(f_path), rstride = 1, cstride = 1)       
		plt.savefig("Rosenbrock function.svg")
		plt.show()