SIR Model

Schrick-Noah_Homework-3.R

@@ -233,9 +233,9 @@ my_euler2 <- function(f, y0, t0, tfinal, dt){
   for (i in 2:npts){ # rows are time
     for (j in 1:nvars){
       y_vec_prev <- y[i-1,] # both y values at previous time point
+      dy <- f(t[i-1],y_vec_prev)*dt # f returns a vector, t not used
+      y[i,] <- y_vec_prev + dy
     }
-    dy <- f(t[i-1],y_vec_prev)*dt # f returns a vector, t not used
-    y[i,] <- y_vec_prev + dy
   }
   return(list(t=tspan,y=y))
 }
@@ -250,12 +250,14 @@ my_rk4_2 <- function(f, y0, t0, tfinal, h){
   for (i in 2:npts){ # rows are time
     for (j in 1:nvars){
       y_vec_prev <- y[i-1,] # both y values at previous time point
+    
+      k1 <- f(t[i-1], y_vec_prev[i-1])
+      k2 <- f(t[i-1] + 0.5 * h, y_vec_prev[i-1] + 0.5 * k1 * h)
+      k3 <- f(t[i-1] + 0.5 * h, y_vec_prev[i-1] + 0.5 * k2 * h)
+      k4 <- f(t[i-1] + h, y_vec_prev[i-1] + k3*h)
+      
+      y[i,] <- y_vec_prev[i-1] + (h/6)*(k1 + 2*k2 + 2*k3 + k4)
     }
-    k1 <- f(t[i-1], y_vec_prev[i-1])
-    k2 <- f(t[i-1] + 0.5 * h, y_vec_prev[i-1] + 0.5 * k1 * h)
-    k3 <- f(t[i-1] + 0.5 * h, y_vec_prev[i-1] + 0.5 * k2 * h)
-    k4 <- f(t[i-1] + h, y_vec_prev[i-1] + k3*h)
-    y[i,] <- y_vec_prev[i-1] + (h/6)*(k1 + 2*k2 + 2*k3 + k4)
   }
   return(list(t=tspan,y=y))
 }
@@ -263,7 +265,6 @@ my_rk4_2 <- function(f, y0, t0, tfinal, h){
 pp.euler.sol <- my_euler2(f = function(t,y){pp.f(t,y,k=pp.params)},
                          y = c(prey0, pred0),
                          t0 = tmin, tfinal = tmax, dt=h)
-
 plot.pred.prey(pp.euler.sol, "Euler")
 
 pp.rk4.sol <- my_rk4_2(f = function(t,y){pp.f(t,y,k=pp.params)},
@@ -282,12 +283,50 @@ pp.sol.c <- ode45(f = function(t,y){pp.f(t,y,k=pp.params)},
 plot.pred.prey(pp.sol.c, "ODE45 with k3=0.02")
 
 ## 5. Solve SIR model numerically from t=0 -> 20
+sir.f <- function(t, y, k) {
+  dS <- -k[1] * y[1] * y[2]
+  dI <-  k[1] * y[1] * y[2] - k[2] * y[2]
+  dR <-  k[2] * y[2]
+  return(as.matrix(c(dS, dI, dR)))
+}
 
 # a) a=0.5, b=1, S(0)=0.9, I(0)=0.1, R(0)=0
+sir.params <- c(0.5, 1)
+S0 <- 0.9
+I0 <- 0.1
+R0 <- 0.0
+tmin <- 0
+tmax <- 20
+
+sir.ode.sol <- ode45(f = function(t,y){sir.f(t,y,k=sir.params)},
+                y = c(S0, I0, R0),
+                t0 = tmin, tfinal = tmax)
 
 # plot
+plot.sir <- function(sol, method){
+  # - pred/prey columns melted to 1 column called "value" 
+  sol.melted <- melt(data.frame(time=sol$t,
+                                   S=sol$y[,1],
+                                   I=sol$y[,2]), 
+                        id = "time") # melt based on time
+  # New column created with variable names, called “variable”
+  colnames(sol.melted)[2] <- "Group" # used in legend  
+  
+  g <- ggplot(data = sol.melted, 
+              aes(x = time, y = value, color = Group))
+  g <- g + geom_point(size=2) 
+  g <- g + xlab("time") +  ylab("Population") 
+  g <- g + ggtitle(paste("SIR Model Using", method))  
+  show(g)
+}
 
+plot.sir(sir.ode.sol, "ODE45")
 # b) a=3
+sir.params.b <- c(3, 1)
+sir.ode.sol.b <- ode45(f = function(t,y){sir.f(t,y,k=sir.params.b)},
+                     y = c(S0, I0, R0),
+                     t0 = tmin, tfinal = tmax)
+plot.sir(sir.ode.sol.b, "ODE45 with a=3")
 
 ## 6. Decomp of dinitrogen pentoxygen into nitrogen dioxide and molecular oxygen