Generating Scenario Results

Building on the quick start material, let's take on the common task of needing to simulate for several different parameter scenarios. We can do this by adding a scenarios block to pipeline configuration.

Configuration File - configs/scenarios_config.yaml

---
system:
  - module: wrapper
    state_change: flow
    script: model_input/plugins/SIR.py
    model_state:
      parameter_names: [s0, i0, r0]
      labels: [S, I, R]
engine:
  - module: wrapper
    state_change: flow
    script: model_input/plugins/solve_ivp.py
backend:
  - module: csv
simulate:
  quick_start:
    system: default
    engine: default
    backend: default
    times: "0.0:1.0:150.0"
    scenario: default
parameter:
  beta:
    module: fixed
    value: 0.3
  gamma:
    module: fixed
    value: 0.1
  s0:
    module: fixed
    value: 999
  i0:
    module: fixed
    value: 1
  r0:
    module: fixed
    value: 0
scenarios:
  - module: grid
    parameters:
      beta: [0.1, 0.3, 0.5]
      gamma: [0.05, 0.1, 0.2]
process:
  - module: shell
    command: Rscript scripts/plot_scenarios.R
    args:
      - configs/scenarios_config.yaml
      - model_output
      - figures/scenarios_plot.png

flepimop2 simulate scenarios_config.yaml # run the scenarios
ls model_output # should show 9 entries
flepimop2 process scenarios_config.yaml

The processing step uses the outputs to create a plot.

Processing Script - model_input/plot_scenarios.R

library(yaml)
library(data.table)
library(ggplot2)

.args <- if (interactive()) {
  c("configs/scenarios_config.yaml", "model_output", "output.png")
} else {
  commandArgs(trailingOnly = TRUE)
}

config <- read_yaml(.args[1])
results_dir <- .args[2]
results_dt <- list.files(results_dir, full.names = TRUE) |>
  lapply(fread) |>
  rbindlist(idcol = TRUE) |>
  setnames(c("scenario", "time", "S", "I", "R")) |>
  melt.data.table(id.vars = c("scenario", "time"), variable.name = "compartment", value.name = "count")

# r expand grid expansion is the opposite of the one in python:
# expands first item fully
# so we need to reverse the order of the parameters
scenarios <- as.data.table(expand.grid(rev(config$scenarios[[1]]$parameters)))[, scenario := seq_len(.N)]

p <- ggplot(results_dt[scenarios, on = .(scenario)]) + aes(x = time, y = count, color = compartment) +
    facet_grid(beta ~ gamma) +
    geom_line() +
    theme_bw()

output_file <- tail(.args, 1)

ggsave(output_file, p, width = 10, height = 8)