#(c) Rory Clune 05/16/2011

from numpy import *
from matplotlib import*

class NelderMeadConstrained:
	'''Implementation of Constrained Nelder Mead Algorithm'''			
	def __init__(self, ndim, structure, start_pt, UB = None, LB = None, tol = 1e-15, maxsteps = 1000):
		self.ndim = ndim
		self.function = structure.Objective
		self.constraint = structure.Constraint
		self.tol = tol
		self.maxsteps = maxsteps
		# Apply the bounds, if any are given
		self.UB = UB
		self.LB = LB
		if (UB != None) or (LB != None): 
			self.bounded = True
			start_pt = maximum(minimum(start_pt, self.UB), self.LB)
		else: self.bounded = False
		self.x0 = start_pt		
		self.scale = ones_like(start_pt)
		self.penalty = 0*ones(self.constraint(start_pt).shape[0])
		self.s = 0.001	
		
	def Run(self):				
		if self.bounded: print 'Bounded Nelder-Mead Algorithm is running.. Please wait'  
		else: print 'Unbounded Nelder-Mead Algorithm is running.. Please wait'
		
		#----------------------
		# 0. Data Setup
		#----------------------		
		function = self.function
		constraint = self.constraint
		penalty = self.penalty; s = self.s		
		scale = self.scale
		ndim = self.ndim
		npts = ndim + 1 # a simplex
		nsteps = 0
		
		#----------------------
		# 1. Define the starting simplex and evaluate the function at its vertices
		#----------------------				
		x = zeros((npts, ndim))		
		if self.x0 == None:
			x[0,:] = 0 # Arbitrary starting point
		else:
			x[0,:] = self.x0
		for i in range(1,npts,1):
			x[i,:] = x[i-1,:] + arange(1.,ndim+1,1)*i/100
			
		print x.shape[1]		
		# Make sure all the new points stay within the allowed bounds
		if self.bounded:
			for i in range(1,npts,1):
				x[i,:] = maximum(minimum(x[i,:], self.UB), self.LB)
		f = apply_along_axis(function,1,x*scale)	 
		
		#----------------------
		# 1. (b) Add penalties for violated constraints and increase multipliers
		#----------------------					
		c = apply_along_axis(constraint,1,x*scale)
		f += (c*penalty).sum(axis = 1)		
		penalty += s*(c.sum(axis = 0))		
		
		# Sort the simplex and the vector of function values
		index = argsort(f, axis = 0) # best - middle - worst
		x = x[index,:] 
		f = f[index]
		nsteps += npts - 1
		# set up storage for simplex path plotting
		self.path = copy(x[0,:])
		
		self.obj_history = f[0]
		self.con_history = c[0]
		
		converged = False
		while (converged==False): #Do the following until convergence
			#----------------------			
			# 2. Find the center of the face of the simplex (except the worst point)
			#----------------------		
			x_mean = mean(x[0:(npts-1),:], axis = 0)    
			
			#----------------------			
			# 3. Reflect the worst point across the face
			#----------------------			
			xr = 2*x_mean-x[(npts-1),:]
			if self.bounded:
				xr = maximum(minimum(xr, self.UB), self.LB)
			nsteps += 1
			fxr = function(xr*scale)
			c = constraint(xr*scale)			
			fxr += sum(penalty*c)
			penalty += s*c
			
			#----------------------			
			# 4. Check if the reflected point is the new best, and go further if so 
			#----------------------			     
			if fxr <= f[0]: # is the new point the new best?
				x2r = 3*x_mean-2*x[(npts-1),:] # grow further in that direction
				if self.bounded:
					x2r = maximum(minimum(xr, self.UB), self.LB)				
				fx2r = function(x2r*scale)				
				c = constraint(x2r*scale)			
				fx2r += sum(penalty*c)
				penalty += s*c						
				nsteps += 1
				if fx2r < fxr: # did the growing improve?
					x[(npts-1),:] = x2r # accept the grown point and go to the next iteration
					f[(npts-1)] = fx2r	
					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)] = fxr            
					print 'accepted reflected point (it was the best)'
			
			#----------------------			
			# 4. Check if the reflected point is still the worst, and try reflecting & shrinking, just shrinking, and then shrinking everything 
			#----------------------			
			elif fxr >= f[(npts-2)]: 
				xs = 1.5*x_mean - 0.5*x[(npts-1),:] # reflect + shrink
				if self.bounded:
					xs = maximum(minimum(xs, self.UB), self.LB)
				fxs = function(xs*scale)				
				c = constraint(xs*scale)			
				fxs += sum(penalty*c)
				penalty += s*c
				nsteps += 1
				if fxs < f[(npts-2)]: # better than the second worst?
					x[(npts-1),:] = xs # accept and move on
					f[(npts-1)] = fxs
					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 self.bounded:
						xs = maximum(minimum(xs, self.UB), self.LB)						
					fxs = function(xs*scale)				
					c = constraint(xs*scale)			
					fxs += sum(penalty*c)
					penalty += s*c
					nsteps += 1
					if fxs > f[(npts-2)]:
						x[1:npts,:] = 0.5*(x[1:npts,:] + x[0,:]) # shrink everything (except the best, 0) towards the best
						if self.bounded:
							for i in range(1,npts,1):
								x[i,:] = maximum(minimum(x[i,:], self.UB), self.LB)
						f = apply_along_axis(function, 1, x*scale)
						c = apply_along_axis(constraint, 1, x*scale)
						f += (c*penalty).sum(axis = 1)		
						penalty += s*(c.sum(axis = 0))
						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)] = fxr
				print 'accepted reflected point (though it was not a new best)'   
			# Sort the simplex and the vector of function values
			index = argsort(f, axis = 0) # best - middle - worst
			x = x[index,:]
			f = f[index]			
			# Record the history of the simplex:			
			self.obj_history = append(self.obj_history, f[0])
			# Check for convergence or maximum iterations
			if (fabs(f[0]- f[(npts-1)]) < self.tol) or (nsteps > self.maxsteps):
				converged = True						
		print '\nconverged at: ' + str(scale*x[0,:]) + ', with a function value of ' + str(f[0]) + '\n' + str(nsteps) + ' points generated' + '\n'
		self.optimum = x[0,:]	



