# Newton form of the interpolating polynomial, with divided differences.
#
#   p(x) = c_0 + c_1 (x - x_0) + c_2 (x - x_0)(x - x_1) + ... + c_n (x - x_0)...(x - x_{n-1}),   c_k = f[x_0, ..., x_k]
#
# The k-th coefficient only depends on the first k + 1 nodes, so adding a node adds one term and changes nothing before it.
#
# Needs: pip install numpy matplotlib      Run:  python newton_interp.py
# Edit the settings just below and run again.
import numpy as np
import matplotlib.pyplot as plt
from math import comb

# ------------------------------------------------------------------ settings: play with these
SOURCE = "function"                    # "function": y_k = f(x_k)   or   "data": your own points below, IN THE ORDER GIVEN
def f(x):                              # any numpy expression of x
    return np.exp(x) * np.sin(3 * x)
A, B = -1.0, 1.0
N = 6                                  # degree n: n + 1 nodes
NODES = "chebyshev"                    # "equispaced", "chebyshev" or "lobatto"
ORDER = "given"                        # order in which the nodes are used: "given", "reversed", "leja", "random"
DATA = [(0, 1), (2, 3), (4, 2.5), (6, 4.5), (8, 3), (10, 6)]
SAVE = True
# ------------------------------------------------------------------


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 leja(xs):
    """Leja order: start at the node of largest modulus, then add the node that maximizes prod |x - chosen|."""
    rest, out = list(xs), []
    out.append(rest.pop(int(np.argmax(np.abs(rest)))))
    while rest:
        score = [np.sum(np.log(np.abs(r - np.array(out)) + 1e-300)) for r in rest]
        out.append(rest.pop(int(np.argmax(score))))
    return np.array(out)


def divided_differences(xs, ys):
    """Return the whole triangle T[k][i] = f[x_i, ..., x_{i+k}]; the Newton coefficients are T[k][0]."""
    T = [np.array(ys, dtype=float)]
    for k in range(1, len(xs)):
        prev = T[-1]
        T.append((prev[1:] - prev[:-1]) / (xs[k:] - xs[:-k]))
    return T


def newton_eval(t, xs, c, k=None):
    """Nested (Horner-like) evaluation of the first k + 1 terms: c_0 + (t-x_0)(c_1 + (t-x_1)(c_2 + ...))."""
    k = len(c) - 1 if k is None else k
    s = np.full_like(t, c[k], dtype=float)
    for j in range(k - 1, -1, -1):
        s = s * (t - xs[j]) + c[j]
    return s


if SOURCE == "data":
    pts = np.array(DATA, dtype=float)               # NOT sorted: the order is the order of the nodes
    xs, ys = pts[:, 0], pts[:, 1]
    A, B = xs.min(), xs.max()
else:
    xs = make_nodes(NODES, N, A, B)
    if ORDER == "reversed":
        xs = xs[::-1]
    elif ORDER == "leja":
        xs = leja(xs)
    elif ORDER == "random":
        xs = np.random.default_rng(1).permutation(xs)
    ys = f(xs)
n = len(xs) - 1
T = divided_differences(xs, ys)
c = np.array([T[k][0] for k in range(n + 1)])

# 1. the divided-difference table
print("nodes (in this order):", np.array2string(xs, precision=4))
print("\ndivided differences  (column k holds f[x_i, ..., x_{i+k}]; the first row is the list of coefficients c_k)")
for i in range(n + 1):
    print("  " + "  ".join(f"{T[k][i]:11.5f}" if i <= n - k else " " * 11 for k in range(n + 1)))
print("\nNewton coefficients c_k =", np.array2string(c, precision=5))

# 2. the same polynomial for every order of the nodes, but not the same coefficients
t = np.linspace(A, B, 1000)
p = newton_eval(t, xs, c)
order2 = np.argsort(xs)
c2 = np.array([divided_differences(xs[order2], ys[order2])[k][0] for k in range(n + 1)])
p2 = newton_eval(t, xs[order2], c2)
print("\nsame polynomial with the nodes sorted:  max |p - p_sorted| =", np.abs(p - p2).max())

# 3. the error: f(x) - p(x) = f[x_0, ..., x_n, x] * omega(x), and for smooth f this is f^(n+1)(xi)/(n+1)! * omega(x)
if SOURCE == "function":
    print("max |f - p| on the grid =", np.abs(f(t) - p).max())

# 4. equispaced nodes: divided differences are forward differences divided by k! h^k  (Gregory-Newton)
if NODES == "equispaced" and SOURCE == "function" and ORDER == "given":
    h = xs[1] - xs[0]
    D = [ys.copy()]
    for k in range(1, n + 1):
        D.append(np.diff(D[-1]))
    s = (t - xs[0]) / h
    pg = sum(D[k][0] * np.array([np.prod([(si - j) for j in range(k)]) / np.prod(range(1, k + 1)) for si in s]) for k in range(n + 1))
    print("Gregory-Newton with Delta^k y_0 reproduces p:  max difference =", np.abs(p - pg).max())

# 5. the pictures: p_0, p_1, ..., p_n  (each new node adds one term)
fig, ax = plt.subplots(1, 2, figsize=(13, 4.8))
cmap = plt.get_cmap("viridis")
for k in range(n + 1):
    ax[0].plot(t, newton_eval(t, xs, c, k), color=cmap(k / max(n, 1)), lw=1.1 + (k == n), label=f"$p_{k}$" if n <= 7 else None)
if SOURCE == "function":
    ax[0].plot(t, f(t), "--", color="gray", lw=1.3, label="$f$")
ax[0].plot(xs, ys, "o", color="crimson", zorder=5)
for k, xk in enumerate(xs):
    ax[0].annotate(str(k), (xk, ys[k]), textcoords="offset points", xytext=(4, 6), fontsize=8)
ax[0].set_title("partial sums $p_k$: node $x_k$ adds the term $c_k(x-x_0)\\cdots(x-x_{k-1})$"); ax[0].legend(fontsize=8, ncol=2); ax[0].grid(alpha=.3)
for k in range(n + 1):
    nb = np.prod([t - xs[j] for j in range(k)], axis=0) if k else np.ones_like(t)
    ax[1].plot(t, nb, color=cmap(k / max(n, 1)), lw=1.2)
ax[1].set_title("the Newton basis $N_k(x)=\\prod_{j<k}(x-x_j)$: $N_k$ vanishes at $x_0,\\dots,x_{k-1}$")
lo, hi = np.percentile(np.concatenate([np.prod([t - xs[j] for j in range(k)], axis=0) if k else np.ones_like(t) for k in range(n + 1)]), [1, 99])
ax[1].set_ylim(lo * 1.2, hi * 1.2); ax[1].grid(alpha=.3)
plt.tight_layout()
if SAVE:
    plt.savefig("newton_pieces.png", dpi=140)

# 6. clustering the nodes: the divided differences become derivatives, f[x_0..x_k] -> f^(k)(x_0)/k!  (Taylor, 1715)
if SOURCE == "function":
    print("\nclustering the nodes around x_0 = 0: the first coefficients settle, then rounding errors (cancellation) take over")
    u = np.arange(7) / 6.0
    for spread in (1.0, 1e-1, 1e-2, 1e-3, 1e-4, 1e-6, 1e-8):
        xc = spread * u
        cc = [divided_differences(xc, f(xc))[k][0] for k in range(5)]
        print(f"  spread {spread:7.0e}:  " + "  ".join(f"{v:12.6f}" for v in cc))
plt.show()