%   
%                  SIR_sensitivities.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_delta0 = 0;
  I_delta0 = 0;
  S_gamma0 = 0;
  I_gamma0 = 0;
  S_r0 = 0;
  I_r0 = 0;
  N = 1000;

  gamma = 0.01;
  r = 0.3;
  delta = 0.35;
  params = [gamma r delta];

  tf = 6;
  dt = 0.1;
  t_vals = 0:dt:tf;
  Y0 = [S0; I0; S_delta0; I_delta0; S_gamma0; I_gamma0; S_r0; I_r0];

%
% 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_delta_sen = Y(:,3);
  I_delta_sen = Y(:,4);
  S_gamma_sen = Y(:,5);
  I_gamma_sen = Y(:,6);
  S_r_sen = Y(:,7);
  I_r_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];

  delta_complex = complex(delta,h);
  params = [gamma r delta_complex];
  [t,Y] = ode45(@SIR_rhs_complex,t_vals,Y0,ode_options,params);
  S_delta = imag(Y(:,1))/h;
  I_delta = imag(Y(:,2))/h;

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

  r_complex = complex(r,h);
  params = [gamma r_complex delta];
  [t,Y] = ode45(@SIR_rhs_complex,t_vals,Y0,ode_options,params);
  S_r = imag(Y(:,1))/h;
  I_r = imag(Y(:,2))/h;

%
% Compute the relative errors.
% 

  S_delta_err = max(abs(S_delta_sen-S_delta))/max(abs(S_delta_sen));
  I_delta_err = max(abs(I_delta_sen-I_delta))/max(abs(I_delta_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));
  S_r_err = max(abs(S_r_sen-S_r))/max(abs(S_r_sen));
  I_r_err = max(abs(I_r_sen-I_r))/max(abs(I_r_sen));
 
%
% Plot the 6 sensitivities as a function of time.
%

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

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

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

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

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

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

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