class NelderMead:
	'''Implementation of Nelder Mead Algorithm'''			
	def __init__(self, ndim, structure, start_pt, UB = None, LB = None, tol = 1e-15, maxsteps = 1000):
		self.ndim = ndim
		self.function = structure.Objective
		self.tol = tol
		self.maxsteps = maxsteps
		# Apply the bounds, if any are given
		self.UB = UB
		self.LB = LB
		if (UB != None) or (LB != None): 
			self.bounded = True
			start_pt = maximum(minimum(start_pt, self.UB), self.LB)
		else: self.bounded = False
		self.x0 = start_pt
		self.scale = ones_like(start_pt)
		
	def Run(self):				
		if self.bounded: print 'Bounded Nelder-Mead Algorithm is running.. Please wait'  
		else: print 'Unbounded Nelder-Mead Algorithm is running.. Please wait'
		
		#----------------------
		# 0. Data Setup
		#----------------------		
		function = self.function		
		ndim = self.ndim
		npts = ndim + 1 # a simplex
		nsteps = 0
		scale = self.scale
		#----------------------
		# 1. Define the starting simplex and evaluate the function at its vertices
		#----------------------				
		x = zeros((npts, ndim))		
		if self.x0 == None:
			x[0,:] = 0 # Arbitrary starting point
		else:
			x[0,:] = self.x0
		for i in range(1,npts,1):			
			x[i,:] = x[i-1,:] + arange(1.,ndim+1,1)*i/100
		# Make sure all the new points stay within the allowed bounds
		if self.bounded:
			for i in range(1,npts,1):
				x[i,:] = maximum(minimum(x[i,:], self.UB), self.LB)
		f = apply_along_axis(function,1,x*scale)
		# Sort the simplex and the vector of function values
		index = argsort(f, axis = 0) # best - middle - worst
		x = x[index,:] 
		f = f[index]
		nsteps += npts - 1
		# set up storage for simplex path plotting
		self.path = copy(x[0,:])
		self.obj_history = f[0]
		
		converged = False
		while (converged==False): #Do the following until convergence
			#----------------------			
			# 2. Find the center of the face of the simplex (except the worst point)
			#----------------------		
			x_mean = mean(x[0:(npts-1),:], axis = 0)    
			
			#----------------------			
			# 3. Reflect the worst point across the face
			#----------------------			
			xr = 2*x_mean-x[(npts-1),:]
			if self.bounded:
				xr = maximum(minimum(xr, self.UB), self.LB)
			nsteps += 1
			fxr = function(xr*scale)
			#----------------------			
			# 4. Check if the reflected point is the new best, and go further if so 
			#----------------------			     
			if fxr <= f[0]: # is the new point the new best?
				x2r = 3*x_mean-2*x[(npts-1),:] # grow further in that direction
				if self.bounded:
					x2r = maximum(minimum(xr, self.UB), self.LB)			
				nsteps += 1
				fx2r = function(x2r*scale)
				if fx2r < fxr: # did the growing improve?
					x[(npts-1),:] = x2r # accept the grown point and go to the next iteration
					f[(npts-1)] = fx2r	
#					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)] = fxr            
#					print 'accepted reflected point (it was the best)'
			
			#----------------------			
			# 4. Check if the reflected point is still the worst, and try reflecting & shrinking, just shrinking, and then shrinking everything 
			#----------------------			
			elif fxr >= f[(npts-2)]: 
				xs = 1.5*x_mean - 0.5*x[(npts-1),:] # reflect + shrink
				if self.bounded:
					xs = maximum(minimum(xs, self.UB), self.LB)
#					print 'shrunk and reflected point only'
				nsteps += 1
				fxs = function(xs*scale)
				if fxs < f[(npts-2)]: # better than the second worst?
					x[(npts-1),:] = xs # accept and move on
					f[(npts-1)] = fxs
#					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 self.bounded:
						xs = maximum(minimum(xs, self.UB), self.LB)
#						print 'shrunk point'
					nsteps += 1
					fxs = function(xs*scale)				
					if fxs > f[(npts-2)]:
						x[1:npts,:] = 0.5*(x[1:npts,:] + x[0,:]) # shrink everything (except the best, 0) towards the best
						if self.bounded:
							for i in range(1,npts,1):
								x[i,:] = maximum(minimum(x[i,:], self.UB), self.LB)
						f = apply_along_axis(function, 1, x*scale)
#						print 'shrank everything towards the best'
					else:
						x[(npts-1),:] = xs # accept and move on
						f[(npts-1)] = fxs
#						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)] = fxr
#				print 'accepted reflected point (though it was not a new best)'   
			# Sort the simplex and the vector of function values
			index = argsort(f, axis = 0) # best - middle - worst
			x = x[index,:]
			f = f[index]			
			
			# Record the history of the simplex:			
			self.obj_history = append(self.obj_history, f[0])
						
			if (fabs(f[0]- f[(npts-1)]) < self.tol) or (nsteps > self.maxsteps):
				converged = True
		print '\nconverged at: ' + str(scale*x[0,:]) + ', with a function value of ' + str(f[0]) + '\n' + str(nsteps) + ' points generated' + '\n'
		self.optimum = x[0,:]