% SIWR Main

clear

%% Load Angola Cholera Data

CholeraData;

%% Setup

% A note---I often find it's convenient to make small functions that take
% your parameters/data/model output and return (for example) the initial 
% conditions or measurement equation for your model. 
% This is because the initial conditions and the measurement equation both 
% depend on the parameter k. So we'll make two functions:
%
%   x0fcn: a function that takes in the current parameter values and
%   returns the initial conditions, i.e. I(0) = data(1)*k, S(0) = 1- I(0),
%   R(0) = 0. We'll assume the parameter order is: betaI, betaW, xi, k.

x0fcn = @(params) [1-params(end)*data(1); params(end)*data(1); 0; 0];

%   yfcn: a function that takes in the current model output and parameter
%   values, and returns the measured cases, y = I/k. Note that k =
%   params(end), and I = x(:,2).

yfcn = @(x,params) x(:,2)/params(end);

% Note the "params" and "x" here are just dummy variables--they will be
% filled once we call these functions! However, "data(1)" is actually
% referencing the CholeraData we just loaded. Be careful--if you change the
% data later, this will not be reflected unless you reload the x0 function!
% You can change this if you want, e.g. to make x0fcn = @(params,data) etc.


% Starting parameter values:
params = [0.75; 0.75; 0.01; 1/89193];


%% Problem 1: Simulate and Plot the Model

[t,x] = ode45(@siwrODE,times,x0fcn(params),[],params);
y = yfcn(x,params);

figure(1)
    set(gca,'LineWidth',1,'FontSize',16,'FontName','Arial')
    hold on
    plot(t,y,'b','LineWidth',2)
    plot(times,data,'ko','LineWidth',2)
    ylabel('Infected Population') 
    xlabel('Time (days)')
    title('Problem 1: Model Simulation')

%% Problem 2: Parameter Estimation (OLS + Nelder Mead)

% Estimate the model parameters 
[paramests, fval] = fminsearch(@(p) siwrCost(times,p,data,x0fcn,yfcn),params,optimset('Display','iter','MaxFunEvals',5000,'MaxIter',5000));
paramests = abs(paramests);

%In Matlab2014b, this yields: paramests = [0.293567334915841;3.52977516648091;0.00195048360857802;1.07563814751018e-05]

% Re-simulate the model with the final parameter estimates
[test,xest] = ode45(@siwrODE,times,x0fcn(paramests),[],paramests);
yest = yfcn(xest,paramests);

% Plot results
figure(2)  % Model Fit
    set(gca,'LineWidth',1,'FontSize',16,'FontName','Arial')
    hold on
    plot(test,yest,'k','LineWidth',2)
    plot(times,data,'ko','LineWidth',2)
    ylabel('Infected Population')
    xlabel('Time (days)')
    title('Problem 2: Parameter Estimation (Model Fit to Data)')

    
figure(3)  %Residuals of the fit
    set(gca,'LineWidth',1,'FontSize',16,'FontName','Arial')
    hold on
    plot(test,(yest-data),'o')
    ylabel('Residuals')
    xlabel('Time (days)')
    title('Problem 2: Parameter Estimation (Residuals)')
    
    
    
%% Problem 3: Exploring betaW & xi

% Takehome message: the parameters betaW and xi appear to be able to
% compensate for one another such that approximately the same output is
% generated for different values of betaW & xi. You can try this for other
% numbers too, e.g. 2/3 and 3/2, etc. and mostly get similar results. 

paramestsalt = paramests.*[1; 2/3; 3/2; 1;]; %paramests.*[1; 5/6; 6/5; 1;];
[talt,xalt] = ode45(@siwrODE,times,x0fcn(paramestsalt),[],paramestsalt);
yalt = yfcn(xalt,paramestsalt);


figure(4)
    set(gca,'LineWidth',1,'FontSize',16,'FontName','Arial')
    hold on
    plot(test,yest,'k','LineWidth',4)
    plot(talt,yalt,'r','LineWidth',2)
    plot(times,data,'ko','LineWidth',2)
    ylabel('Infected Population')  
    xlabel('Time (days)')
    title('Problem 3: Tradeoff Between \beta_W and \xi')

    
    
%% Problem 4: Fisher Information Matrix

% Note: you can also do this analytically using the sensitivity equations,
% as you saw in yesterday's lecture/practical! Feel free to use and modify
% that code instead if you prefer.

FIM = MiniFisher(times,paramests,x0fcn,@siwrODE,yfcn);

%Calculate & print rank of FIM
rank(FIM)

%Calculate SEs & CVs, print CVs
SEs = sqrt(diag(((length(data) - length(paramests))/fval)*inv(FIM)));
CVs = SEs./paramests  
% the CV's normalize the SE's so that you can more easily compare across
% parameters with very different values. (e.g. an SE of 0.01 might mean
% something very different for betaW than for k).

% Takehome message: in Matlab2014b, this gives a rank deficient matrix
% (rank = 3), but is numerically invertible anyhow. The CV's for betaI and
% k are both a couple orders of magnitude smaller than those for betaW and
% xi, further underscoring the idea that these two might be involved in a
% practically identifiable combination.


%% Problem 5: Refit the model to various noisy data sets

% You might try changing this so you fit the model from different starting 
% values too! (for example, see commented line below)

noiseparams = [];
noisefvals = [];
startparams = paramests;
paramlist = {'\beta_I','\beta_W','\xi','k'};
    
for i=1:30
    i %track progress
    
    % startparams = normrnd(paramests,0.4*paramests);
    
    data = poissrnd(yest);
    % data = normrnd(yest,0.1*yest);
    
    % Note we need to update the x0fcn since now the data is changing each
    % time! We could also make x0fcn a function of the data too.
    x0fcn = @(params) [1-params(end)*data(1); params(end)*data(1); 0; 0;];
    
    % Estimate parameters with the new data
    [currparams, currfval] = fminsearch(@(p) siwrCost(times,p,data,x0fcn,yfcn),startparams,optimset('MaxFunEvals',5000,'MaxIter',5000));
    
    % Keep track of the estimates
    noiseparams = [noiseparams abs(currparams)];
    noisefvals = [noisefvals currfval];
end


figure(6)
pairsplot_2(noiseparams',paramlist,paramests);


% Takehome message: Here, note the relationship between betaW and xi---all the other
% parameters seem not to have much in the way of relationships with each
% other. (This can be seen as the central clumped dots in the middle, as
% well as the horizontal and vertical lines--the lines show that even when
% the parameters vary, there's no relationship between them.) It may be
% good to zoom in on the betaW/xi plots so you can see the curve they form.
    



