refactor prevalence module

references #805
EpiModel · Mar 7, 2023 · a52469f · a52469f
1 parent 3b7186f
commit a52469f
Showing 1 changed file with 24 additions and 113 deletions.
diff --git a/R/net.mod.prevalence.R b/R/net.mod.prevalence.R
@@ -24,140 +24,51 @@ prevalence.net <- function(dat, at) {
   type <- if (is.null(type)) "None" else type
 
   groups <- get_param(dat, "groups")
+  if (groups > 1) {
+    group <- get_attr(dat, "group")[get_attr(dat, "active") == 1]
+  } else {
+    group <- rep(1, length.out = sum(active == 1))
+  }
 
-  # Subset attr to active == 1
-  l <- lapply(seq_along(dat$attr), function(x) dat$attr[[x]][active == 1])
-  names(l) <- names(dat$attr)
-  l$active <- l$infTime <- NULL
+  suffixes <- c("", if (groups > 1) paste0(".g", seq_len(groups)[-1]))
+  statuses <- c("s", "i", if (type == "SIR") "r")
 
-  status <- l$status
+  status <- get_attr(dat, "status")[get_attr(dat, "active") == 1]
 
   ## Subsetting for epi.by control
   eb <- !is.null(dat$control$epi.by)
   if (eb == TRUE) {
     ebn <- get_control(dat, "epi.by")
     ebv <- dat$temp$epi.by.vals
     ebun <- paste0(".", ebn, ebv)
-    assign(ebn, l[[ebn]])
-  }
-
-  if (groups == 1) {
-    dat <- set_epi(dat, "s.num", at, sum(status == "s"))
-    if (eb == TRUE) {
-      for (i in seq_along(ebun)) {
-        ebn.temp <- paste0("s.num", ebun[i])
-        dat <- set_epi(dat, ebn.temp, at, sum(status == "s" &
-          get(ebn) == ebv[i]))
-      }
-    }
-
-    dat <- set_epi(dat, "i.num", at, sum(status == "i"))
-    if (eb == TRUE) {
-      for (i in seq_along(ebun)) {
-        ebn.temp <- paste0("i.num", ebun[i])
-        dat <- set_epi(dat, ebn.temp, at, sum(status == "i" &
-          get(ebn) == ebv[i]))
-      }
-    }
-
-    if (type == "SIR") {
-      dat <- set_epi(dat, "r.num", at, sum(status == "r"))
-      if (eb == TRUE) {
-        for (i in seq_along(ebun)) {
-          ebn.temp <- paste0("r.num", ebun[i])
-          dat <- set_epi(dat, ebn.temp, at, sum(status == "r" &
-            get(ebn) == ebv[i]))
-        }
-      }
-    }
-    dat <- set_epi(dat, "num", at, length(status))
-    if (eb == TRUE) {
-      for (i in seq_along(ebun)) {
-        ebn.temp <- paste0("num", ebun[i])
-        dat <- set_epi(dat, ebn.temp, at, sum(get(ebn) == ebv[i]))
-      }
+    if (ebn %in% c("active", "infTime")) {
+      ebvec <- NULL
+    } else {
+      ebvec <- get_attr(dat, ebn)[get_attr(dat, "active") == 1]
     }
   }
 
-  if (groups == 2) {
-    group <- get_attr(dat, "group")
-    group <- group[active == 1]
-
-    dat <- set_epi(dat, "s.num", at, sum(status == "s" & group == 1))
-    if (eb == TRUE) {
-      for (i in seq_along(ebun)) {
-        ebn.temp <- paste0("s.num", ebun[i])
-        dat <- set_epi(dat, ebn.temp, at, sum(status == "s" &
-          group == 1 &
-          get(ebn) == ebv[i]))
-      }
-    }
-    dat <- set_epi(dat, "i.num", at, sum(status == "i" & group == 1))
-    if (eb == TRUE) {
-      for (i in seq_along(ebun)) {
-        ebn.temp <- paste0("i.num", ebun[i])
-        dat <- set_epi(dat, ebn.temp, at, sum(status == "i" &
-          group == 1 &
-          get(ebn) == ebv[i]))
-      }
-    }
-    if (type == "SIR") {
-      dat <- set_epi(dat, "r.num", at, sum(status == "r" & group == 1))
-      if (eb == TRUE) {
-        for (i in seq_along(ebun)) {
-          ebn.temp <- paste0("r.num", ebun[i])
-          dat <- set_epi(dat, ebn.temp, at, sum(status == "r" &
-            group == 1 &
-            get(ebn) == ebv[i]))
-        }
-      }
-    }
-    dat <- set_epi(dat, "num", at, sum(group == 1))
-    dat$epi$num[at] <- sum(group == 1)
-    if (eb == TRUE) {
-      for (i in seq_along(ebun)) {
-        ebn.temp <- paste0("num", ebun[i])
-        dat <- set_epi(dat, ebn.temp, at, sum(group == 1 &
-          get(ebn) == ebv[i]))
-      }
-    }
-    dat <- set_epi(dat, "s.num.g2", at, sum(status == "s" & group == 2))
-    if (eb == TRUE) {
-      for (i in seq_along(ebun)) {
-        ebn.temp <- paste0("s.num.g2", ebun[i])
-        dat <- set_epi(dat, ebn.temp, at, sum(status == "s" &
-          group == 2 &
-          get(ebn) == ebv[i]))
-      }
-    }
-    dat <- set_epi(dat, "i.num.g2", at, sum(status == "i" & group == 2))
-    if (eb == TRUE) {
-      for (i in seq_along(ebun)) {
-        ebn.temp <- paste0("i.num.g2", ebun[i])
-        dat <- set_epi(dat, ebn.temp, at, sum(status == "i" &
-          group == 2 &
-          get(ebn) == ebv[i]))
-      }
-    }
-    if (type == "SIR") {
-      dat <- set_epi(dat, "r.num.g2", at, sum(status == "r" & group == 2))
+  for (g in seq_len(groups)) {
+    for (status_value in statuses) {
+      dat <- set_epi(dat, paste0(status_value, ".num", suffixes[g]), at,
+                     sum(status == status_value & group == g))
       if (eb == TRUE) {
         for (i in seq_along(ebun)) {
-          ebn.temp <- paste0("r.num.g2", ebun[i])
-          dat <- set_epi(dat, ebn.temp, at, sum(status == "r" &
-            group == 2 &
-            get(ebn) == ebv[i]))
+          ebn.temp <- paste0(status_value, ".num", suffixes[g], ebun[i])
+          dat <- set_epi(dat, ebn.temp, at, sum(status == status_value &
+            group == g & ebvec == ebv[i]))
         }
       }
     }
-    dat <- set_epi(dat, "num.g2", at, sum(group == 2))
+
+    dat <- set_epi(dat, paste0("num", suffixes[g]), at, sum(group == g))
     if (eb == TRUE) {
       for (i in seq_along(ebun)) {
-        ebn.temp <- paste0("num.g2", ebun[i])
-        dat <- set_epi(dat, ebn.temp, at, sum(group == 2 &
-          get(ebn) == ebv[i]))
+        ebn.temp <- paste0("num", suffixes[g], ebun[i])
+        dat <- set_epi(dat, ebn.temp, at, sum(ebvec == ebv[i] & group == g))
       }
     }
   }
+
   return(dat)
 }