%% Model Mumbai Plague
% Using Levenberg's Quasi-Newton Method, model the data w=[5, 10, 17, 22, 30, 50, 51, 90, 120, 180, 292, 395, 445, 775, 780,...\\
%     700, 698, 880, 925, 800, 578, 400, 350, 202, 105, 65, 55, 40, 30, 20];
% by the function w(t)=A\sech^2[B(t-C)].
% First, we make sure levenberg is working. Recall the nonlinear system from class:

x1 = [0;0;0];   
x = levenberg(@nlsystem2,x1)

%
% We first plot the data. Note that the population was rescaled by 1000 to keep the data range small. The time period was not specified except it was implied they were equally spaced units.

clear
global t w
w=[5, 10, 17, 22, 30, 50, 51, 90, 120, 180, 292, 395, 445, 775, 780,...
    700, 698, 880, 925, 800, 578, 400, 350, 202, 105, 65, 55, 40, 30, 20];
w=w/1000;
t=1:30;
plot(t,w,'.','MarkerSize',15)

%
% We then need a guess for gthe parameters. It looks like the data peaks near w=1and t=17. Plot the model function and play with B until it looks satisfactory.

A=1;
B=.2;
C=17;
z=linspace(0,30,100);
plot(t,w,'.','MarkerSize',15)
hold on
plot(z,A*sech(B*(z-C)).^2);
hold off

%
% Run Levenberg  using the misfit function. \mathbf{f}(\mathbf{c}) =\sum_{i=1}^m \bigl[ g(t_i;\mathbf{c})-y_i \bigr]^2. Then, Levenberg seeks to minimize the norm, \min_{\mathbf{c}\in\mathbb{R}^n} \bigl\| \mathbf{f}(\mathbf{c}) \bigr\|^2=\min_{\mathbf{c}\in\mathbb{R}^n} \sum_{i=1}^m \bigl[ g(t_i;\mathbf{c})-y_i \bigr]^2.

x1 = [A;B;C];   
x = levenberg(@misfit,x1)

%
% Oops! It did not converge. Looking at the reason, we see that the message is printed when norm(f_k) > 10^{-3}.  But if you look at the parameter values in x, then the values do not appear to be changing a lot. They just are not as precise as the routine would like.
% Let's use the last values to create a model function and see how well it fits.

a = x(1,end),  b = x(2,end), c=x(3,end)   % final values
model = @(x) a*sech(b*(x-c)).^2;

final_misfit_norm = norm(model(t)-w) 
hold on
fplot( model, [0 30],'LineWidth',2 )
title('Mumbai Plague Fit')   
ylabel('Deaths Per Week'), xlabel('Time')    
hold off

% Due to the scattering of the data, it seems that this is the best we can do.
%
% What if we used only the first 12 or 13 values to predict the future peak in the plague?

% Redo using only 12 or 13 values
N = 12;
if length(t)>13
    told = t;
    wold = w;
end
w = wold(1,1:N);
t = told(1,1:N);

figure
plot(t,w,'.','MarkerSize',15)

%

x1 = [A;B;C];   
x = levenberg(@misfit,x1)

a = x(1,end),  b = x(2,end), c=x(3,end)   % final values
model = @(x) a*sech(b*(x-c)).^2;

figure
plot(told,wold,'.','MarkerSize',15) 
hold on
fplot( model, [0 30],'LineWidth',2 )
title('Mumbai Plague Fit')   
ylabel('Deaths Per Week'), xlabel('Time')    
hold off

%
% Below are the functions used.
%

function f = nlsystem2(x)
    f = zeros(3,1);   % ensure a column vector output
    f(1) = exp(x(2)-x(1)) - 2;
    f(2) = x(1)*x(2) + x(3);
    f(3) = x(2)*x(3) + x(1)^2 - x(2);
end

function f = misfit(p)
global t w      
    a = p(1); b = p(2); c=p(3);
    f = norm(a*sech(b*(t-c)).^2 - w); % with norm
    %f = (a*sech(b*(t-c)).^2 - w)';    % w/o norm
end

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