import matplotlib.pyplot as plt
from matplotlib.patches import Circle
import numpy as np
from scipy.special import erfinv

# Set up the grid and other parameters
M, N = 50, 100
A, B = -2.2, 7
C, D = -4, 4
a, b, c = 1., 3, 1
p, q, r, s = 1, 7.5, 1.3, 0.16
p = s*p
q = s*q
r = s*r

X = np.linspace(A, B, N, dtype=np.float)
XX = np.linspace(A, B, N+1, dtype=np.float)
YY = p + q * XX + r * (XX**2)

Xr = 7.0 * np.arange(N, dtype=np.float)
Yr = Xr.copy()
for i in range(N):
    rd = np.random.rand()
    Xr[i] = XX[i] * rd + XX[i+1] * (1 - rd)
    Yr[i] = p + q * XX[i] + r * (XX[i]**2) + 0.4 * erfinv(2*np.random.rand() - 1)
    
# least squares fitting
Mat = np.array([(0*Xr+1), Xr, Xr**2]).T
V = np.matmul(Mat.T, Yr)
V = np.linalg.solve(np.matmul(Mat.T, Mat), V)
pe = V[0]; qe = V[1]; re = V[2];

fig, ax = plt.subplots(1, 1, figsize=(2,2), edgecolor='k')

ax.plot(X, pe + qe * X + re * (X**2), 'b', lw=2)
myrad = 0.05
for i in range(len(Xr)):
    sample = Circle((Xr[i], Yr[i]), myrad, color='r')
    ax.add_patch(sample)

ax.grid(True)
ax.set_xticks([-2, -1, 0, 1, 2])
ax.set_yticks([-1, 0, 1, 2, 3])
ax.axis('square')
ax.axis([-2, 2, -1.7, 3.5])
fig.savefig('least squares.svg', transparent=True, bbox_inches='tight')