refactor recovery module

references #805

R/net.mod.recovery.R

@@ -16,74 +16,18 @@
 #' @keywords internal
 #'
 recovery.net <- function(dat, at) {
-
-  ## Only run with SIR/SIS
-  type <- get_control(dat, "type", override.null.error = TRUE)
-  type <- if (is.null(type)) "None" else type
-
-  if (!(type %in% c("SIR", "SIS"))) {
-    return(dat)
-  }
-
-  # Variables ---------------------------------------------------------------
-  active <- get_attr(dat, "active")
-  status <- get_attr(dat, "status")
-  infTime <- get_attr(dat, "infTime")
-  recovState <- ifelse(type == "SIR", "r", "s")
-
-  rec.rate <- get_param(dat, "rec.rate")
-
-  nRecov <- 0
-  idsElig <- which(active == 1 & status == "i")
-  nElig <- length(idsElig)
-
-
-  # Time-Varying Recovery Rate ----------------------------------------------
-  infDur <- at - infTime[active == 1 & status == "i"]
-  infDur[infDur == 0] <- 1
-  lrec.rate <- length(rec.rate)
-  if (lrec.rate == 1) {
-    ratesElig <- rec.rate
-  } else {
-    ratesElig <- ifelse(infDur <= lrec.rate, rec.rate[infDur],
-                        rec.rate[lrec.rate])
-  }
-
-
-  # Process -----------------------------------------------------------------
-  if (nElig > 0) {
-    vecRecov <- which(rbinom(nElig, 1, ratesElig) == 1)
-    if (length(vecRecov) > 0) {
-      idsRecov <- idsElig[vecRecov]
-      nRecov <- length(idsRecov)
-      status[idsRecov] <- recovState
-    }
-  }
-  dat <- set_attr(dat, "status", status)
-
-  # Output ------------------------------------------------------------------
-  outName <- ifelse(type == "SIR", "ir.flow", "is.flow")
-  dat <- set_epi(dat, outName, at, nRecov)
-
-  return(dat)
+  recovery_with_ngroups(dat, at, 1)
 }
 
-#' @title Recovery: netsim Module
-#'
-#' @description This function simulates recovery from the infected state
-#'              either to a distinct recovered state (SIR model type) or back
-#'              to a susceptible state (SIS model type), for use in
-#'              \code{\link{netsim}}.
-#'
-#' @inheritParams recovery.net
-#'
-#' @inherit recovery.net return
-#'
+#' @rdname recovery.net
 #' @export
 #' @keywords internal
 #'
 recovery.2g.net <- function(dat, at) {
+  recovery_with_ngroups(dat, at, 2)
+}
 
+recovery_with_ngroups <- function(dat, at, ngroups) {
   ## Only run with SIR/SIS
   type <- get_control(dat, "type", override.null.error = TRUE)
   type <- if (is.null(type)) "None" else type
@@ -96,61 +40,52 @@ recovery.2g.net <- function(dat, at) {
   active <- get_attr(dat, "active")
   status <- get_attr(dat, "status")
   infTime <- get_attr(dat, "infTime")
-  group <- get_attr(dat, "group")
+  if (ngroups > 1) {
+    group <- get_attr(dat, "group")
+  } else {
+    group <- rep(1, length.out = length(active))
+  }
   recovState <- ifelse(type == "SIR", "r", "s")
 
-  rec.rate <- get_param(dat, "rec.rate")
-  rec.rate.g2 <- get_param(dat, "rec.rate.g2")
+  suffixes <- c("", if (ngroups > 1) paste0(".g", seq_len(ngroups)[-1L]))
+
+  rec.rates <- lapply(seq_len(ngroups), function(i) get_param(dat, paste0("rec.rate", suffixes[i])))
 
-  nRecov <- nRecovG2 <- 0
+  nRecov <- integer(ngroups)
   idsElig <- which(active == 1 & status == "i")
   nElig <- length(idsElig)
 
   # Time-Varying Recovery Rate ----------------------------------------------
-  infDur <- at - infTime[active == 1 & status == "i"]
-  infDur[infDur == 0] <- 1
-  lrec.rate <- length(rec.rate)
-  if (lrec.rate == 1) {
-    gElig <- group[idsElig]
-    rates <- c(rec.rate, rec.rate.g2)
-    ratesElig <- rates[gElig]
-  } else {
-    gElig <- group[idsElig]
-    if (is.null(rec.rate.g2)) {
-      rates <- ifelse(infDur <= lrec.rate, rec.rate[infDur],
-                      rec.rate[lrec.rate])
-    } else {
-      rates <- rep(NA, length(infDur))
-      rates[gElig == 1] <- ifelse(infDur[gElig == 1] <= lrec.rate,
-                                  rec.rate[infDur[gElig == 1]],
-                                  rec.rate[lrec.rate])
-      rates[gElig == 2] <- ifelse(infDur[gElig == 2] <= length(rec.rate.g2),
-                                  rec.rate.g2[infDur[gElig == 2]],
-                                  rec.rate.g2[length(rec.rate.g2)])
-    }
-    ratesElig <- rates
-  }
+  infDur <- pmax(1, at - infTime[active == 1 & status == "i"])
 
+  gElig <- group[idsElig]
+  ratesElig <- numeric(length(gElig))
+  for(g in seq_len(ngroups)) {
+    rec.rate <- NVL(rec.rates[[g]], rec.rates[[1]]) # allow rec.rates.g2 to be NULL, for backwards compatibility
+    ratesElig[gElig == g] <- rec.rate[pmin(infDur[gElig == g], length(rec.rate))]
+  }
 
   # Process -----------------------------------------------------------------
   if (nElig > 0) {
     vecRecov <- which(rbinom(nElig, 1, ratesElig) == 1)
     if (length(vecRecov) > 0) {
       idsRecov <- idsElig[vecRecov]
-      nRecov <- sum(group[idsRecov] == 1)
-      nRecovG2 <- sum(group[idsRecov] == 2)
+      nRecov <- tabulate(group[idsRecov], nbins = ngroups)
       status[idsRecov] <- recovState
     }
   }
   dat <- set_attr(dat, "status", status)
 
   # Output ------------------------------------------------------------------
-  outName <- ifelse(type == "SIR", "ir.flow", "is.flow")
-  outName[2] <- paste0(outName, ".g2")
-
-  dat <- set_epi(dat, outName[1], at, nRecov)
-  dat <- set_epi(dat, outName[2], at, nRecovG2)
+  if (type == "SIR") {
+    outName <- "ir.flow"
+  } else {
+    outName <- "is.flow"
+  }
 
+  for (g in seq_len(ngroups)) {
+    dat <- set_epi(dat, paste0(outName, suffixes[g]), at, nRecov[g])
+  }
 
   return(dat)
 }