%% Chapter 6 IVPs for ODEs
% This is a MATLAB Live Editor version of the Fundamentals of Numerical Computation Jupyter Notebooks found at https://tobydriscoll.net/project/fnc/.
%
% Example 6.1.2
% The equation u'=\sin[(u+t)^2] also has a solution that can be found numerically with ease, even though no formula exists for its solution.

f = @(t,u) sin( (t+u).^2 );
[t,u] = ode45(f,[0,4],-1);
plot(t,u)
xlabel('t'), ylabel('u(t)'), title('Solution of u''=sin[(t+u)^2]')     % ignore this line 

%
% Example 6.1.3
% We return to he equation u'=\sin[(u+t)^2] to see two variations on how to obtain numerical solutions for it. In the first one, we can supply a vector of time nodes and receive the solution at exactly those times.

f = @(t,u) sin( (t+u).^2 );
t = linspace(0,4,60)';
[t,u] = ode45(f,t,-1);
plot(t,u,'.')
xlabel('t'), ylabel('u(t)'), title('Solution at 60 points')     % ignore this line

%
% In the second variation, we see that it's also possible to create a callable function for the solution. Essentially this is a high-quality interpolant of the computed particular values.

u = @(t) deval( ode45(f,[0,4],-1), t );
fplot(u,[0,4])
xlabel('t'), ylabel('u(t)'), title('Solution as a callable function')     % ignore this line

%
% Example 6.1.5
% The equation u'=(u+t)^2, gives us some trouble.

warning on    % ignore this line
f = @(t,u) (t+u).^2;
[t,u] = ode45(f,[0,1],1);
semilogy(t,u)
xlabel('t'), ylabel('u(t)'), title('Finite time blowup')

% The warning message we received can mean that there is a bug in the formulation of the problem. But if everything has been done correctly, it suggests that the solution does not exist past the indicated time.
%
% Example 6.1.6
% Consider the ODEs u'=u and u'=-u. In each case we compute \partial f/\partial u = \pm 1, so the condition number bound is e^{(b-a)} in both problems. However, they behave quite differently. In the case of exponential growth, u'=u, the bound is the actual condition number.

for u0 = [0.7 1 1.3]
    fplot(@(t) exp(t)*u0,[0 3]), hold on
end
xlabel('t'), ylabel('u(t)')   % ignore this line
title('Exponential divergence of solutions')   % ignore this line

%
% But with u'=-u, solutions actually get closer together with time.

clf
for u0 = [0.7 1 1.3]
    fplot(@(t) exp(-t)*u0,[0 3]), hold on
end
xlabel('t'), ylabel('u(t)')   % ignore this line
title('Exponential convergence of solutions')   % ignore this line
hold off

% In this case the actual condition number is one, due to the difference of solutions at the initial time.
%
% Example 6.2.1
% We consider the IVP u'=\sin[(u+t)^2] over 0\le t \le 4, with u(0)=-1.

f = @(t,u) sin( (t+u).^2 );
a = 0;  b = 4;
u0 = -1;
[t,u] = eulerivp(f,[a,b],u0,20);
plot(t,u,'.-')
xlabel('t'), ylabel('u(t)'), title('Solution by Euler''s method')  % ignore this line

%
% We could define a different interpolant to get a smoother picture above, but the derivation assumed the piecewise linear interpolant, so it is the most meaningful one. We can instead request more steps to make the interpolant look smoother.

[t,u] = eulerivp(f,[a,b],u0,200);
hold on, plot(t,u,'-')
legend('n=20','n=200','location','southwest')  % ignore this line

% Increasing n changed the solution noticeably. Since we know that interpolants and finite differences become more accurate as h\to 0, we should expect that from Euler's method too.
%
% We don't have an exact solution to compare to, so we will use the built-in solver ode113 with settings chosen to construct an accurate solution.

opt = odeset('abstol',5e-14,'reltol',5e-14);  
uhat = ode113(f,[a,b],u0,opt);
u_exact = @(t) deval(uhat,t)';
fplot(u_exact,[a,b],'k-')
legend('n=20','n=200','accurate','location','southwest')  % ignore this line

%
% Now we can perform a convergence study.

n = 50*2.^(0:5)';
err = 0*n;
for j = 1:length(n)
    [t,u] = eulerivp(f,[a,b],u0,n(j));
    err(j) = max(abs(u_exact(t)-u));
end
table(n,err)

%
% The error is almost perfectly halved at each step, so we expect that a log-log plot will reveal first-order convergence.

