%   
%                  SIR_sensitivity.m
%

  clear all

  global N 

%
%  Compute the sensitivities of the 3 parameter SIR model using the sensitivity
%  equations and complex-step derivative approximations.  Reduce to a system of
%  two differential equations using the relation R(t) = N - S(t) - I(t).
%  Note: p=3, N=2
%

%
% Set parameters and initial conditions
%

  S0 = 900;
  I0 = 100;
  R0 = 0;
  S_mu0 = 0;
  I_mu0 = 0;
  S_beta0 = 0;
  I_beta0 = 0;
  S_gamma0 = 0;
  I_gamma0 = 0;
  N = 1000;

  beta = 0.005;
  gamma = 0.3;
  mu = 0.1;
  params = [beta gamma mu];

  tf = 15;
  dt = 0.1;
  t_vals = 0:dt:tf;
  Y0 = [S0; I0; S_mu0; I_mu0; S_beta0; I_beta0; S_gamma0; I_gamma0];

%
% Integrate the coupled system using ode45.m.  There are now N(p+1) coupled
% differential equations comprised of the state and sensitivity equations
%

  ode_options = odeset('RelTol',1e-8);
  [t,Y] = ode45(@SIR_rhs,t_vals,Y0,ode_options,params);

%
% Extract the states and sensitivities.
%

  S = Y(:,1);  
  I = Y(:,2);
  S_mu_sen = Y(:,3);
  I_mu_sen = Y(:,4);
  S_beta_sen = Y(:,5);
  I_beta_sen = Y(:,6);
  S_gamma_sen = Y(:,7);
  I_gamma_sen = Y(:,8); 

%
% Compute the sensitivities using complex-step derivative approximations.  This
% requires p=3 solutions of the ODE system.
%

  clear Y0
  h = 1e-16;
  Y0 = [S0; I0];

  mu_complex = complex(mu,h);
  params = [beta gamma mu_complex];
  [t,Y] = ode45(@SIR_rhs_complex,t_vals,Y0,ode_options,params);
  S_mu = imag(Y(:,1))/h;
  I_mu = imag(Y(:,2))/h;

  beta_complex = complex(beta,h);
  params = [beta_complex gamma mu];
  [t,Y] = ode45(@SIR_rhs_complex,t_vals,Y0,ode_options,params);
  S_beta = imag(Y(:,1))/h;
  I_beta = imag(Y(:,2))/h;

  gamma_complex = complex(gamma,h);
  params = [beta gamma_complex mu];
  [t,Y] = ode45(@SIR_rhs_complex,t_vals,Y0,ode_options,params);
  S_gamma = imag(Y(:,1))/h;
  I_gamma = imag(Y(:,2))/h;

%
% Compute the relative errors.
% 

  S_mu_err = max(abs(S_mu_sen-S_mu))/max(abs(S_mu_sen));
  I_mu_err = max(abs(I_mu_sen-I_mu))/max(abs(I_mu_sen));
  S_beta_err = max(abs(S_beta_sen-S_beta))/max(abs(S_beta_sen));
  I_beta_err = max(abs(I_beta_sen-I_beta))/max(abs(I_beta_sen));
  S_gamma_err = max(abs(S_gamma_sen-S_gamma))/max(abs(S_gamma_sen));
  I_gamma_err = max(abs(I_gamma_sen-I_gamma))/max(abs(I_gamma_sen));
 
%
% Plot the 6 sensitivities as a function of time.
%

  figure(1)
  plot(t,I,'linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('Number of Infectious')

  figure(3)
  plot(t,S_mu_sen,'-b',t,S_mu,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('S_\mu')
  legend('Sensitivity Eq','Complex-Step','Location','SouthEast')
%  print -depsc S_mu

  figure(4)
  plot(t,I_mu_sen','-b',t,I_mu,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('I_\mu')
  legend('Sensitivity Eq','Complex-Step','Location','SouthEast')
%  print -depsc I_mu

  figure(5)
  plot(t,S_beta_sen','-b',t,S_beta,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('S_\beta')
  legend('Sensitivity Eq','Complex-Step','Location','SouthEast')
%  print -depsc S_beta

  figure(6)
  plot(t,I_beta_sen,'-b',t,I_beta,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('I_\beta')
  legend('Sensitivity Eq','Complex-Step','Location','NorthEast')
%  print -depsc I_beta

  figure(7)
  plot(t,S_gamma_sen,'-b',t,S_gamma,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('S_\gamma')
  legend('Sensitivity Eq','Complex-Step','Location','SouthEast')
%  print -depsc S_gamma

  figure(8)
  plot(t,I_gamma_sen,'-b',t,I_gamma,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
%  ylabel('$\displaystyle\frac{\partial I}{\partial gamma}$','interpreter','latex')
  ylabel('I_\gamma')
  legend('Sensitivity Eq','Complex-Step','Location','NorthEast')
%  print -depsc I_gamma