# Finite differences, the Newton-Gregory polynomial and finite-difference derivatives.
#
#   forward:   p(x_0 + s h) = sum_k binom(s, k)         Delta^k y_0        (s = (x - x_0) / h)
#   backward:  p(x_n + s h) = sum_k binom(s + k - 1, k) nabla^k y_n        (s = (x - x_n) / h)
#   derivative:  h D = ln(1 + Delta) = Delta - Delta^2/2 + Delta^3/3 - ...
#
# Needs: pip install numpy matplotlib      Run:  python finite_differences.py
# Edit the settings just below and run again.
from fractions import Fraction
from math import comb, factorial
import numpy as np
import matplotlib.pyplot as plt

# ------------------------------------------------------------------ settings: play with these
SOURCE = "function"                    # "function": y_i = f(x_i)   or   "data": your own values below
def f(x):                              # any numpy expression of x
    return np.exp(x)
X0, H, N = -1.0, 0.25, 8               # nodes x_i = X0 + i H, i = 0..N (the data list sets N when SOURCE = "data")
DATA = [1, 2, 0, 5, 4.5, 8]
K = 3                                  # degree used in the plot of the polynomial
EVAL_AT = -0.6                         # point where both formulas are evaluated
NOISE = 1e-3                           # size of the disturbance in the noise experiment (alternating signs: the worst case)
DERIVATIVE_ORDER = 1                   # 1 or 2: the derivative in the step experiment
SAVE = True
# ------------------------------------------------------------------


def difference_table(y):
    """T[k][i] = Delta^k y_i; T[k] has N + 1 - k entries."""
    T = [np.array(y, dtype=float)]
    for _ in range(len(y) - 1):
        T.append(np.diff(T[-1]))
    return T


def binom_s(s, k):
    """binom(s, k) for real s."""
    b = 1.0
    for j in range(k):
        b *= (s - j) / (j + 1)
    return b


def forward_eval(x, x0, h, T, k):
    s = (x - x0) / h
    return sum(binom_s(s, j) * T[j][0] for j in range(k + 1))


def backward_eval(x, xn, h, T, k):
    n = len(T) - 1
    s = (x - xn) / h
    return sum(binom_s(s + j - 1, j) * T[j][n - j] for j in range(k + 1))


def lagrange_check(x, xs, ys):
    """The same polynomial from the Lagrange form, to compare (n + 1 nodes, degree n)."""
    tot = 0.0
    for j in range(len(xs)):
        term = ys[j]
        for m in range(len(xs)):
            if m != j:
                term *= (x - xs[m]) / (xs[j] - xs[m])
        tot += term
    return tot


def forward_weights(m, K):
    """Weights w_j with  D^m f(x_0) ~ h^-m sum_j w_j y_j  from differentiating the forward polynomial, in exact fractions."""
    def binom_poly(k):                  # coefficients of binom(s, k) in powers of s
        c = [Fraction(1)]
        for j in range(k):
            nc = [Fraction(0)] * (len(c) + 1)
            for i, ci in enumerate(c):
                nc[i + 1] += ci
                nc[i] -= j * ci
            c = [q / (j + 1) for q in nc]
        return c
    w = [Fraction(0)] * (K + 1)
    for k in range(m, K + 1):
        a_k = binom_poly(k)[m] * factorial(m)         # m-th derivative at s = 0
        for j in range(k + 1):
            w[j] += a_k * (-1) ** (k - j) * comb(k, j)
    return w


def central_weights(m, r):
    """Weights on the offsets -r..r from  sum_j w_j o_j^q = m! delta_{qm}."""
    o = np.arange(-r, r + 1)
    A = np.vander(o, increasing=True).T.astype(float)
    b = np.zeros(2 * r + 1)
    b[m] = factorial(m)
    return np.linalg.solve(A, b)