clf
loglog(n,err,'.-'), hold on
loglog(n,0.05*(n/n(1)).^(-1),'--')
axis tight, xlabel('n'), ylabel('inf-norm error')  % ignore this line
title('Convergence of Euler''s method')  % ignore this line
legend('error','1st order','location','southwest')  % ignore this line
hold off

%
% Example 6.3.2
% Let A be an m\times n time dependent matrix. Then, \frac{d\mathbf{u}}{dt}=A(t)\mathbf{u} is a linear system of differential equations. If A is time independnet, then ther solution of the system is given by the matrix exponential, \mathbf{u}(t)=e^{tA}\mathbf{u}_0, where \mathbf{u}_0=\mathbf{u}(0).
% Let

A = [ -2 5; -1 0 ]

%

u0 = [1;0];
t = linspace(0,6,600);  % times for plotting
u = zeros(length(t),length(u0));
for j=1:length(t)
    ut = expm(t(j)*A)*u0;  
    u(j,:) = ut';
end

%

plot(t,u)

%
% Example 6.3.3
% We encode the predator--prey equations via an included function called predprey, which is found at the end of this file. To solve the IVP we must also provide the initial condition, which is a 2-vector here, and the interval for the independent variable.

u0 = [1; 0.01];
t = linspace(0,80,2001)';       % specify where the time outputs will be
[t,u] = ode45(@predprey,t,u0);

%
% Each row of the output u is the value of the solution vector at a single node in time. Each column of u is the entire history of one component of the solution.

size_u = size(u)
y = u(:,1);  z = u(:,2);  
plot(t,y,t,z)
xlabel('t'), ylabel('u(t)'), title('Predator-prey solution')  % ignore this line
legend('prey','predator')  % ignore this line

%
% When there are just two components, it's common to plot the solution in the phase plane, i.e., with u_1 and u_2 along the axes and time as a parameterization of the curve.

plot(y,z)
xlabel('y'), ylabel('z'), title('Predator-prey phase plane')  % ignore this line

% From this plot we can deduce that the solution approaches a periodic one, which in the phase plane is reprepresented by a closed loop.
%
% Example 6.4.1
% We consider the IVP u'=\sin[(u+t)^2] over 0\le t \le 4, with u(0)=-1.

f = @(t,u) sin( (t+u).^2 );
a = 0;  b = 4;
u0 = -1;

%
% We use the built-in solver ode113 to construct an accurate approximation to the exact solution.

opt = odeset('abstol',5e-14,'reltol',5e-14); 
uhat = ode113(f,[a,b],u0,opt);
u_exact = @(t) deval(uhat,t)';

%
% Now we perform a convergence study of our two Runge--Kutta implementations.

n = 50*2.^(0:5)';
err_IE2 = 0*n;
err_RK4 = 0*n;
for j = 1:length(n)
    [t,u] = ie2(f,[a,b],u0,n(j));
    err_IE2(j) = max(abs(u_exact(t)-u));
    [t,u] = rk4(f,[a,b],u0,n(j));
    err_RK4(j) = max(abs(u_exact(t)-u));
end

%
% The amount of computational work at each time step is assumed to be proportional to the number of stages. Let's compare on an apples-to-apples basis by using the number of f-evaluations on the horizontal axis.

loglog(2*n,err_IE2,'.-')
hold on, loglog(4*n,err_RK4,'.-')
loglog(2*n,0.01*(n/n(1)).^(-2),'--')
loglog(4*n,1e-6*(n/n(1)).^(-4),'--')
xlabel('n'), ylabel('inf-norm error')  % ignore this line
title('Convergence of the improved Euler method')  % ignore this line
axis tight  % ignore this line
legend('IE2','RK4','2nd order','4th order','location','southwest')  % ignore this line
hold off

% The fourth-order variant is more efficient in this problem over a wide range of accuracy.
%
%
% Example 6.5.1
% Let's run adaptive RK on  u'=e^{t-u\sin u}.

f = @(t,u) exp(t-u*sin(u));
[t,u] = rk23(f,[0,5],0,1e-5);
plot(t,u,'.')
xlabel('t'),  ylabel('u(t)')  % ignore this line
title('Adaptive IVP solution')  % ignore this line

%
% The solution makes a very abrupt change near t=2.4. The resulting time steps vary over three orders of magnitude.

semilogy(t(1:end-1),diff(t))
axis tight, title('Adaptive step sizes')  % ignore this line  
xlabel('t'),  ylabel('step size')  % ignore this line

%
% If we had to run with a uniform step size to get this accuracy, it would be

h_min = min(diff(t))

%
% On the other hand, the average step size that was actually taken was

h_avg = mean(diff(t))

% We took fewer steps by a factor of almost 1000! Even accounting for the extra stage per step and the occasional rejected step, the savings are clear.
%
% Example 6.5.2

