%=======================================================================
% Chapter 10, code 3: a nonlinear boundary value problem.
% Solve y'' = f(x,y,y'), with y(a) = alpha and y(b) = beta.
% Method: centered finite differences and Newton iteration.
% Each Newton correction is found with the Thomas algorithm.
%
% Example on [0,1]:
%   y'' = y^2 - sin(pi*x)^2 - pi^2*sin(pi*x)
%   y(0) = y(1) = 0; exact solution y = sin(pi*x).
% The forcing was chosen so that the exact solution is known.
%
% MATLAB alternatives: solve each Newton system with a sparse matrix and
% correction = J\(-R), or use bvp4c to solve the BVP by collocation.
% bvp4c requires the first-order system and an initial solution guess.
%
% To solve another problem:
%   1. Edit the interval, endpoint values, grid, and initial guess.
%   2. Edit equation below, including its two partial derivatives.
%   3. Disable check_exact or replace exact_solution.
%=======================================================================
clear;

%-------------------- Problem and controls: edit here --------------------
a = 0.0;
b = 1.0;
alpha = 0.0;                 % y(a)
beta = 0.0;                  % y(b)
n = 20;                     % Number of grid intervals
tolerance = 1.0e-10;         % Maximum absolute discrete-equation residual
max_iterations = 30;
check_exact = true;
make_plot = true;
%------------------------------------------------------------------------
if n < 2 || n ~= floor(n) || b <= a
    error('Use an integer n >= 2 and b > a.');
end
if tolerance <= 0 || max_iterations < 1 || max_iterations ~= floor(max_iterations)
    error('Use positive tolerance and an integer iteration limit.');
end
h = (b-a)/n;
x = linspace(a,b,n+1)';

% Initial guess: a straight line between the endpoint values.
% For another initial guess, change this expression.
y = alpha+(beta-alpha)*(x-a)/(b-a);
y(1) = alpha;
y(end) = beta;

fprintf('Nonlinear boundary value problem: finite differences and Newton iteration\n');
fprintf('Grid intervals:   %8d\n',n);
fprintf('Grid spacing:     %12.6f\n',h);
fprintf('Newton tolerance: %12.4e\n',tolerance);
[y,iterations,residual_norm] = newton_bvp(@equation,n,x,y, ...
                                       tolerance,max_iterations);

filename = 'ch10_nonlinear_bvp.dat';
fid = fopen(filename,'w');
if fid == -1
    error('Cannot open the data file.');
end

fprintf('\n');

if check_exact
    exact = exact_solution(x);
    abs_error = abs(y-exact);
    fprintf('%10s %16s %16s %12s\n', ...
            'x','numerical y','exact y','abs error');
    fprintf(fid,'# x numerical_y exact_y abs_error\n');
    for i = 1:n+1
        fprintf('%10.6f %16.8e %16.8e %12.4e\n', ...
                x(i),y(i),exact(i),abs_error(i));
        fprintf(fid,' %.16e %.16e %.16e %.16e\n', ...
                x(i),y(i),exact(i),abs_error(i));
    end
else
    fprintf('%10s %16s\n','x','numerical y');
    fprintf(fid,'# x numerical_y\n');
    for i = 1:n+1
        fprintf('%10.6f %16.8e\n',x(i),y(i));
        fprintf(fid,' %.16e %.16e\n',x(i),y(i));
    end
end
fclose(fid);

fprintf('\n');
fprintf('Newton iterations: %8d\n',iterations);
fprintf('Final residual:    %12.4e\n',residual_norm);
if check_exact
    fprintf('Largest grid error: %12.4e\n',max(abs_error));
end
fprintf('Data file: %s\n',filename);

if make_plot
    figure;
    if check_exact
        subplot(2,1,1);
    end
    plot(x,y,'ko-','LineWidth',1.0,'MarkerSize',4);
    hold on;
    if check_exact
        x_plot = linspace(a,b,401)';
        plot(x_plot,exact_solution(x_plot), ...
             'k--','LineWidth',1.2);
        legend('Finite differences','Exact solution','Location','best');
    end
    hold off;
    xlabel('x');
    ylabel('y(x)');
    title('Nonlinear boundary value problem');
    grid on;

    if check_exact
        subplot(2,1,2);
        plot(x,abs_error,'ko-','LineWidth',1.0,'MarkerSize',4);
        xlabel('x');
        ylabel('Absolute error');
        grid on;
    end
end


function [f,dfdy,dfdv] = equation(x,y,v)
% Supply y'' = f(x,y,v), where v = y'.
% Also supply df/dy and df/dv for Newton's method. These derivatives treat
% x, y, and v as independent variables. Set dfdv = 0 if f is independent of v.
s = sin(pi*x);
f = y*y-s*s-pi*pi*s;
dfdy = 2*y;
dfdv = 0.0;
end


function y = exact_solution(x)
% Exact solution for the supplied equation and default boundary values.
y = sin(pi*x);
end


function [y,iterations,residual_norm] = newton_bvp(fcn,n,x,y,tolerance,max_iterations)
%=======================================================================
% At each interior point define
%   R = (y_left-2*y_center+y_right)/h^2 - f(x,y_center,v),
%   v = (y_right-y_left)/(2*h).
%
% Newton's method solves J*correction = -R and updates y.
% J is tridiagonal. Multiplying the Newton equations by h^2 gives:
%   lower    =  1 + h*(df/dv)/2
%   diagonal = -2 - h^2*(df/dy)
%   upper    =  1 - h*(df/dv)/2
%   rhs      = -h^2*R.
%
% MATLAB interior unknown i lies at x(i+1). Endpoint corrections are zero.
% Stop when max|R| <= tolerance. This measures convergence of the discrete
% equations; grid refinement checks discretization error.
% Full Newton steps are used. An unsuitable initial guess may cause
% failure or convergence to a different solution when several exist.
%=======================================================================
h = (x(end)-x(1))/n;
m = n-1;
lower = zeros(m,1);
diagonal = zeros(m,1);
upper = zeros(m,1);
rhs = zeros(m,1);
change = 0;
fprintf('%10s %18s %18s\n','iteration','max residual','max correction');

for iterations = 0:max_iterations
    residual_norm = 0;
    for i = 1:m
        j = i+1;
        v = (y(j+1)-y(j-1))/(2*h);
        [f,dfdy,dfdv] = fcn(x(j),y(j),v);
        if ~all(isfinite([f,dfdy,dfdv]))
            error('Nonfinite equation or partial derivative; check the initial guess.');
        end
        residual = (y(j-1)-2*y(j)+y(j+1))/(h*h)-f;
        residual_norm = max(residual_norm,abs(residual));
        lower(i) = 1+h*dfdv/2;
        diagonal(i) = -2-h*h*dfdy;
        upper(i) = 1-h*dfdv/2;
        rhs(i) = -h*h*residual;
    end
    % The correction shown is the one applied in the preceding iteration.
    fprintf('%10d %18.10e %18.10e\n',iterations,residual_norm,change);
    if residual_norm <= tolerance
        return;
    end
    if iterations == max_iterations
        break;
    end

    lower(1) = 0;
    upper(m) = 0;
    correction = thomas(lower,diagonal,upper,rhs);
    change = max(abs(correction));
    if ~all(isfinite(correction))
        error('Nonfinite Newton correction; check the initial guess.');
    end
    if change <= eps*max(1,max(abs(y)))
        error('Newton iteration cannot advance; check the tolerance and initial guess.');
    end
    y(2:n) = y(2:n)+correction;
end
error('Newton iteration did not converge within the iteration limit.');
end


function y = thomas(lower,diagonal,upper,rhs)
%=======================================================================
% Solve a tridiagonal system by forward elimination and backward
% substitution. The local copies of diagonal and rhs are overwritten.
%
% No pivoting is used. A near-zero pivot stops the calculation.
% Such a failure can require a solver with pivoting or can indicate
% a singular boundary value problem.
%=======================================================================
m = length(rhs);
y = zeros(m,1);

% Check each pivot before division.
for i = 1:m
    scale = max([abs(lower(i)),abs(diagonal(i)),abs(upper(i))]);
    if abs(diagonal(i)) <= eps*scale
        error('Thomas algorithm: zero or near-zero pivot.');
    end
    if i == m
        break;
    end
    multiplier = lower(i+1)/diagonal(i);
    diagonal(i+1) = diagonal(i+1)-multiplier*upper(i);
    rhs(i+1) = rhs(i+1)-multiplier*rhs(i);
end

y(m) = rhs(m)/diagonal(m);
for i = m-1:-1:1
    y(i) = (rhs(i)-upper(i)*y(i+1))/diagonal(i);
end
end