%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Solution to CSE Your Homework Assignment Project 6    %
%  More Models of Infection: It's Epidemic               %
%  solutionb.m Dianne P. O'Leary   11/03                 %
%                                                        %
%  This program solves Problem 5.                        %
%                                                        %
% We model the spread of an infection through a          %
% population, assuming that a person once infected       %
% cannot be reinfected.                                  %
%                                                        %
% We model two parts of the population:                  %
%    I(t) = proportion of population that is infected    %
%           at time t                                    %
%    S(t) = proportion of population that, at time t,    %
%           has never been infected                      %
% These quantities satisfy 0 <= I(t) <= 1 for t >= 0     %
%                          0 <= S(t) <= 1 for t >= 0     %
% The third part, the proportion of the population that  %
% was once infected but has now recovered, can be        %
% derived from these two:                                %
%                                                        %
%    R(t) = 1 - I(t) - S(t).                             %
% Each of these quantities is also a function of (x,y)   %
% where 0 <= x <= 1 and 0 <= y <= 1.                     %
%                                                        %
% The model is                                           %
%    dI/dt =   mytrans * I * S - I/myrecov               %
%    dS/dt = - mytrans * I * S                           %
%                                                        %
% To account for locality of the epidemic, we introduce  %
% a spacial dependence and diffusion of the infected     %
% population.  This adds a term                          %
%      mymobil * ( (d^2 I)/(dx^2) + (d^2 I)/(dy^2) ) S   %
% to dI/dt, and subtracts the same term from dS/dt.      %
%                                                        %
% We assume that the initial values I(0,x,y) and S(0,x,y)%
% are given, that we study the problem for 0 < x < 1,    %
% 0 < y < 1,  and t >= 0, and that diffusion across the  %
% boundaries x=0, x=1, y=0, y=1 can be neglected so that %
% Dirichlet boundary conditions can be used.             %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

global myN mytrans myrecov mymobil mythreshhold myvac myA mynt myy0

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Set the parameters for the model:                      %
%                                                        %
%   myN = number of grid points                          %
%   mytrans = transmission rate for the infection        %
%   myrecov = # days contagious                          %
%   mymobil = mobility rate for the infected (diffusion) %
%   myinitinfec = proportion of population initially     %
%                 infected                               %
%   mythreshhold = tolerance; when infected or           %
%                  susceptible drop below this value,    %
%                  we conclude that the epidemic ended.  %
%   myvac = vaccination rate                             %
%   mynt = default length of simulation                  %
%   myy0 = initial conditions                            %
%   myA = Laplacian matrix with Neumann boundary cond.   %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

n = 11;  
myN = n*n;
mytrans = .8; 
myrecov = 4;
mymobil = .2; 
mythreshhold = 1.e-5;
myvac = 0;   
mynt = 3000;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Construct the Laplacian matrix (negative diagonal).    %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

t = ones(n,3);
t(:,2) = -2;
t(2,3) = 2;
t(n-1,1) = 2;
scale = (n-1)^2;
T = scale*spdiags(t,[-1,0,1],n,n);
I = speye(n);
myA = kron(T,I) + kron(I,T);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Construct the initial conditions.                      %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

susceptible = ones(myN,1);
infected    = zeros(myN,1);

spot = n*(n-1)/2 + floor(n/2)+1;
infected(spot) = .5;
susceptible(spot) = 1 - infected(spot);

myy0 = [infected;susceptible];


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Problem 5a:  Solve the ODE model.                      %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

strng = sprintf('Solution from Differential Equation Model')
tspan = 0:mynt;
options=odeset('Events',@epiend);
[t,y] = ode23s(@epi,tspan,myy0,options);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Compute I(t) and R(t) by summing over space.           %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

for tmy = 1:length(t)
  suscept =  sum(y(tmy,myN+1:2*myN));
  dinf(tmy) =  sum(y(tmy,1:myN))/myN;
  drec(tmy) =  1 - dinf(tmy) - suscept/myN;
end

figure(1)
plot(t(1:tmy),dinf(1:tmy),t(1:tmy),drec(1:tmy))
axis([0 t(end) 0 1])
legend('Infected','Recovered')
xlabel('time')
ylabel('Proportion')
title(strng)
print -depsc prob5a.eps
sprintf('Proportion of population infected: %f',drec(tmy))

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Problem 5b:  Solve the ODE model with vaccination.     %
% We subtract myvac*(susc.*infe)./(infe+susc) from the   %
% derivative of the susceptibles, and add a vaccinated   %
% segment of the population with derivative              %
% myvac*(susc.*infe)./(infe+susc)                        %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

strng = sprintf('Solution from Differential Equation Model with Vaccination')
myvac = .7;

tspan = 0:mynt;
[t,yv] = ode23s(@epi,tspan,[myy0;zeros(myN,1)],options);

for tmy = 1:length(t)
  suscept =  sum(yv(tmy,myN+1:2*myN));
  dvac(tmy) = sum(yv(tmy,2*myN+1:3*myN))/myN;
  dinf(tmy) =  sum(yv(tmy,1:myN))/myN;
  drec(tmy) =  1 - dinf(tmy) - suscept/myN - dvac(tmy);
end

figure(2)
plot(t(1:tmy),dinf(1:tmy),t(1:tmy),drec(1:tmy),t(1:tmy),dvac(1:tmy))
axis([0 t(end) 0 1])
legend('Infected','Recovered','Vaccinated')
xlabel('time')
ylabel('Proportion')
title(strng)
print -depsc prob5b.eps
sprintf('Proportion of population infected: %f',drec(tmy))