Skip to content

Commit

Permalink
refactor arrival and departure modules
Browse files Browse the repository at this point in the history 
references #805
  • Loading branch information
@chad-klumb
chad-klumb committed Mar 8, 2023
1 parent 226b100 commit 06fbe4d
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 217 deletions.
281 changes: 64 additions & 217 deletions R/net.mod.vital.R
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,18 @@
#' @keywords netMod internal
#'
departures.net <- function(dat, at) {
departures_with_ngroups(dat, at, 1)
}

#' @rdname departures.net
#' @export
#' @keywords netMod internal
#'
departures.2g.net <- function(dat, at) {
departures_with_ngroups(dat, at, 2)
}

departures_with_ngroups <- function(dat, at, ngroups) {
# Conditions --------------------------------------------------------------
vital <- get_param(dat, "vital")
if (vital == FALSE) {
Expand All @@ -26,66 +37,41 @@ departures.net <- function(dat, at) {
active <- get_attr(dat, "active")
status <- get_attr(dat, "status")
exitTime <- get_attr(dat, "exitTime")
rates.sus <- get_param(dat, "ds.rate")
rates.inf <- get_param(dat, "di.rate")
if (type == "SIR") {
rates.rec <- get_param(dat, "dr.rate")
}

# Susceptible departures --------------------------------------------------
nDepartures.sus <- 0
idsElig.sus <- which(active == 1 & status == "s")
nElig.sus <- length(idsElig.sus)
if (nElig.sus > 0) {
vecDepartures.sus <- which(rbinom(nElig.sus, 1, rates.sus) == 1)
if (length(vecDepartures.sus) > 0) {
idsDpt.sus <- idsElig.sus[vecDepartures.sus]
nDepartures.sus <- length(idsDpt.sus)
active[idsDpt.sus] <- 0
exitTime[idsDpt.sus] <- at
}
}

# Infected departures -----------------------------------------------------
nDepartures.inf <- 0
idsElig.inf <- which(active == 1 & status == "i")
nElig.inf <- length(idsElig.inf)
if (nElig.inf > 0) {
vecDepartures.inf <- which(rbinom(nElig.inf, 1, rates.inf) == 1)
if (length(vecDepartures.inf) > 0) {
idsDpt.inf <- idsElig.inf[vecDepartures.inf]
nDepartures.inf <- length(idsDpt.inf)
active[idsDpt.inf] <- 0
exitTime[idsDpt.inf] <- at
}
}

# Recovered departures ----------------------------------------------------
if (type == "SIR") {
nDepartures.rec <- 0
idsElig.rec <- which(active == 1 & status == "r")
nElig.rec <- length(idsElig.rec)
if (nElig.rec > 0) {
vecDepartures.rec <- which(rbinom(nElig.rec, 1, rates.rec) == 1)
if (length(vecDepartures.rec) > 0) {
idsDpt.rec <- idsElig.rec[vecDepartures.rec]
nDepartures.rec <- length(idsDpt.rec)
active[idsDpt.rec] <- 0
exitTime[idsDpt.rec] <- at
statuses <- c("s", "i", if (type == "SIR") "r")
suffixes <- c("", if (ngroups > 1) paste0(".g", seq_len(ngroups)[-1L]))
if (ngroups > 1) {
group <- get_attr(dat, "group")
} else {
group <- rep(1, length.out = length(active))
}

for (status_value in statuses) {
gDepartures <- integer(ngroups)
idsElig <- which(active == 1 & status == status_value)
nElig <- length(idsElig)
if (nElig > 0) {
rates <- unlist(get_param_list(dat, paste0("d", status_value, ".rate", suffixes)))
gElig <- group[idsElig]
gRates <- rates[gElig]
vecDepartures <- which(rbinom(nElig, 1, gRates) == 1)
if (length(vecDepartures) > 0) {
idsDpt <- idsElig[vecDepartures]
gDepartures <- tabulate(group[idsDpt], nbins = ngroups)
active[idsDpt] <- 0
exitTime[idsDpt] <- at
}
}
for (i in seq_along(suffixes)) {
dat <- set_epi(dat, paste0("d", status_value, ".flow", suffixes[i]), at, gDepartures[i])
}
}

# Output ------------------------------------------------------------------

dat <- set_attr(dat, "active", active)
dat <- set_attr(dat, "exitTime", exitTime)

dat <- set_epi(dat, "ds.flow", at, nDepartures.sus)
dat <- set_epi(dat, "di.flow", at, nDepartures.inf)
if (type == "SIR") {
dat <- set_epi(dat, "dr.flow", at, nDepartures.rec)
}
return(dat)
}

Expand All @@ -105,200 +91,61 @@ departures.net <- function(dat, at) {
#' @keywords netMod internal
#'
arrivals.net <- function(dat, at) {

# Conditions --------------------------------------------------------------
vital <- get_param(dat, "vital")
if (vital == FALSE) {
return(dat)
}

# Variables ---------------------------------------------------------------
a.rate <- get_param(dat, "a.rate")
index <- at - 1
nOld <- get_epi(dat, "num", index)
nArrivals <- 0

# Add Nodes ---------------------------------------------------------------
if (nOld > 0) {
nArrivals <- rbinom(1, nOld, a.rate)
if (nArrivals > 0) {
dat <- append_core_attr(dat, at, nArrivals)
dat <- append_attr(dat, "status", "s", nArrivals)
dat <- append_attr(dat, "infTime", NA, nArrivals)
}
}

# Output ------------------------------------------------------------------
dat <- set_epi(dat, "a.flow", at, nArrivals)

return(dat)
}


#' @title Departures: netsim Module
#'
#' @description This function simulates departure for use in \link{netsim}
#' simulations.
#'
#' @inheritParams recovery.net
#'
#' @inherit recovery.net return
#'
#' @seealso \code{\link{netsim}}
#'
#' @export
#' @keywords netMod internal
#'
departures.2g.net <- function(dat, at) {

# Conditions --------------------------------------------------------------
vital <- get_param(dat, "vital")
if (vital == FALSE) {
return(dat)
}

# Variables ---------------------------------------------------------------

type <- get_control(dat, "type", override.null.error = TRUE)
type <- if (is.null(type)) "None" else type


active <- get_attr(dat, "active")
status <- get_attr(dat, "status")
exitTime <- get_attr(dat, "exitTime")
group <- get_attr(dat, "group")
rates.sus <- c(get_param(dat, "ds.rate"), get_param(dat, "ds.rate.g2"))
rates.inf <- c(get_param(dat, "di.rate"), get_param(dat, "di.rate.g2"))
if (type == "SIR") {
rates.rec <- c(get_param(dat, "dr.rate"), get_param(dat, "dr.rate.g2"))
}

# Susceptible departures --------------------------------------------------
nDepartures.sus <- nDeparturesG2.sus <- 0
idsElig.sus <- which(active == 1 & status == "s")
nElig.sus <- length(idsElig.sus)
if (nElig.sus > 0) {
gElig.sus <- group[idsElig.sus]
ratesElig.sus <- rates.sus[gElig.sus]
vecDepartures.sus <- which(rbinom(nElig.sus, 1, ratesElig.sus) == 1)
if (length(vecDepartures.sus) > 0) {
idsDpt.sus <- idsElig.sus[vecDepartures.sus]
nDepartures.sus <- sum(group[idsDpt.sus] == 1)
nDeparturesG2.sus <- sum(group[idsDpt.sus] == 2)
active[idsDpt.sus] <- 0
exitTime[idsDpt.sus] <- at
}
}

# Infected departures -----------------------------------------------------
nDepartures.inf <- nDeparturesG2.inf <- 0
idsElig.inf <- which(active == 1 & status == "i")
nElig.inf <- length(idsElig.inf)
if (nElig.inf > 0) {
gElig.inf <- group[idsElig.inf]
ratesElig.inf <- rates.inf[gElig.inf]
vecDepartures.inf <- which(rbinom(nElig.inf, 1, ratesElig.inf) == 1)
if (length(vecDepartures.inf) > 0) {
idsDpt.inf <- idsElig.inf[vecDepartures.inf]
nDepartures.inf <- sum(group[idsDpt.inf] == 1)
nDeparturesG2.inf <- sum(group[idsDpt.inf] == 2)
active[idsDpt.inf] <- 0
exitTime[idsDpt.inf] <- at
}
}

# Recovered departures ----------------------------------------------------
if (type == "SIR") {
nDepartures.rec <- nDeparturesG2.rec <- 0
idsElig.rec <- which(active == 1 & status == "r")
nElig.rec <- length(idsElig.rec)
if (nElig.rec > 0) {
gElig.rec <- group[idsElig.rec]
ratesElig.rec <- rates.rec[gElig.rec]
vecDepartures.rec <- which(rbinom(nElig.rec, 1, ratesElig.rec) == 1)
if (length(vecDepartures.rec) > 0) {
idsDpt.rec <- idsElig.rec[vecDepartures.rec]
nDepartures.rec <- sum(group[idsDpt.rec] == 1)
nDeparturesG2.rec <- sum(group[idsDpt.rec] == 2)
active[idsDpt.rec] <- 0
exitTime[idsDpt.rec] <- at
}
}
}

# Output ------------------------------------------------------------------
dat <- set_attr(dat, "active", active)
dat <- set_attr(dat, "exitTime", exitTime)

dat <- set_epi(dat, "ds.flow", at, nDepartures.sus)
dat <- set_epi(dat, "di.flow", at, nDepartures.inf)
if (type == "SIR") {
dat <- set_epi(dat, "dr.flow", at, nDepartures.rec)
}
dat <- set_epi(dat, "ds.flow.g2", at, nDeparturesG2.sus)
dat <- set_epi(dat, "di.flow.g2", at, nDeparturesG2.inf)
if (type == "SIR") {
dat <- set_epi(dat, "dr.flow.g2", at, nDeparturesG2.rec)
}
return(dat)
arrivals_with_ngroups(dat, at, 1)
}


#' @title Arrivals: netsim Module
#'
#' @description This function simulates new arrivals into the network
#' for use in \code{\link{netsim}} simulations.
#'
#' @inheritParams recovery.net
#'
#' @inherit recovery.net return
#'
#' @seealso \code{\link{netsim}}
#'
#' @rdname arrivals.net
#' @export
#' @keywords netMod internal
#'
arrivals.2g.net <- function(dat, at) {
arrivals_with_ngroups(dat, at, 2)
}

arrivals_with_ngroups <- function(dat, at, ngroups) {
# Conditions --------------------------------------------------------------
vital <- get_param(dat, "vital")
if (vital == FALSE) {
return(dat)
}

suffixes <- c("", if (ngroups > 1) paste0(".g", seq_len(ngroups)[-1L]))

# Variables ---------------------------------------------------------------
a.rate <- get_param(dat, "a.rate")
a.rate.g2 <- get_param(dat, "a.rate.g2")
a.rate <- unlist(get_param_list(dat, paste0("a.rate", suffixes)))
a.rate[is.na(a.rate)] <- a.rate[1] # allow a.rate.g2 to be NA, for backwards compatibility

index <- at - 1
nOld <- get_epi(dat, "num", index)
nOldG2 <- get_epi(dat, "num.g2", index)
totArr <- nArrivals <- nArrivalsG2 <- 0
nOld <- integer(ngroups)
for (i in seq_len(ngroups)) {
nOld[i] <- get_epi(dat, paste0("num", suffixes[i]), index)
}

# Add Nodes ---------------------------------------------------------------
if (nOld > 0) {

if (is.na(a.rate.g2)) {
nArrivals <- rbinom(1, nOld, a.rate)
nArrivalsG2 <- rbinom(1, nOld, a.rate)
totArr <- nArrivals + nArrivalsG2
} else {
nArrivals <- rbinom(1, nOld, a.rate)
nArrivalsG2 <- rbinom(1, nOldG2, a.rate.g2)
totArr <- nArrivals + nArrivalsG2
nArrivals <- integer(ngroups)
if (nOld[1] > 0) { # should this be any(nOld > 0)?
for (i in seq_len(ngroups)) {
nArrivals[i] <- rbinom(1, nOld[i], a.rate[i])
}
totArr <- sum(nArrivals)

if (totArr > 0) {
dat <- append_core_attr(dat, at, totArr)
dat <- append_attr(dat, "group", 1, nArrivals)
dat <- append_attr(dat, "group", 2, nArrivalsG2)
if (ngroups > 1) {
for (i in seq_len(ngroups)) {
dat <- append_attr(dat, "group", i, nArrivals[i])
}
}
dat <- append_attr(dat, "status", "s", totArr)
dat <- append_attr(dat, "infTime", NA, totArr)
}
}

# Output ------------------------------------------------------------------
dat <- set_epi(dat, "a.flow", at, nArrivals)
dat <- set_epi(dat, "a.flow.g2", at, nArrivalsG2)

for (i in seq_len(ngroups)) {
dat <- set_epi(dat, paste0("a.flow", suffixes[i]), at, nArrivals[i])
}

return(dat)
}
3 changes: 3 additions & 0 deletions man/arrivals.net.Rd
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions man/departures.net.Rd
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 06fbe4d

Please sign in to comment.