%   
%                  SIR_sensitivities_IC.m
%

  clear all

  global N 

%
%  Compute the sensitivities of the 3 parameter SIR model to the 2 initial conditions
%  with fixed parameters in the differential equation 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: N=2
%

%
% Set parameters and initial conditions
%

  S0 = 900;
  I0 = 100;
  R0 = 0;
  S_S0 = 1;
  I_S0 = 0;
  S_I0 = 0;
  I_I0 = 1;
  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_S0; I_S0; S_I0; I_I0];

%
% 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_IC,t_vals,Y0,ode_options,params);

%
% Extract the states and sensitivities.
%

  S = Y(:,1);  
  I = Y(:,2);
  S_S0_sen = Y(:,3);
  I_S0_sen = Y(:,4);
  S_I0_sen = Y(:,5);
  I_I0_sen = Y(:,6);

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

  clear Y0
  h = 1e-16;

  S0_complex = complex(S0,h);
  Y0 = [S0_complex; I0];
  [t,Y] = ode45(@SIR_rhs_complex,t_vals,Y0,ode_options,params);
  S_S0 = imag(Y(:,1))/h;
  I_S0 = imag(Y(:,2))/h;

  I0_complex = complex(I0,h); 
  Y0 = [S0; I0_complex]; 
  [t,Y] = ode45(@SIR_rhs_complex,t_vals,Y0,ode_options,params);
  S_I0 = imag(Y(:,1))/h;
  I_I0 = imag(Y(:,2))/h;

 
%
% Plot the 4 sensitivities as a function of time.
%

  figure(1)
  plot(t,S_S0_sen,'-b',t,S_S0,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('S_{S_0}')
  legend('Sensitivity Eq','Complex-Step','Location','NorthEast')
%  print -depsc S_S0

  figure(2)
  plot(t,I_S0_sen','-b',t,I_S0,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('I_{S_0}')
  legend('Sensitivity Eq','Complex-Step','Location','NorthEast')
%  print -depsc I_S0

  figure(3)
  plot(t,S_I0_sen','-b',t,S_I0,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('S_{I_0}')
  legend('Sensitivity Eq','Complex-Step','Location','SouthEast')
%  print -depsc S_I0

  figure(4)
  plot(t,I_I0_sen,'-b',t,I_I0,'--r','linewidth',3)
  set(gca,'Fontsize',[22]);
  xlabel('Time (Days)')
  ylabel('I_{I_0}')
  legend('Sensitivity Eq','Complex-Step','Location','NorthEast')
%  print -depsc I_I0