def main():
    if SOURCE == "data":
        y = np.array(DATA, dtype=float)
        n = len(y) - 1
        x = X0 + H * np.arange(n + 1)
    else:
        n = N
        x = X0 + H * np.arange(n + 1)
        y = f(x)
    T = difference_table(y)

    print("difference table (column k = Delta^k y_i):")
    for k, col in enumerate(T):
        print(f"  k={k}:", np.array2string(col, precision=6, suppress_small=True, max_line_width=110))
    scale = np.max(np.abs(y))
    for k in range(1, n + 1):
        if np.max(np.abs(T[k])) <= 64 * 2.3e-16 * scale * 2 ** k:
            print(f"column {k} vanishes to rounding error: the data are those of a polynomial of degree {k - 1}")
            break

    # ---- the two formulas against the interpolating polynomial
    pf, pb = forward_eval(EVAL_AT, X0, H, T, n), backward_eval(EVAL_AT, x[n], H, T, n)
    pl = lagrange_check(EVAL_AT, x, y)
    print(f"\nat x = {EVAL_AT}: forward {pf:.12g}   backward {pb:.12g}   Lagrange {pl:.12g}")
    for k in range(n + 1):
        print(f"  k={k}:  forward {forward_eval(EVAL_AT, X0, H, T, k):.10g}   backward {backward_eval(EVAL_AT, x[n], H, T, k):.10g}")

    # ---- noise: alternating signs give the worst case 2^k
    yn = y + NOISE * (-1.0) ** np.arange(n + 1)
    Tn = difference_table(yn)
    print(f"\nlargest entry of each column, clean / disturbed (noise {NOISE}):")
    for k in range(n + 1):
        print(f"  k={k}:  {np.max(np.abs(T[k])):.3e}   {np.max(np.abs(Tn[k])):.3e}   bound 2^k * noise = {2 ** k * NOISE:.3e}")

    # ---- derivative formulas
    m = DERIVATIVE_ORDER
    print(f"\nforward formulas for the derivative of order {m}:")
    for kk in range(m, m + 5):
        w = forward_weights(m, kk)
        den = np.lcm.reduce([q.denominator for q in w])
        print(f"  {kk + 1} points: (1/({den} h^{m})) * {[int(q * den) for q in w]}")
    print("central formulas:")
    for r in (1, 2, 3):
        w = central_weights(m, r)
        print(f"  {2 * r + 1} points:", np.round(w, 8))

    # ---- plots
    fig, ax = plt.subplots(1, 3, figsize=(15, 4.2))
    xs = np.linspace(x[0], x[n], 500)
    if SOURCE == "function":
        ax[0].plot(xs, f(xs), "k", label="f")
    ax[0].plot(xs, [forward_eval(t, X0, H, T, K) for t in xs], "r", label=f"forward, degree {K}")
    ax[0].plot(xs, [backward_eval(t, x[n], H, T, K) for t in xs], "b--", label=f"backward, degree {K}")
    ax[0].plot(x, y, "o", color="gray")
    ax[0].legend(); ax[0].set_title("Newton-Gregory polynomials")
    ks = np.arange(n + 1)
    ax[1].semilogy(ks, [max(np.max(np.abs(c)), 1e-17) for c in T], "o-", label="clean")
    ax[1].semilogy(ks, [max(np.max(np.abs(c)), 1e-17) for c in Tn], "o-", label="disturbed")
    ax[1].semilogy(ks, 2.0 ** ks * NOISE, "k--", label="2^k noise")
    ax[1].set_xlabel("k"); ax[1].set_title("largest |difference| in column k"); ax[1].legend()
    if SOURCE == "function":
        hs = 10.0 ** np.linspace(-12, -0.5, 200)
        hr = 0.005 * max(abs(n * H), 1e-3)
        wr = central_weights(m, 6)
        ref = sum(wr[j + 6] * f(X0 + j * hr) for j in range(-6, 7)) / hr ** m
        for kk in range(m, m + 4):
            w = [float(q) for q in forward_weights(m, kk)]
            err = [abs(sum(w[j] * f(X0 + j * h) for j in range(kk + 1)) / h ** m - ref) for h in hs]
            ax[2].loglog(hs, err, label=f"forward, order {kk + 1 - m}")
        for r in (1, 2):
            if 2 * r + 1 >= m + 2:
                w = central_weights(m, r)
                err = [abs(sum(w[j + r] * f(X0 + j * h) for j in range(-r, r + 1)) / h ** m - ref) for h in hs]
                ax[2].loglog(hs, err, "--", label=f"central, order {2 * r}")
        ax[2].set_xlabel("h"); ax[2].set_ylabel("|error|"); ax[2].set_title(f"derivative of order {m} at x0"); ax[2].legend()
    plt.tight_layout()
    if SAVE:
        plt.savefig("finite_differences.png", dpi=130)
    plt.show()


if __name__ == "__main__":
    main()