"""
MATH3476 Example 1.7
Newton divided difference for u=sqrt(x)
"""
import numpy as np
import matplotlib.pyplot as plt

n = 6
a = 2
b = 2.4

# n+1 eqaully-spaced nodes
xn_unif = np.linspace(a,b,n+1)

# n+1 Chebsshev nodes
mid = (a+b)/2.
r = (b-a)/2.
theta = np.pi/n*np.linspace(0,n,n+1)
xn_cheb = mid - r*np.cos(theta)

# set up the x-axis
n_pt = 500
x = np.linspace(a,b,n_pt)

Psi_unif = np.ones(n_pt)
Psi_cheb = np.ones(n_pt)

for i in range(n+1):
    Psi_unif *= (x - xn_unif[i])
    Psi_cheb *= (x - xn_cheb[i])

# plotting
plt.figure(1,figsize=(4,4)); plt.clf()
plt.plot(x,Psi_unif,'b')
plt.plot(x,Psi_cheb,'r--')
plt.xlim((a,b))
plt.xlabel('$x$')
plt.axhline(linewidth=1,color='k')

plt.savefig('Psi.pdf',dpi=600,bbox_inches='tight')
