import numpy as np
import matplotlib.pyplot as plt

m = 0.2 ## mesh on the abscissa
s = 3 ## standard deviation of errors

def pdesign(X, d):
    """Generate a polynomial design matrix on X of order d."""
    V = X[:,np.newaxis]
    F = [V**k for k in range(d+1)]
    D = np.concatenate(F, axis=1)
    return D

def regfit(Y, D):
    """Regress Y on D using least squares."""
    U,S,Vt = np.linalg.svd(D,0)
    V = np.transpose(Vt)
    return np.dot(U, np.dot(np.transpose(U), Y))

X = np.arange(-2, 2, m, dtype=np.float64)

D1 = pdesign(X, 3)
D2 = pdesign(X, 13)

EY = X + X**3
Y1 = EY + np.random.normal(size=len(X))*s
Y2 = EY + np.random.normal(size=len(X))*s

Yhat1 = regfit(Y1, D1)
Yhat2 = regfit(Y1, D2)

plt.clf()
plt.figure(figsize=(8,3))
ax1 = plt.axes([0.06,0.1,0.4,0.8])
plt.title("Training set")
plt.plot(X, Y1, 'o')
plt.hold(True)
plt.plot(X, Yhat1, '-', color='green')
plt.plot(X, Yhat2, '-', color='orange')
ax1.set_ylim(-10, 10)
ax1.set_xticks([-2,-1,0,1,2])
ax2 = plt.axes([0.56,0.1,0.4,0.8])
plt.title("Test set")
plt.plot(X, Y2, 'o')
plt.plot(X, Yhat1, '-', color='green')
plt.plot(X, Yhat2, '-', color='orange')
ax2.set_xticks([-2,-1,0,1,2])
ax2.set_ylim(-10, 10)
plt.savefig("traintest.png")
plt.savefig("traintest.svg")

print ((Yhat1-Y1)**2).mean()
print ((Yhat2-Y1)**2).mean()

print ((Yhat1-Y2)**2).mean()
print ((Yhat2-Y2)**2).mean()