Skip to content

Commit 64313e9

Browse files
pearsoncapearsonca
committed
fifelse optimizations
1 parent 9b8d67c commit 64313e9

File tree

3 files changed

+58
-83
lines changed

3 files changed

+58
-83
lines changed

‎R/aux_functions.R‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,7 @@ inf_fn <- function(inc_samp = NULL, k = NULL) {
1717
omega = 2,
1818
alpha = k)
1919

20-
out <- ifelse(out < 1, 1, out)
21-
22-
return(out)
20+
return(pmax(1, out))
2321
}
2422

2523
#' Calculate proportion of runs that have controlled outbreak

‎R/outbreak_setup.R‎

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -34,21 +34,23 @@
3434
#' out
3535
outbreak_setup <- function(num.initial.cases, incubation_period, onset_to_isolation, k, prop.asym) {
3636
# Set up table of initial cases
37-
inc_samples <- incubation_period(num.initial.cases)
38-
39-
case_data <- data.table(exposure = rep(0, num.initial.cases), # Exposure time of 0 for all initial cases
40-
asym = as.logical(rbinom(num.initial.cases, 1, prop.asym)),
41-
caseid = 1:(num.initial.cases), # set case id
42-
infector = 0,
43-
missed = TRUE,
44-
onset = inc_samples,
45-
new_cases = NA)
37+
case_data <- data.table(
38+
exposure = 0, # Exposure time of 0 for all initial cases
39+
asym = runif(num.initial.cases) < prop.asym,
40+
caseid = seq_len(num.initial.cases), # set case id
41+
infector = 0,
42+
isolated = FALSE,
43+
missed = TRUE,
44+
onset = incubation_period(num.initial.cases),
45+
new_cases = NA,
46+
isolated_time := Inf
47+
)
4648

4749
# set isolation time for cluster to minimum time of onset of symptoms + draw from delay distribution
48-
case_data <- case_data[, isolated_time := onset + onset_to_isolation(1)
49-
][, isolated := FALSE]
50-
51-
case_data$isolated_time[case_data$asym] <- Inf
50+
case_data <- case_data[
51+
asym == FALSE,
52+
isolated_time := pmin(onset + onset_to_isolation(.N))
53+
]
5254

5355
# return
5456
return(case_data)

‎R/outbreak_step.R‎

Lines changed: 42 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -61,15 +61,11 @@ outbreak_step <- function(case_data = NULL, disp.iso = NULL, disp.com = NULL,
6161

6262
# For each case in case_data, draw new_cases from a negative binomial distribution
6363
# with an R0 and dispersion dependent on if isolated=TRUE
64-
case_data[, new_cases := purrr::map2_dbl(
65-
ifelse(vect_isTRUE(isolated),
66-
disp.iso,
67-
ifelse(vect_isTRUE(asym), disp.subclin, disp.com)),
68-
ifelse(vect_isTRUE(isolated),
69-
r0isolated,
70-
ifelse(vect_isTRUE(asym),r0subclin, r0community)),
71-
~ rnbinom(1, size = .x, mu = .y))
72-
]
64+
case_data[, new_cases := rnbinom(
65+
.N,
66+
size = fifelse(isolated, disp.iso, fifelse(asym, disp.subclin, disp.com)),
67+
mu = fifelse(isolated, r0isolated, fifelse(asym, r0subclin, r0community))
68+
)]
7369

7470
# Select cases that have generated any new cases
7571
new_case_data <- case_data[new_cases > 0]
@@ -90,70 +86,56 @@ outbreak_step <- function(case_data = NULL, disp.iso = NULL, disp.com = NULL,
9086
}
9187

9288
# Compile a data.table for all new cases, new_cases is the amount of people that each infector has infected
93-
inc_samples <- incubation_period(total_new_cases)
94-
95-
prob_samples <- data.table(
89+
prob_samples <- new_case_data[, .(
9690
# time when new cases were exposed, a draw from serial interval based on infector's onset
97-
exposure = unlist(purrr::map2(new_case_data$new_cases, new_case_data$onset,
98-
function(x, y) {
99-
inf_fn(rep(y, x), k)
100-
})),
91+
exposure = inf_fn(rep(onset, new_cases), k),
10192
# records the infector of each new person
102-
infector = unlist(purrr::map2(new_case_data$caseid, new_case_data$new_cases,
103-
function(x, y) {
104-
rep(as.integer(x), as.integer(y))
105-
})),
106-
# records when infector was isolated
107-
infector_iso_time = unlist(purrr::map2(new_case_data$isolated_time, new_case_data$new_cases,
108-
function(x, y) {
109-
rep(x, as.integer(y))
110-
})),
93+
infector = rep(caseid, new_cases),
94+
# records when infector was isolated
95+
infector_iso_time = rep(isolated_time, new_cases),
11196
# records if infector asymptomatic
112-
infector_asym = unlist(purrr::map2(new_case_data$asym, new_case_data$new_cases,
113-
function(x, y) {
114-
rep(x, y)
115-
})),
116-
# draws a sample to see if this person is asymptomatic
117-
asym = as.logical(rbinom(n = total_new_cases, 1, prob = prop.asym)),
118-
# draws a sample to see if this person is traced
119-
missed = as.logical(rbinom(n = total_new_cases, 1, prob = 1 - prop.ascertain)),
97+
infector_asym = rep(asym, new_cases),
98+
# cases whose parents are asymptomatic are automatically missed;
99+
# will draw this for infector_asym == FALSE
100+
missed = TRUE,
120101
# sample from the incubation period for each new person
121-
incubfn_sample = inc_samples,
122-
isolated = FALSE,
123-
new_cases = NA
124-
)
125-
126-
127-
prob_samples <- prob_samples[exposure < infector_iso_time][, # filter out new cases prevented by isolation
128-
`:=`(# onset of new case is exposure + incubation period sample
129-
onset = exposure + incubfn_sample)]
130-
131-
132-
# cases whose parents are asymptomatic are automatically missed
133-
prob_samples$missed[vect_isTRUE(prob_samples$infector_asym)] <- TRUE
134-
135-
# If you are asymptomatic, your isolation time is Inf
136-
prob_samples[, isolated_time := ifelse(vect_isTRUE(asym), Inf,
137-
# If you are not asymptomatic, but you are missed,
138-
# you are isolated at your symptom onset
139-
ifelse(vect_isTRUE(missed), onset + onset_to_isolation(1),
140-
# If you are not asymptomatic and you are traced,
141-
# you are isolated at max(onset,infector isolation time) # max(onset,infector_iso_time)
142-
ifelse(!vect_isTRUE(rep(quarantine, total_new_cases)),
143-
pmin(onset + onset_to_isolation(1), pmax(onset, infector_iso_time)),
144-
infector_iso_time)))]
145-
102+
incubfn_sample = incubation_period(total_new_cases),
103+
isolated = FALSE, new_cases = NA
104+
)][,
105+
# draws a sample to see if this person is asymptomatic
106+
asym := runif(.N) < prop.asym
107+
][
108+
exposure < infector_iso_time
109+
][, # filter out new cases prevented by isolation
110+
`:=`(# onset of new case is exposure + incubation period sample
111+
onset = exposure + incubfn_sample
112+
)]
113+
114+
# draw a sample for missing
115+
prob_samples[infector_asym == FALSE, missed := runif(.N) > prop.ascertain]
116+
117+
prob_samples[, isolated_time := {
118+
ref_time <- onset + onset_to_isolation(.N)
119+
fifelse(
120+
# If asymptomatic, never isolated: time is Inf
121+
asym == TRUE, Inf, fifelse(
122+
# If not asymptomatic, but are missed, isolated at your symptom onset
123+
missed == TRUE, ref_time,
124+
if (quarantine == TRUE) infector_iso_time else
125+
# if symptomatic & traced, and infectors not quarantined
126+
pmin(onset + ref_time, pmax(onset, infector_iso_time)
127+
)))}]
146128

147129
# Chop out unneeded sample columns
148130
prob_samples[, c("incubfn_sample", "infector_iso_time", "infector_asym") := NULL]
149131
# Set new case ids for new people
150-
prob_samples$caseid <- (nrow(case_data) + 1):(nrow(case_data) + nrow(prob_samples))
132+
prob_samples[, caseid := case_data[.N, caseid] + seq_len(.N) ]
151133

152134
## Number of new cases
153135
cases_in_gen <- nrow(prob_samples)
154136

155137
## Estimate the effective r0
156-
effective_r0 <- nrow(prob_samples) / nrow(case_data[!vect_isTRUE(case_data$isolated)])
138+
effective_r0 <- cases_in_gen / case_data[isolated == FALSE, .N]
157139

158140
# Everyone in case_data so far has had their chance to infect and are therefore considered isolated
159141
case_data$isolated <- TRUE
@@ -168,10 +150,3 @@ outbreak_step <- function(case_data = NULL, disp.iso = NULL, disp.com = NULL,
168150

169151
return(out)
170152
}
171-
172-
173-
174-
# A vectorised version of isTRUE
175-
vect_isTRUE <- function(x) {
176-
purrr::map_lgl(x, isTRUE)
177-
}

0 commit comments

Comments
 (0)