#!/usr/bin/env python3
"""Chapter 10, code 1: a linear boundary value problem.

Solve y'' + p(x)*y' + q(x)*y = r(x), with y(a) = alpha and y(b) = beta.
Method: centered finite differences and the Thomas algorithm.
Example: y'' - y = 0, y(0) = y(1) = 1.
Exact solution: cosh(x-0.5)/cosh(0.5).

Requires NumPy and Matplotlib. Run: python3 ch10_finite_difference.py
SciPy alternatives: solve_banded can solve the assembled tridiagonal
system in banded storage, or spsolve can solve a sparse matrix system.
The Thomas routine is retained here to show the elimination steps.

To solve another problem:
  1. Edit the interval, boundary values, and grid in main.
  2. Edit coefficients immediately below the driver.
  3. Disable check_exact or replace exact_solution.
"""

import numpy as np
import matplotlib.pyplot as plt


def main():
    # ------------------- Problem and grid: edit here ------------------
    a = 0.0
    b = 1.0
    alpha = 1.0                 # y(a)
    beta = 1.0                  # y(b)
    n = 20                      # n intervals, n+1 points, n-1 unknowns
    check_exact = True
    make_plot = True
    # -----------------------------------------------------------------
    if n < 2 or n != int(n) or b <= a:
        raise ValueError('Use integer n >= 2 and b > a.')
    h = (b-a)/n
    x = np.linspace(a, b, n+1)
    y = finite_difference(coefficients, n, x, alpha, beta)

    print('Linear boundary value problem: finite differences')
    print(f'Grid intervals: {n:8d}')
    print(f'Grid spacing:   {h:12.6f}\n')
    filename = 'ch10_finite_difference.dat'
    if check_exact:
        exact = exact_solution(x, a, b, alpha, beta)
        abs_error = np.abs(y-exact)
    with open(filename, 'w', encoding='utf-8') as data_file:
        if check_exact:
            print(f"{'x':>10} {'numerical y':>16} {'exact y':>16} {'abs error':>12}")
            data_file.write('# x numerical_y exact_y abs_error\n')
        else:
            print(f"{'x':>10} {'numerical y':>16}")
            data_file.write('# x numerical_y\n')
        for i in range(n+1):
            if check_exact:
                print(f'{x[i]:10.6f} {y[i]:16.8e} {exact[i]:16.8e} {abs_error[i]:12.4e}')
                row = (x[i], y[i], exact[i], abs_error[i])
            else:
                print(f'{x[i]:10.6f} {y[i]:16.8e}')
                row = (x[i], y[i])
            data_file.write(' '.join(f'{value:.16e}' for value in row)+'\n')
    print()
    if check_exact:
        print(f'Largest grid error: {np.max(abs_error):12.4e}')
    print(f'Data file: {filename}')

    if make_plot:
        panels = 2 if check_exact else 1
        fig, axes = plt.subplots(panels, 1, figsize=(7, 7 if check_exact else 4),
                                 squeeze=False)
        ax = axes[0, 0]
        ax.plot(x, y, 'ko-', linewidth=1, markersize=4, label='Finite differences')
        if check_exact:
            x_plot = np.linspace(a, b, 401)
            ax.plot(x_plot, exact_solution(x_plot, a, b, alpha, beta),
                    'k--', linewidth=1.2, label='Exact solution')
            ax.legend()
        ax.set_xlabel('x')
        ax.set_ylabel('y(x)')
        ax.set_title('Linear boundary value problem')
        ax.grid(True)
        if check_exact:
            ax = axes[1, 0]
            ax.plot(x, abs_error, 'ko-', linewidth=1, markersize=4)
            ax.set_xlabel('x')
            ax.set_ylabel('Absolute error')
            ax.grid(True)
        fig.tight_layout()
        plt.show()


def coefficients(x):
    """Supply p, q, r in y'' + p(x)*y' + q(x)*y = r(x).

    Default example: y'' - y = 0.
    """
    p = 0.0
    q = -1.0
    r = 0.0
    return p, q, r


def exact_solution(x, a, b, alpha, beta):
    """Exact solution of y'' - y = 0 with the supplied endpoint values.

    Replace when changing the equation, or disable the exact check.
    """
    return (alpha*np.sinh(b-x)+beta*np.sinh(x-a))/np.sinh(b-a)


def finite_difference(fcn, n, x, alpha, beta):
    """Assemble the interior tridiagonal system using centered differences.

    y'' = (y_left-2*y_center+y_right)/h^2
    y'  = (y_right-y_left)/(2*h).
    Multiply the equation by h^2. Interior unknown i lies at x[i+1]
    because Python indices start at zero. Move known endpoint terms
    to the right-hand side before solving.
    """
    h = (x[-1]-x[0])/n
    m = n-1
    lower = np.zeros(m)
    diagonal = np.zeros(m)
    upper = np.zeros(m)
    rhs = np.zeros(m)
    for i in range(m):
        p, q, r = fcn(x[i+1])
        lower[i] = 1-h*p/2
        diagonal[i] = -2+h*h*q
        upper[i] = 1+h*p/2
        rhs[i] = h*h*r

    # Both corrections apply when there is only one interior point.
    rhs[0] -= lower[0]*alpha
    rhs[-1] -= upper[-1]*beta
    lower[0] = 0
    upper[-1] = 0
    interior = thomas(lower, diagonal, upper, rhs)
    y = np.zeros(n+1)
    y[0] = alpha
    y[1:n] = interior
    y[n] = beta
    return y


def thomas(lower, diagonal, upper, rhs):
    """Forward elimination and backward substitution for a tridiagonal system.

    Copies preserve the caller's diagonal and rhs arrays.
    No pivoting is used. A zero or near-zero pivot stops the calculation;
    this may require a solver with pivoting or indicate a singular BVP.
    """
    diagonal = diagonal.copy()
    rhs = rhs.copy()
    m = len(rhs)
    y = np.zeros(m)
    for i in range(m):
        scale = max(abs(lower[i]), abs(diagonal[i]), abs(upper[i]))
        if abs(diagonal[i]) <= np.finfo(float).eps*scale:
            raise RuntimeError('Thomas algorithm: zero or near-zero pivot.')
        if i == m-1:
            break
        multiplier = lower[i+1]/diagonal[i]
        diagonal[i+1] -= multiplier*upper[i]
        rhs[i+1] -= multiplier*rhs[i]
    y[-1] = rhs[-1]/diagonal[-1]
    for i in range(m-2, -1, -1):
        y[i] = (rhs[i]-upper[i]*y[i+1])/diagonal[i]
    return y


if __name__ == '__main__':
    main()