Visualize compartment summaries from odin

compartment-viz-functions.R

library(data.table)
library(tidyverse)

path_seirds_model <- system.file("examples/discrete_stochastic_seirds.R", package = "odin")

seirds_generator <- odin::odin(path_seirds_model)

x <- seirds_generator()

x_res <- as.data.frame(replicate(100, x$run(0:365)[, -1]))

names(x_res)

compartment_name <- "R"

x_res <- as.data.table(x_res) 

x_comp <- x_res[, .SD, .SDcols = patterns(compartment_name)]

quantz <- c(.025, .25, .5, .75, .0975)

summary_values <- purrr::map_dfc(quantz, 
														 ~matrixStats::rowQuantiles(as.matrix(x_comp), probs = .x, na.rm = T) )
names(summary_values) <- sprintf("q_%s",gsub(pattern = "\\.", "", quantz))

summary_values <- cbind(summary_values, mu = rowMeans(as.matrix(x_comp)))

#' @param dat the data 
#' @param compartment_name names of one of the compartments
#' @param quantz the returned quantiles with x being returned by default

summarise_compartments <- function(dat, compartment_name, quantz, market_share = 1){
	x_res <- data.table::as.data.table(dat) 
	
	x_comp <- x_res[, .SD, .SDcols = patterns(compartment_name)]
	
	quantz <- c(.025, .25, .5, .75, .975)
	
	summary_values <- purrr::map_dfc(quantz, 
																	 ~ceiling(matrixStats::rowQuantiles(as.matrix(x_comp), 
																	 																	 probs = .x, na.rm = T)*market_share) )
	names(summary_values) <- sprintf("q_%s",gsub(pattern = "\\.", "", quantz))
	
	summary_values <- cbind(summary_values, 
													mu = rowMeans(as.matrix(x_comp)),
													max_obs = matrixStats::rowMaxs(as.matrix(x_comp)))
	summary_values
}
out <- summarise_compartments(x_res, compartment_name = "E")

cone_colors <- c(
	`teal` = "#00a2b2",
	`navy` = "#00205b",
	`lime` = "#b7d866",
	`light orange` = "#f1bd51",
	`cyan` = "#63ccff",
	`purple` = "#7750a9"
)

horizon <- 14 # Days
first_date <- Sys.Date()
find_max <- max(matrixStats::rowMaxs(as.matrix(out[1:horizon,])))
out$date <-seq.Date(first_date, first_date +nrow(out)-1, "day")
cone_dark <- "#00A2B2"
op <- par(mar = c(1,1,1,1), col.main = cone_colors["navy"])
op
plot(NULL, xlim = c(first_date-1,first_date+horizon), 
		 ylim = c(-3,ceiling(find_max/5)*5), frame.plot = FALSE,axes = FALSE, 
		 xlab = '', ylab = 'Count', 
		 main = sprintf("Evolution of %s",compartment_name), adj=0, xpd = NA)
lines(out[["date"]], out[[1]], lty = 2)
lines(out[["date"]], out[[2]], lty = 2, col = cone_colors["cyan"])
#lines(out[["date"]], out[[3]], lty = 2)
lines(out[["date"]], out[[4]], lty = 2, col = cone_colors["cyan"])
lines(out[["date"]], out[[5]], lty = 2)
lines(out[["date"]], out[["mu"]], col = cone_colors["teal"])
lines(out[["date"]], out[["max_obs"]], col = "grey90")

x_labs <- seq.Date(first_date, first_date +horizon-1, "day")
y_labs <- pretty(seq(0,find_max+5,5))
abline(h = y_labs, lty = 3, col = "grey90")
text(x = x_labs,
		 y = -2.5, adj=0,
		 labels = format(x_labs, "%d-%B"), srt = 90, col = "grey50")

text(x = first_date-1, y = y_labs, adj=1,
		 labels = y_labs, srt = 0, col = "grey50")
