import numpy as np
from numpy.linalg import inv
import matplotlib.pyplot as plt

# generate "noisy" data
n = 20 # number of points
x = 5*np.random.randn(n)
y = np.random.randn(1)*x + np.random.randn(1) # first-order polynomial
y = y + np.random.randn(n) # add noise

# your least-squares estimation code here
X = np.stack( (x, np.ones((n))), axis=1)
u = inv(np.transpose(X)@X)@np.transpose(X)@y

# your weighted least-squares estimation code here
w  = 1 + np.sqrt(np.square(x) + np.square(y)) # weighting
w  = np.square(w)
W  = np.diag(w) # diagonal weighting matrix
W2 = W@W # square of diagonal weighting matrix
uw = inv(np.transpose(X)@W2@X)@np.transpose(X)@W2@y

# plot data and least-squares fit
x0 = np.arange(-20,20,0.1)
y1 = u[0]*x0 + u[1]
y2 = uw[0]*x0 + uw[1]

ax = plt.subplot(1,1,1)
plt.plot(x,y,'bo', label='data') # original data
plt.plot( x0, y1, 'r-', label='LS') # least-squares fit
plt.plot( x0, y2, 'g-', label='WLS' ) # weighted least-squares fit
ax.legend()
plt.show()