# Lagrange interpolation: the polynomial through points, built from special tables.
#
# Lagrange's insight: for the SAME nodes x_0..x_n, the polynomial l_j of degree <= n through the special
# table (x_k, delta_jk)  -- value 1 at x_j and 0 at every other node -- is  l_j(x) = prod_{k != j} (x - x_k)/(x_j - x_k).
# Any table y is the sum of y_j times these special tables, and interpolation is linear, so
#       p(x) = y_0 l_0(x) + y_1 l_1(x) + ... + y_n l_n(x).
#
# Needs: pip install numpy matplotlib      Run:  python lagrange_interp.py
# Edit the settings just below and run again.
import numpy as np
import matplotlib.pyplot as plt

# ------------------------------------------------------------------ settings: play with these
SOURCE = "function"                    # "function": y_k = f(x_k)   or   "data": your own points below
def f(x):                              # any numpy expression of x, for example np.exp(x), np.abs(x), np.sin(2*np.pi*x)
    return 1 / (1 + 25 * x**2)         # Runge's function
A, B = -1.0, 1.0                       # the interval
N = 8                                  # degree n: there are n + 1 nodes
NODES = "chebyshev"                    # "equispaced", "chebyshev" (zeros of T_{n+1}) or "lobatto" (extrema of T_n)
DATA = [(0, 1), (2, 3), (4, 2.5), (6, 4.5), (8, 3), (10, 6)]   # (x, y) pairs, used when SOURCE == "data"
SAVE = True                            # also write the pictures as lagrange_*.png
# ------------------------------------------------------------------


def make_nodes(kind, n, a, b):
    m, h, i = (a + b) / 2, (b - a) / 2, np.arange(n + 1)
    if kind == "chebyshev":
        return m - h * np.cos((2 * i + 1) * np.pi / (2 * n + 2))
    if kind == "lobatto" and n > 0:
        return m - h * np.cos(i * np.pi / n)
    return np.linspace(a, b, n + 1)


def lagrange_basis(xs, j, t):
    """l_j(t) straight from the definition: the product of the n factors (t - x_k)/(x_j - x_k), k != j."""
    out = np.ones_like(t, dtype=float)
    for k, xk in enumerate(xs):
        if k != j:
            out *= (t - xk) / (xs[j] - xk)
    return out


def interpolate(xs, ys, t):
    """p(t) = sum_j y_j l_j(t): the data table written as a sum of special tables."""
    return sum(yj * lagrange_basis(xs, j, t) for j, yj in enumerate(ys))


def barycentric(xs, ys, t):
    """Same polynomial, second barycentric formula: O(n) per point and numerically stable."""
    w = np.array([1 / np.prod(xj - np.delete(xs, j)) for j, xj in enumerate(xs)])
    out = np.empty_like(t, dtype=float)
    for i, ti in enumerate(t):
        d = ti - xs
        if np.any(d == 0):
            out[i] = ys[np.argmin(np.abs(d))]
        else:
            q = w / d
            out[i] = np.dot(q, ys) / q.sum()
    return out


if SOURCE == "data":
    pts = np.array(sorted(DATA), dtype=float)
    xs, ys = pts[:, 0], pts[:, 1]
    A, B = xs[0], xs[-1]
else:
    xs = make_nodes(NODES, N, A, B)
    ys = f(xs)
n = len(xs) - 1
t = np.linspace(A, B, 1200)

# 1. the special tables: l_j(x_k) must be the identity matrix
L = np.array([[lagrange_basis(xs, j, np.array([xk]))[0] for k, xk in enumerate(xs)] for j in range(n + 1)])
print("max |l_j(x_k) - delta_jk| =", np.abs(L - np.eye(n + 1)).max())
print("max |p(x_k) - y_k|        =", np.abs(interpolate(xs, ys, xs) - ys).max())
print("max |Lagrange - barycentric| on the grid =", np.abs(interpolate(xs, ys, t) - barycentric(xs, ys, t)).max())

# 2. the Lebesgue function sum_j |l_j(x)| and its maximum: the condition number of interpolation
leb = sum(np.abs(lagrange_basis(xs, j, t)) for j in range(n + 1))
print("Lebesgue constant Lambda_n =", leb.max(), " (|p(x)| <= Lambda_n * max|y_j|)")
print("Vandermonde condition number =", np.linalg.cond(np.vander(xs, increasing=True)))

# 3. the pictures
p = interpolate(xs, ys, t)
fig, ax = plt.subplots(1, 2, figsize=(13, 4.8))
for j in range(n + 1):
    ax[0].plot(t, ys[j] * lagrange_basis(xs, j, t), lw=1, alpha=.8, label=f"$y_{j}\\,\\ell_{j}$" if n <= 6 else None)
ax[0].plot(t, p, "k", lw=2.2, label="$p=\\sum y_j\\ell_j$")
if SOURCE == "function":
    ax[0].plot(t, f(t), "--", color="gray", lw=1.2, label="$f$")
ax[0].plot(xs, ys, "o", color="crimson", zorder=5)
lo, hi = ys.min(), ys.max()
ax[0].set_ylim(lo - 1.5 * (hi - lo), hi + 1.5 * (hi - lo))
ax[0].set_title("your table = sum of $y_j$ times the special tables"); ax[0].legend(fontsize=8, ncol=2); ax[0].grid(alpha=.3)
for j in range(n + 1):
    ax[1].plot(t, lagrange_basis(xs, j, t), lw=1.3)
    ax[1].plot(xs, (np.arange(n + 1) == j).astype(float), "o", ms=3, color="k")
ax[1].set_ylim(-1.2, 1.6)
ax[1].set_title("$\\ell_j$ interpolates the table $(x_k,\\delta_{jk})$: a 1 at $x_j$, 0 elsewhere"); ax[1].grid(alpha=.3)
plt.tight_layout()
if SAVE:
    plt.savefig("lagrange_pieces.png", dpi=140)

if SOURCE == "function":
    fig2, ax2 = plt.subplots(figsize=(7, 4))
    ax2.semilogy(t, np.abs(f(t) - p) + 1e-18, label="$|f-p_n|$")
    ax2.semilogy(t, leb * 1e-3, "--", label="Lebesgue function / 1000")
    ax2.set_title(f"error, n = {n}, {NODES} nodes"); ax2.legend(); ax2.grid(alpha=.3)
    plt.tight_layout()
    if SAVE:
        plt.savefig("lagrange_error.png", dpi=140)
plt.show()