f = @(t,u) (t+u).^2;
warning on
[t,u] = rk23(f,[0,1],1,1e-5);
semilogy(t,u)
xlabel('t'), ylabel('u(t)'), title('Finite time blowup')  % ignore this line
warning off

%
% Example 6.7.1
% We consider the IVP u'=\sin[(u+t)^2] over 0\le t \le 4, with u(0)=-1.

f = @(t,u) sin( (t+u).^2 );
a = 0;  b = 4;
u0 = -1;

%
% We use the built-in solver ode113 to construct an accurate approximation to the exact solution.

opt = odeset('abstol',5e-14,'reltol',5e-14); 
uhat = ode113(f,[a,b],u0,opt);
u_exact = @(t) deval(uhat,t)';  % callable function

%
% Now we perform a convergence study of the AB4 code.

n = 10*2.^(0:5)';
err = 0*n;
for j = 1:length(n)
    [t,u] = ab4(f,[a,b],u0,n(j));
    err(j) = max(abs(u_exact(t)-u));
end

%
% The method should converge as O(h^4), so a log-log scale is appropriate for the errors.

loglog(n,err,'.-')
xlabel('n'), ylabel('error')  % ignore this line
title('Convergence of AB4')   % ignore this line
hold on, loglog(n,0.1*(n/n(1)).^(-4),'--')
axis tight   % ignore this line
legend('AB4','4th order','location','southwest')   % ignore this line
hold off

%
% Example 6.7.2
% The following simple ODE uncovers a surprise.

f = @(t,u) u.^2 - u.^3;
u0 = 0.005;

%
% We will solve the problem first with the implicit AM2 method using n=200 steps.

[tI,uI] = am2(f,[0 400],u0,200);
plot(tI,uI)
xlabel('t'), ylabel('u(t)')   % ignore this line
title('AM2 solution')   % ignore this line

%
% Now we repeat the process using the explicit AB4 method.

[tE,uE] = ab4(f,[0 400],u0,200);
hold on, plot(tE,uE), ylim([-1 2])
title('AM2 and AB4 solutions')  % ignore this line

%
% Once the solution starts to take off, the AB4 result goes catastrophically wrong.

format short e, uE(105:111)

%
% We hope that AB4 will converge in the limit h\to 0, so let's try using more steps.

for n = [1000 1600]
    [tE,uE] = ab4(f,[0 400],u0,n);
    plot(tE,uE)
end

% So AB4, which is supposed to be more accurate than AM2, actually needs something like 8 times as many steps to get a reasonable-looking answer!
%
% Example 6.8.1
% Consider the ridiculously simple IVP u'=u, u(0)=1, whose solution is e^t.

dudt = @(t,u) u;
u_exact = @exp;
a = 0;  b = 1;

%
% Let's apply the LIAF method to this problem for varying fixed step sizes. We'll measure the error at the time t=1.

n = [5;10;20;40;60];
err = 0*n;
for j = 1:length(n)
    h = (b-a)/n(j);
    t = a + h*(0:n(j))';
    u = [1; u_exact(h); zeros(n(j)-1,1)];
    f = [dudt(t(1),u(1)); zeros(n(j)-2,1)];
    for i = 2:n(j)
        f(i) = dudt(t(i),u(i));
        u(i+1) = -4*u(i) + 5*u(i-1) + h*(4*f(i)+2*f(i-1));
    end
    err(j) = abs(u_exact(b) - u(end));
end

h = (b-a)./n;
table(n,h,err)

%
% The error starts out promisingly, but things explode from there. A graph of the last numerical attempt yields a clue.

figure
semilogy(t,abs(u),'.-')
xlabel('t'),  ylabel('|u|')  % ignore this line
title('LIAF solution'), axis tight   % ignore this line

% It's clear that the solution is growing exponentially in time.
%
%
%% Functions for Chapter 6
% Note that the function for the ODE must accept both t and u inputs, even though there is no explicit dependence on t.

function dudt = predprey(t,u)
    alpha = 0.1;  beta = 0.25;
    y = u(1);  z = u(2);
    s = (y*z) / (1+beta*y);  % appears in both equations
    dudt = [ y*(1-alpha*y) - s;  -z + s ];
end

function [t,u] = eulerivp(dudt,tspan,u0,n)
% EULERIVP   Euler's method for a scalar initial-value problem.
% Input:
%   dudt    defines f in u'(t)=f(t,u). (function)
%   tspan   endpoints of time interval (2-vector)
%   u0      initial value (scalar) 
%   n       number of time steps (integer)
% Output:
%   t       selected nodes  (vector, length n+1)
%   u       solution values   (vector, length n+1)

