Updating virus methods

UofUEpiBio · Aug 28, 2023 · aa0ada2 · aa0ada2
1 parent ac36976
commit aa0ada2
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 6 deletions.
diff --git a/NAMESPACE b/NAMESPACE
@@ -6,8 +6,12 @@ S3method(add_tool_n,epiworld_model)
 S3method(add_virus,epiworld_model)
 S3method(add_virus,epiworld_seir)
 S3method(add_virus,epiworld_seirconn)
+S3method(add_virus,epiworld_seird)
+S3method(add_virus,epiworld_seirdconn)
 S3method(add_virus,epiworld_sir)
 S3method(add_virus,epiworld_sirconn)
+S3method(add_virus,epiworld_sird)
+S3method(add_virus,epiworld_sirdconn)
 S3method(add_virus_n,epiworld_model)
 S3method(add_virus_n,epiworld_seir)
 S3method(add_virus_n,epiworld_seirconn)

diff --git a/R/ModelSEIRDCONN.R b/R/ModelSEIRDCONN.R
@@ -41,8 +41,8 @@
 #' plot(model_seirdconn)
 #' 
 #' # Adding the flu
-#' flu <- virus("Flu", .9, 1/7)
-#' add_virus(model_seirdconn, flu, .001)
+#' flu <- virus("Flu", prob_infecting = .3, recovery_rate = 1/7, prob_death = 0.001)
+#' add_virus(model = model_seirdconn, virus = flu, proportion = .001)
 #' 
 #' #' # Running and printing
 #' run(model_seirdconn, ndays = 100, seed = 1912)

diff --git a/R/virus.R b/R/virus.R
@@ -160,6 +160,15 @@ add_virus.epiworld_sir <- function(model, virus, proportion) {
 
 }
 
+#' @export
+add_virus.epiworld_sird <- function(model, virus, proportion) {
+
+  stopifnot_virus(virus)
+  virus_set_state(virus, init = 1, end = 2, removed = 3)
+  invisible(add_virus_cpp(model, virus, proportion))
+
+}
+
 #' @export
 add_virus.epiworld_sirconn <- function(model, virus, proportion) {
 
@@ -168,6 +177,15 @@ add_virus.epiworld_sirconn <- function(model, virus, proportion) {
 
 }
 
+#' @export
+add_virus.epiworld_sirdconn <- function(model, virus, proportion) {
+
+  stopifnot_virus(virus)
+  add_virus.epiworld_sird(model, virus, proportion)
+
+}
+
+
 #' @export
 add_virus.epiworld_seir <- function(model, virus, proportion) {
 
@@ -177,6 +195,15 @@ add_virus.epiworld_seir <- function(model, virus, proportion) {
 
 }
 
+#' @export
+add_virus.epiworld_seird <- function(model, virus, proportion) {
+
+  stopifnot_virus(virus)
+  virus_set_state(virus, init = 1, end = 3, removed = 4)
+  invisible(add_virus_cpp(model, virus, proportion))
+
+}
+
 #' @export
 add_virus.epiworld_seirconn <- function(model, virus, proportion) {
 
@@ -185,6 +212,14 @@ add_virus.epiworld_seirconn <- function(model, virus, proportion) {
 
 }
 
+#' @export
+add_virus.epiworld_seirdconn <- function(model, virus, proportion) {
+
+  stopifnot_virus(virus)
+  add_virus.epiworld_seird(model, virus, proportion)
+
+}
+
 #' @export
 #' @rdname virus
 #' @returns 

diff --git a/man/ModelSEIRD.Rd b/man/ModelSEIRD.Rd
diff --git a/man/ModelSEIRDCONN.Rd b/man/ModelSEIRDCONN.Rd