Python: Least Squares Fit

# Least squares fit
import numpy as np
import matplotlib.pyplot as plt

def least_squares(x, y, k):

	if k >= 7: print('Warning: Higher order polynomials \
				    \nhave a tendency to become unbounded \
				    \nat beginning or end of user-supplied \
				    \ndata ranges.\n')
	# Least squares fit
	n = len(x)

	x_nk = np.array(
		   np.mat(
		   np.zeros((n, k+1))
		   )
		   )

	#for i in xrange(0, n):

	for j in xrange(0, k+1):

		x_nk[:, j] = x**(j)
	
	# a = inv(x'x)*x'y
	a = np.dot(
			np.dot(
				np.linalg.inv(
					np.dot(
						np.transpose(x_nk), x_nk
				 	       )
		 	                 ),
				np.transpose(x_nk)
		       	  ),
			np.transpose(y)
		      )

	print('Coefficients using least-squares method:')
	print(np.poly1d(np.flipud(a)))

	xf = np.linspace(0, np.max(x))
	yls = sum([a[i]*xf**i for i in range(0, k+1)])
	
	# Using numpy polyfit() to compare
	pfit = np.poly1d(np.polyfit(x, y, k))

	print('Coefficients using built-in polyfit() method:')
	print(np.poly1d(pfit))

	plt.plot(x,  y,  'ob',  
	         xf, yls, '-.r', 
	         xf, pfit(xf), '+k')

	plt.grid(True)
	plt.legend(['Data Points',
		        'Least square fit',
		        'Using polyfit()'])
	plt.show()

# User-supplied data
x = np.array([5, 10, 15, 20, 28, 30, 35, 40, 44, 50])
y = np.array([17, 24, 31, 33, 37, 37, 40, 40, 42, 41])

# Order of least squares fit
k = 3

least_squares(x, y, k)