from __future__ import print_function
from matplotlib.pyplot import *
from numpy import *

# Define the functions
def f(x): return sin(x)
def g(x): return 1./(1.+x*x)
def h(x): return x**(-12) - x**(-6)

# Interpolation
from scipy.interpolate import lagrange

# Create 5 equidistant points between 0 and 2*PI
x5 = linspace(0.0, 2.0 * pi, 5)
# Generate the interpolating polynomial
fp5 = lagrange(x5, f(x5))
# fp5 can now be used like a function
print(fp5(pi/2))

# Plot the function
xs = linspace(0.0, 2.0 * pi, 1000)
plot(xs, f(xs), label="$f(x)$")
# Plot the interpolating polynomial
plot(xs, fp5(xs), label="$P_5(x)$")
plot(x5, f(x5), 'o')
legend()

# A function to simplify plotting
def plot_interp(f, min, max, interp, ns=[5,10,15], nodes=linspace):
    """Create a plot of the function f between min and max plus plots 
    of the interpolating functions interp with ns different interpolation 
    points. 
    nodes should be a function that generates n supporting points in the 
    interval [min, max]. 
    """
    xp = linspace(min, max, 1000)
    plot(xp, f(xp), label="$f(x)$")

    for n in ns:
        x = nodes(min, max, n)
        ip = interp(x, f(x))
        plot(xp, ip(xp), label="$P_{{{0}}}(x)$".format(n))
    legend(loc="best")

# Use the function to generate plots of the interpolating function
figure()
plot_interp(f, 0.0, 2*pi, lagrange)
title("Sine")

figure()
plot_interp(g, -5., 5., lagrange)
title("Runge")

figure()
plot_interp(h, 1., 5., lagrange)
title("LJ")