a = tspan(1);  b = tspan(2);
h = (b-a)/n;
t = a + (0:n)'*h;
u = zeros(n+1,1);
u(1) = u0;
for i = 1:n
  u(i+1) = u(i) + h*dudt(t(i),u(i));
end
end

function [t,u] = ie2(dudt,tspan,u0,n)
% IE2    Improved Euler method for an IVP.
% Input:
%   dudt    defines f in u'(t)=f(t,u) (function)
%   tspan   endpoints of time interval (2-vector)
%   u0      initial value (vector, length m)
%   n       number of time steps (integer)
% Output:
%   t       selected nodes (vector, length N+1)
%   u       solution values (array, (n+1)-by-m)

% Time discretization.
a = tspan(1);  b = tspan(2);
h = (b-a)/n;   
t = a + h*(0:n)';

% Initialize solution array. 
u = zeros(length(u0),n+1);  
u(:,1) = u0;

% Time stepping. 
for i = 1:n 
    uhalf = u(:,i) + h/2 * dudt(t(i),u(:,i));
    u(:,i+1) = u(:,i) + h * dudt(t(i)+h/2,uhalf);
end

u = u.';   % conform with MATLAB output convention
end

function [t,u] = rk4(dudt,tspan,u0,n)
% RK4    Fourth-order Runge-Kutta for an IVP.
% Input:
%   dudt    defines f in u'(t)=f(t,u)  (function)
%   tspan   endpoints of time interval (2-vector)
%   u0      initial value (vector, length m)
%   n       number of time steps (integer)
% Output:
%   t       selected nodes (vector, length n+1)
%   u       solution values (array, (n+1) by m)

% Define time discretization.
a = tspan(1);  b = tspan(2);
h = (b-a)/n;
t = a + (0:n)'*h;

% Initialize solution array.
u = zeros(length(u0),n+1);
u(:,1) = u0(:);

% Time stepping.
for i = 1:n
  k1 = h*dudt( t(i),     u(:,i) );
  k2 = h*dudt( t(i)+h/2, u(:,i)+k1/2 );
  k3 = h*dudt( t(i)+h/2, u(:,i)+k2/2 );
  k4 = h*dudt( t(i)+h,   u(:,i)+k3 );
  u(:,i+1) = u(:,i) + (k1 + 2*(k2 + k3) + k4)/6;
end

u = u.';   % conform to MATLAB output convention
end

function [t,u] = rk23(dudt,tspan,u0,tol)
% RK23   Adaptive IVP solver based on embedded RK formulas.
% Input:
%   dudt    defines f in u'(t)=f(t,u)  (function)
%   tspan   endpoints of time interval (2-vector)
%   u0      initial value (vector, length m)
%   tol     global error target (positive scalar)
% Output:
%   t       selected nodes (vector, length n+1)
%   u       solution values (array, size m by n+1)

% Initialize for the first time step.
t = tspan(1);
u(:,1) = u0(:);   i = 1;
h = 0.5*tol^(1/3);
s1 = dudt(t(1), u(:,1));

% Time stepping.
while t(i) < tspan(2)
    
    % Detect underflow of the step size.
    if t(i)+h == t(i)
        warning('Stepsize too small near t=%.6g.',t(i))
        break  % quit time stepping loop
    end
    
    % New RK stages.
    s2 = dudt( t(i)+h/2,   u(:,i)+(h/2)*s1   );
    s3 = dudt( t(i)+3*h/4, u(:,i)+(3*h/4)*s2 );
    unew2 = u(:,i) + h*(2*s1 + 3*s2 + 4*s3)/9;   % 2rd order solution
    s4 = dudt( t(i)+h,     unew2 );
    err = h*(-5*s1/72 + s2/12 + s3/9 - s4/8);    % 2nd/3rd order difference
    E = norm(err,inf);                           % error estimate
    maxerr = tol*(1 + norm(u(:,i),inf));  % relative/absolute blend
    
    % Accept the proposed step? 
    if E < maxerr     % yes 
        t(i+1) = t(i) + h;
        u(:,i+1) = unew2;
        i = i+1;
        s1 = s4;      % use FSAL property
    end
    
    % Adjust step size. 
    q = 0.8*(maxerr/E)^(1/3);       % conservative optimal step factor
    q = min(q,4);                   % limit stepsize growth
    h = min(q*h,tspan(2)-t(i));     % don't step past the end
end

t = t';  u = u.';   % conform to MATLAB output convention
end