mtext(text = "CONE HEALTH", 
			side = 1, 
			adj = 0,
			col = "black", padj = -.5)
mtext(text = "CONE", 
			side = 1, 
			adj =0, col = cone_colors[1], padj = -.5)
mtext(text = "Count", side = 2, col = "grey50")

points(x = history[["date"]][!violations],
			 y = history[["cases"]][!violations], pch = 19, col = cone_colors["light orange"])
points(x = history[["date"]][violations],
			 y = history[["cases"]][violations], pch = 19, col = "red")


history <- data.frame(
	date = seq.Date(Sys.Date(), Sys.Date()+4, "day"),
	cases = c(1,1,1,3,4)
)

violations <- history$cases>out$q_0975[1:nrow(history)]

plot_compartment_summary <- function(out, horizon = 14, first_date = Sys.Date(), history = NULL){
	cone_colors <- c(
		`teal` = "#00a2b2",
		`navy` = "#00205b",
		`lime` = "#b7d866",
		`light orange` = "#f1bd51",
		`cyan` = "#63ccff",
		`purple` = "#7750a9"
	)
	
	find_max <- max(matrixStats::rowMaxs(as.matrix(out[1:horizon,])))
	out$date <-seq.Date(first_date, first_date +nrow(out)-1, "day")
	
	op <- par(mar = c(1,1,1,1), col.main = cone_colors["navy"])
	op
	plot(NULL, xlim = c(first_date-1,first_date+horizon), 
			 ylim = c(-3,ceiling(find_max/5)*5), frame.plot = FALSE,axes = FALSE, 
			 xlab = '', ylab = 'Count', 
			 main = sprintf("Evolution of %s",compartment_name), adj=0, xpd = NA)
	polygon(c(out$date[1:horizon],rev(out$date[1:horizon])),
					c(out[[1]][1:horizon],rev(out[[5]][1:horizon])),
					col=ggplot2::alpha("grey90",.5), border = FALSE)
	lines(out[["date"]], out[[1]], lty = 2)
	lines(out[["date"]], out[[2]], lty = 2, col = cone_colors["cyan"])
	#lines(out[["date"]], out[[3]], lty = 2)
	lines(out[["date"]], out[[4]], lty = 2, col = cone_colors["cyan"])
	lines(out[["date"]], out[[5]], lty = 2)
	lines(out[["date"]], out[["mu"]], col = cone_colors["teal"])
	lines(out[["date"]], out[["max_obs"]], col = "grey90")
	
	x_labs <- seq.Date(first_date, first_date +horizon-1, "day")
	y_labs <- pretty(seq(0,find_max+5,5))
	abline(h = y_labs, lty = 3, col = "grey90")
	text(x = x_labs,
			 y = -2.5, adj=0,
			 labels = format(x_labs, "%d-%B"), srt = 90, col = "grey50")
	
	text(x = first_date-1, y = y_labs, adj=1,
			 labels = y_labs, srt = 0, col = "grey50")
	mtext(text = "CONE HEALTH", 
				side = 1, 
				adj = 0,
				col = "black", padj = -.5)
	mtext(text = "CONE", 
				side = 1, 
				adj =0, col = cone_colors[1], padj = -.5)
	mtext(text = "Count", side = 2, col = "grey50")
	
	
	if(!is.null(history)){
		violations <- history$cases>out$q_0975[1:nrow(history)]
		points(x = history[["date"]][!violations],
					 y = history[["cases"]][!violations], pch = 19, col = cone_colors["light orange"])
		points(x = history[["date"]][violations],
					 y = history[["cases"]][violations], pch = 19, col = "red")
	}
}

history <- data.frame(
	date = seq.Date(Sys.Date(), Sys.Date()+4, "day"),
	cases = c(1,1,1,3,4)
)

out <- summarise_compartments(x_res, compartment_name = "E")

plot_compartment_summary(out)

plot_compartment_summary(out, history = history)