text: #Code written by Zachary Montgomery.
  #Altered by Cory Goates to handle "NaN" values.
  #Altered by Cory Goates to allow the user to specify if some power coefficients
  #should be set to zero.
  
  import numpy as np
  from scipy.optimize import least_squares
  
  def poly_fit(n,xx,yy,forcezero=None):
  	good = np.where(~(np.isnan(xx)) & ~(np.isnan(yy)))
  	x = np.copy(xx[good])
  	y = np.copy(yy[good])
   
  	if forcezero is None:
  		# initialize variables
  		A = np.zeros((n,n))
  		b = np.zeros(n)
  		powers = range(n)
  	else:
  		size = n - len(forcezero)
  		A = np.zeros((size,size))
  		b = np.zeros(size)
  		powers = np.delete(range(n), forcezero)
    
  	for i,e in enumerate(powers):
  		for j,f in enumerate(powers):
  			A[i,j] = sum( x ** float(e + f) )
  		b[i] = sum( y * x ** float(e) )
    
  	a = np.linalg.solve(A,b)
   
  	if forcezero is not None:
  		for zero in forcezero:
  			a = np.insert(a,zero,0.)
  	r = r2(a,x,y)
  	return a, r
  
  def poly_func(a,x):
  	f = 0.
  	n = len(a)
  	for i in range(n):
  		f += a[i] * x ** float(i)
  	#print(f)
  	return f
  
  def r2(aa,xx,yy):
  	a = np.copy(aa)
  	x = np.copy(xx)
  	y = np.copy(yy)
  	n = len(y)
  	y_ = sum(y) / float(n)
  	SSt = sum((y-y_)**2.)
  	f = np.zeros(n)
  	for i in range(n):
  		f[i] = poly_func(a,x[i])
  	SSr = sum((y-f)**2.)
  	return 1. - SSr / SSt
  
  def force_good_poly_fit(n,xx,yy,tol,sym=False):
  	x = np.copy(xx)
  	y = np.copy(yy)
  	
  	a, r = poly_fit(n,x,y,sym=sym)
  	x_bad = []
  	y_bad = []
  	print('R2 is {:7.5f}'.format(r))
  	while r < tol:
  		
  		e_loc = 0
  		l = y.size
  		e_max = 0.
  		for i in range(l):
  			e = abs(poly_func(n,a,x[i]) - y[i])
  			if e > e_max:
  				e_max = e
  				e_loc = i
  		x_bad.append(x[e_loc])
  		y_bad.append(y[e_loc])
  		x = np.delete(x,e_loc)
  		y = np.delete(y,e_loc)
  		
  		a, r = poly_fit(n,x,y)
  		print('R2 is {:7.5f}'.format(r))
  	
  	mx = max(y)
  	
  	return a, x, y, x_bad, y_bad, r, mx
  
  def constrainedFit(n,xx,yy,forcezero=None):
  # Fits a polynomial, ensuring the highest-order coefficient is negative so that end behavior is as expected.
  	def fun(a):
  		res = poly_func(a,xx)-yy
  		res[np.isnan(res)] = 0
  		return res
  
  	x0 = -0.01*np.ones(n).astype(float)
  	upperBounds = np.full(n,np.inf)
  	upperBounds[-1] = 0 # Costrain last coef to be negative
  	result = least_squares(fun,x0,bounds=(-np.inf,upperBounds))
  	return result.x, result.cost
  