function [t,u] = ab4(dudt,tspan,u0,n)
%AB4     4th-order Adams-Bashforth formula for an IVP.
% Input:
%   dudt    defines f in u'(t)=f(t,u) (function)
%   tspan   endpoints of time interval (2-vector)
%   u0      initial value (m-vector)
%   n       number of time steps (integer)
% Output:
%   t       selected nodes (vector, length n+1)
%   u       solution values (array, size n+1 by m)

% Discretize time.
a = tspan(1);  b = tspan(2);
h = (b-a)/n;
t = tspan(1) + (0:n)'*h;

% Constants in the AB4 method.
k = 4;  
sigma = [55; -59; 37; -9]/24;  

% Find starting values by RK4.
u = zeros(length(u0),n+1);
[ts,us] = rk4(dudt,[a, a+(k-1)*h],u0,k-1);
u(:,1:k) = us(1:k,:).';

% Compute history of u' values, from oldest to newest.
f = zeros(length(u0),k);
for i = 1:k-1
  f(:,k-i) = dudt(t(i),u(:,i));
end

% Time stepping.
for i = k:n
  f = [dudt(t(i),u(:,i)), f(:,1:k-1)];   % new value of du/dt
  u(:,i+1) = u(:,i) + h*(f*sigma);       % advance one step
end

u = u.';   % conform to MATLAB output convention

end

function [t,u] = am2(dudt,tspan,u0,n)
% AM2    2nd-order Adams-Moulton (trapezoid) formula for an IVP.
% Input:
%   dudt    f(t,y) for the ODE (function)
%   tspan   endpoints of time interval (2-vector)
%   u0      initial value (m-vector)
%   n       number of time steps (integer)
% Output:
%   t       vector of times (vector, length n+1)
%   u       solution (array, size n+1 by m)

% Discretize time.
a = tspan(1);  b = tspan(2);
h = (b-a)/n;
t = tspan(1) + (0:n)'*h;

m = numel(u0);
u = zeros(m,n+1);
u(:,1) = u0(:);

% Time stepping.
for i = 1:n
  % Data that does not depend on the new value.
  known = u(:,i) + h/2*dudt(t(i),u(:,i));
  % Find a root for the new value. 
  unew = levenberg(@trapzero,known);
  u(:,i+1) = unew(:,end);
end

u = u.';   % conform to MATLAB output convention

    % This function defines the rootfinding problem at each step.
    function F = trapzero(z)
        F = z - h/2*dudt(t(i+1),z) - known;
    end

end  % main function


function x = levenberg(f,x1,tol)
% LEVENBERG   Quasi-Newton method for nonlinear systems.
% Input:
%   f         objective function 
%   x1        initial root approximation
%   tol       stopping tolerance (default is 1e-12)
% Output       
%   x         array of approximations (one per column)

% Operating parameters.
if nargin < 3, tol = 1e-12; end
ftol = tol;  xtol = tol;  maxiter = 40;

x = x1(:);     fk = f(x1);
k = 1;  s = Inf;        
Ak = fdjac(f,x(:,1),fk);   % start with FD Jacobian
jac_is_new = true;
I = eye(length(x));

lambda = 10; 
while (norm(s) > xtol) && (norm(fk) > ftol) && (k < maxiter)
    % Compute the proposed step.
    B = Ak'*Ak + lambda*I;
    z = Ak'*fk;
    s = -(B\z);

    xnew = x(:,k) + s;   fnew = f(xnew);
    
    % Do we accept the result?
    if norm(fnew) < norm(fk)    % accept
        y = fnew - fk;
        x(:,k+1) = xnew;  fk = fnew;  
        k = k+1;
        
        lambda = lambda/10;  % get closer to Newton
        % Broyden update of the Jacobian.
        Ak = Ak + (y-Ak*s)*(s'/(s'*s));
        jac_is_new = false;
    else                       % don't accept
        % Get closer to steepest descent.
        lambda = lambda*4;
        % Re-initialize the Jacobian if it's out of date.
        if ~jac_is_new
            Ak = fdjac(f,x(:,k),fk);
            jac_is_new = true;
        end
    end
end

if (norm(fk) > 1e-3)
    warning('Iteration did not find a root.')
end

end

function J = fdjac(f,x0,y0)
% FDJAC   Finite-difference approximation of a Jacobian.
% Input:
%   f        function to be differentiated
%   x0       evaluation point (n-vector)
%   y0       value of f at x0 (m-vector)
% Output       
%   J        approximate Jacobian (m-by-n)

delta = sqrt(eps);   % FD step size
m = length(y0);  n = length(x0);
J = zeros(m,n);
I = eye(n);
for j = 1:n
    J(:,j) = ( f(x0+delta*I(:,j)) - y0) / delta;
end

end