From 6fcd95984b6af16635a8f776344039fe19ce2b1f Mon Sep 17 00:00:00 2001
From: Giovanni Charles <giovanni.charles@gmail.com>
Date: Thu, 21 Jul 2022 10:44:11 +0100
Subject: [PATCH 1/6] Add ITN usage renderer

---
 R/processes.R                        |  3 ++-
 R/vector_control.R                   | 10 ++++++++++
 tests/testthat/test-vector-control.R | 17 +++++++++++++++++
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/R/processes.R b/R/processes.R
index f5b4b6ee..fc7749ee 100644
--- a/R/processes.R
+++ b/R/processes.R
@@ -200,7 +200,8 @@ create_processes <- function(
         events$throw_away_net,
         parameters,
         correlations
-      )
+      ),
+      net_usage_renderer(variables$net_time, renderer)
     )
   }
 
diff --git a/R/vector_control.R b/R/vector_control.R
index f197a7d1..0537d746 100644
--- a/R/vector_control.R
+++ b/R/vector_control.R
@@ -185,3 +185,13 @@ bednet_decay <- function(t, gamma) {
 spraying_decay <- function(t, theta, gamma) {
   1 / (1 + exp(-(theta + gamma * t)))
 }
+
+net_usage_renderer <- function(net_time, renderer) {
+  function(t) {
+    renderer$render(
+      'net_usage',
+      net_time$get_index_of(-1)$not(TRUE)$size(),
+      t
+    )
+  }
+}
diff --git a/tests/testthat/test-vector-control.R b/tests/testthat/test-vector-control.R
index 8af13a2e..25a72bf3 100644
--- a/tests/testthat/test-vector-control.R
+++ b/tests/testthat/test-vector-control.R
@@ -314,3 +314,20 @@ test_that('prob_bitten correctly combines spraying and net probabilities', {
     tolerance=1e-4
   )
 })
+
+test_that('usage renderer outputs correct values', {
+  timestep <- 150
+
+  all <- individual::IntegerVariable$new(c(100, 50, 5, 50))
+  half <- individual::IntegerVariable$new(c(100, 50, -1, -1))
+  none <- individual::IntegerVariable$new(rep(-1, 4))
+
+  renderer <- list(render = mockery::mock())
+  net_usage_renderer(all, renderer)(timestep)
+  net_usage_renderer(half, renderer)(timestep)
+  net_usage_renderer(none, renderer)(timestep)
+
+  mockery::expect_args(renderer$render, 1, 'net_usage', 4, timestep)
+  mockery::expect_args(renderer$render, 2, 'net_usage', 2, timestep)
+  mockery::expect_args(renderer$render, 3, 'net_usage', 0, timestep)
+})

From 2c71a44fe0ae7db3e64f5d727200d6161c0d24fd Mon Sep 17 00:00:00 2001
From: Giovanni Charles <giovanni.charles@gmail.com>
Date: Fri, 22 Jul 2022 15:04:57 +0100
Subject: [PATCH 2/6] Update ITN usage output name

---
 R/vector_control.R                   | 2 +-
 tests/testthat/test-vector-control.R | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/R/vector_control.R b/R/vector_control.R
index 0537d746..3272b7f8 100644
--- a/R/vector_control.R
+++ b/R/vector_control.R
@@ -189,7 +189,7 @@ spraying_decay <- function(t, theta, gamma) {
 net_usage_renderer <- function(net_time, renderer) {
   function(t) {
     renderer$render(
-      'net_usage',
+      'n_use_net',
       net_time$get_index_of(-1)$not(TRUE)$size(),
       t
     )
diff --git a/tests/testthat/test-vector-control.R b/tests/testthat/test-vector-control.R
index 25a72bf3..47a89d9d 100644
--- a/tests/testthat/test-vector-control.R
+++ b/tests/testthat/test-vector-control.R
@@ -327,7 +327,7 @@ test_that('usage renderer outputs correct values', {
   net_usage_renderer(half, renderer)(timestep)
   net_usage_renderer(none, renderer)(timestep)
 
-  mockery::expect_args(renderer$render, 1, 'net_usage', 4, timestep)
-  mockery::expect_args(renderer$render, 2, 'net_usage', 2, timestep)
-  mockery::expect_args(renderer$render, 3, 'net_usage', 0, timestep)
+  mockery::expect_args(renderer$render, 1, 'n_use_net', 4, timestep)
+  mockery::expect_args(renderer$render, 2, 'n_use_net', 2, timestep)
+  mockery::expect_args(renderer$render, 3, 'n_use_net', 0, timestep)
 })

From 38b3ec4e26fc055e38bb2d54ef45f7fabba6bf1e Mon Sep 17 00:00:00 2001
From: Peter Winskill <p.winskill@imperial.ac.uk>
Date: Mon, 25 Jul 2022 09:42:44 +0100
Subject: [PATCH 3/6] Time varying demography

---
 DESCRIPTION                                 |   1 +
 R/mortality_processes.R                     |   2 +-
 R/parameters.R                              |   2 +
 R/population_parameters.R                   |  46 ++----
 R/processes.R                               |   5 +
 R/render.R                                  |  19 +++
 R/utils.R                                   |   9 ++
 man/run_metapop_simulation.Rd               |   2 +-
 man/run_simulation.Rd                       |   2 +-
 man/set_demography.Rd                       |   5 +-
 tests/testthat/test-demography.R            |   7 +-
 tests/testthat/test-mortality-integration.R |   4 +-
 vignettes/Demography.Rmd                    | 148 ++++++++++----------
 13 files changed, 131 insertions(+), 121 deletions(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index e60e26c9..bf6e19eb 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -64,6 +64,7 @@ Suggests:
   DiagrammeR,
   cowplot,
   ggplot2,
+  tidyr,
   covr,
   mgcv 
 RoxygenNote: 7.2.1
diff --git a/R/mortality_processes.R b/R/mortality_processes.R
index 5884c928..976383c0 100644
--- a/R/mortality_processes.R
+++ b/R/mortality_processes.R
@@ -19,7 +19,7 @@ create_mortality_process <- function(variables, events, renderer, parameters) {
       renderer$render('natural_deaths', died$size(), timestep)
     } else {
       age <- get_age(variables$birth$get_values(), timestep)
-      last_deathrate <- match_timestep(parameters$deathrate_timesteps, timestep)
+      last_deathrate <- match_last_timestep(parameters$deathrate_timesteps, timestep)
       deathrates <- rep(1, pop)
       age_groups <- .bincode(age, c(0, parameters$deathrate_agegroups))
       in_range <- !is.na(age_groups)
diff --git a/R/parameters.R b/R/parameters.R
index 309b0265..53756ee2 100644
--- a/R/parameters.R
+++ b/R/parameters.R
@@ -400,6 +400,8 @@ get_parameters <- function(overrides = list()) {
     severe_prevalence_rendering_max_ages = numeric(0),
     severe_incidence_rendering_min_ages = numeric(0),
     severe_incidence_rendering_max_ages = numeric(0),
+    age_group_rendering_min_ages = numeric(0),
+    age_group_rendering_max_ages = numeric(0),
     # misc
     custom_demography = FALSE,
     human_population = 100,
diff --git a/R/population_parameters.R b/R/population_parameters.R
index 84b00448..13d68264 100644
--- a/R/population_parameters.R
+++ b/R/population_parameters.R
@@ -3,8 +3,6 @@
 #' @param parameters the model parameters
 #' @param agegroups vector of agegroups (in timesteps)
 #' @param timesteps vector of timesteps for each change in demography
-#' @param birthrates vector of new birthrates (number of individuals born per
-#' timestep) for each element of `timesteps`
 #' @param deathrates matrix of deathrates per age group per timestep.
 #' Rows are timesteps from the `timesteps` param. Columns are the age groups
 #' from the `agegroups` param.
@@ -13,49 +11,27 @@ set_demography <- function(
   parameters,
   agegroups,
   timesteps,
-  birthrates,
   deathrates
   ) {
+  
+  if(min(timesteps) != 0){
+    stop("Must include the baseline demography (timesteps must include 0),
+         when setting a custom demography")
+  }
+  
   stopifnot(all(agegroups > 0))
   stopifnot(all(timesteps >= 0))
-  stopifnot(all(birthrates > 0))
-  stopifnot(length(birthrates) == length(timesteps))
   stopifnot(all(deathrates > 0 & deathrates < 1))
   stopifnot(length(agegroups) == ncol(deathrates))
   stopifnot(length(timesteps) == nrow(deathrates))
   stopifnot(!is.unsorted(timesteps, strictly = TRUE))
-  stopifnot(length(timesteps) == 1) # changing population is not yet supported
+
   parameters$custom_demography <- TRUE
   parameters$deathrate_agegroups <- agegroups
   parameters$deathrate_timesteps <- timesteps
   parameters$deathrates <- deathrates
-  parameters$birthrate_timesteps <- timesteps
-  parameters$birthrates <- birthrates
-  n_age <- length(agegroups)
-
-  # set the populations
-  populations <- vapply(
-    seq(timesteps),
-    function(timestep) {
-      get_equilibrium_population(
-        agegroups,
-        birthrates[[timestep]],
-        parameters$deathrates[timestep,]
-      )
-    },
-    numeric(n_age)
-  )
-
-  deathrates <- vnapply(
-    seq(timesteps),
-    function(timestep) {
-      sum(
-         parameters$deathrates[timestep,] * populations[timestep,]
-      )
-    }
-  )
-  parameters$human_population <- round(colSums(populations))
-  parameters$human_population_timesteps <- timesteps
+  parameters$birthrates <- find_birthrates(parameters$human_population, agegroups, deathrates[1,])
+  parameters$birthrate_timesteps <- 0
 
   parameters
 }
@@ -104,11 +80,11 @@ get_birthrate <- function(parameters, timestep) {
   if (!parameters$custom_demography) {
     return(1 / parameters$average_age * get_human_population(parameters, timestep))
   }
-  last_birthrate <- match_timestep(parameters$birthrate_timesteps, timestep)
+  last_birthrate <- match_last_timestep(parameters$birthrate_timesteps, timestep)
   parameters$birthrates[last_birthrate]
 }
 
 get_human_population <- function(parameters, timestep) {
-  last_pop <- match_timestep(parameters$human_population_timesteps, timestep)
+  last_pop <- match_last_timestep(parameters$human_population_timesteps, timestep)
   parameters$human_population[last_pop]
 }
diff --git a/R/processes.R b/R/processes.R
index beedff6c..133ebb49 100644
--- a/R/processes.R
+++ b/R/processes.R
@@ -158,6 +158,11 @@ create_processes <- function(
       parameters,
       renderer
     ),
+    create_age_group_renderer(
+      variables$birth,
+      parameters,
+      renderer
+    ),
     create_compartmental_rendering_process(renderer, solvers, parameters)
   )
 
diff --git a/R/render.R b/R/render.R
index 3dddf94c..a62c4dec 100644
--- a/R/render.R
+++ b/R/render.R
@@ -152,3 +152,22 @@ create_total_M_renderer_compartmental <- function(renderer, solvers, parameters)
     }
   }
 }
+
+create_age_group_renderer <- function(
+    birth,
+    parameters,
+    renderer
+) {
+  function(timestep) {
+    for (i in seq_along(parameters$age_group_rendering_min_ages)) {
+      lower <- parameters$age_group_rendering_min_ages[[i]]
+      upper <- parameters$age_group_rendering_max_ages[[i]]
+      in_age <- in_age_range(birth, timestep, lower, upper)
+      renderer$render(
+        paste0('n_age_', lower, '_', upper),
+        in_age$size(),
+        timestep
+      ) 
+    }
+  }
+}
diff --git a/R/utils.R b/R/utils.R
index 506357cf..f2b3dfd1 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -64,3 +64,12 @@ rtexp <- function(n, m, t) { itexp(runif(n), m, t) }
 match_timestep <- function(ts, t) {
   min(sum(ts < t) + 1, length(ts))
 }
+
+#'@title Find index of the latest timestep in vector of timesteps
+#'@param ts timesteps, assumed to be sorted and unique
+#'@param t current timestep
+#'@noRd
+match_last_timestep <- function(ts, t) {
+  min(sum(ts <= t), length(ts))
+}
+
diff --git a/man/run_metapop_simulation.Rd b/man/run_metapop_simulation.Rd
index 83b83968..a33d6875 100644
--- a/man/run_metapop_simulation.Rd
+++ b/man/run_metapop_simulation.Rd
@@ -15,7 +15,7 @@ run_metapop_simulation(timesteps, parameters, correlations = NULL, mixing)
 (default: NULL)}
 
 \item{mixing}{matrix of mixing coefficients for infectivity towards
-mosquitoes}
+mosquitoes. Each element must be between 0 and 1 and all rows and columns must sum to 1.}
 }
 \value{
 a list of dataframe of results
diff --git a/man/run_simulation.Rd b/man/run_simulation.Rd
index 7e4f795c..e4e54170 100644
--- a/man/run_simulation.Rd
+++ b/man/run_simulation.Rd
@@ -88,7 +88,7 @@ asymptomatic
 subpatent
 \item rate_U_S: rate that humans transition from subpatent to
 susceptible
-\item net_usage: the proportion of the population protected by a bed net
+\item net_usage: the number people protected by a bed net
 \item mosquito_deaths: number of adult female mosquitoes who die this timestep
 }
 }
diff --git a/man/set_demography.Rd b/man/set_demography.Rd
index cdc765b4..faf43038 100644
--- a/man/set_demography.Rd
+++ b/man/set_demography.Rd
@@ -4,7 +4,7 @@
 \alias{set_demography}
 \title{Parameterise variable deathrates}
 \usage{
-set_demography(parameters, agegroups, timesteps, birthrates, deathrates)
+set_demography(parameters, agegroups, timesteps, deathrates)
 }
 \arguments{
 \item{parameters}{the model parameters}
@@ -13,9 +13,6 @@ set_demography(parameters, agegroups, timesteps, birthrates, deathrates)
 
 \item{timesteps}{vector of timesteps for each change in demography}
 
-\item{birthrates}{vector of new birthrates (number of individuals born per
-timestep) for each element of \code{timesteps}}
-
 \item{deathrates}{matrix of deathrates per age group per timestep.
 Rows are timesteps from the \code{timesteps} param. Columns are the age groups
 from the \code{agegroups} param.}
diff --git a/tests/testthat/test-demography.R b/tests/testthat/test-demography.R
index 13350357..c471fe90 100644
--- a/tests/testthat/test-demography.R
+++ b/tests/testthat/test-demography.R
@@ -37,9 +37,8 @@ test_that('calculate_initial_ages calculates truncated exp custom demographic',
   parameters <- set_demography(
     parameters,
     agegroups = ages,
-    timesteps = 1,
-    birthrates = find_birthrates(4, ages, deathrates),
-    deathrates = matrix(deathrates, nrow=1, ncol=2)
+    timesteps = 0,
+    deathrates = matrix(deathrates, nrow = 1, ncol = 2)
   )
   mock_groups <- mockery::mock(c(2, 1, 2, 1))
   mock_rtexp <- mockery::mock(c(25 * 365, 30 * 365), c(25 * 365, 30 * 365))
@@ -51,7 +50,7 @@ test_that('calculate_initial_ages calculates truncated exp custom demographic',
     mockery::mock(c(3, 1))
   )
   ages <- calculate_initial_ages(parameters)
-  mockery::expect_args(mock_groups, 1, 2, 4, replace = TRUE, prob = c(3, 1))
+  mockery::expect_args(mock_groups, 1, 2, 100, replace = TRUE, prob = c(3, 1))
   mockery::expect_args(mock_rtexp, 1, 2, .5, 50 * 365)
   mockery::expect_args(mock_rtexp, 2, 2, .75, 50 * 365)
   expect_setequal(ages, c(25 * 365, 75 * 365, 30 * 365, 80 * 365))
diff --git a/tests/testthat/test-mortality-integration.R b/tests/testthat/test-mortality-integration.R
index 138f3d19..2e1fd64e 100644
--- a/tests/testthat/test-mortality-integration.R
+++ b/tests/testthat/test-mortality-integration.R
@@ -49,13 +49,13 @@ test_that('mortality_process resets humans correctly', {
 test_that('mortality_process samples deaths from a custom demography', {
   timestep <- 2
   parameters <- get_parameters()
+  parameters$human_population <- 4
   ages <- c(50, 100) * 365
   deaths <- c(.5, .75)
   parameters <- set_demography(
     parameters,
     agegroups = ages,
-    timesteps = 1,
-    birthrates = find_birthrates(4, ages, deaths),
+    timesteps = 0,
     deathrates = matrix(deaths, nrow=1, ncol=2)
   )
   events <- create_events(parameters)
diff --git a/vignettes/Demography.Rmd b/vignettes/Demography.Rmd
index 06549855..20c0dac9 100644
--- a/vignettes/Demography.Rmd
+++ b/vignettes/Demography.Rmd
@@ -16,6 +16,7 @@ knitr::opts_chunk$set(
 
 ```{r setup}
 suppressPackageStartupMessages(library(ggplot2))
+library(tidyr)
 library(malariasimulation)
 ```
 
@@ -28,17 +29,13 @@ year <- 365
 month <- 30
 sim_length <- 5 * year
 human_population <- 1000
-starting_EIR <- 50
+starting_EIR <- 5
 
 simparams <- get_parameters(
   list(
     human_population = human_population,
-    model_seasonality = TRUE, # Let's try a bi-modal model
-    g0 = 0.28605,
-    g = c(0.20636, -0.0740318, -0.0009293),
-    h = c(0.173743, -0.0730962, -0.116019),
-    severe_incidence_rendering_min_ages = 2 * year,
-    severe_incidence_rendering_max_ages = 10 * year
+    age_group_rendering_min_ages = seq(0, 80, 5) * 365,
+    age_group_rendering_max_ages = seq(5, 85, 5) * 365
   )
 )
 
@@ -47,92 +44,97 @@ simparams <- set_equilibrium(simparams, starting_EIR)
 
 ## Custom demography
 
-We can set a custom demography and compare the severe outputs:
+We can set a custom demography:
 
 ```{r}
 demography_params <- simparams
 
 # Set a flat demography
 ages <- round(c(
-  .083333,
-  1,
-  5,
-  10,
-  15,
-  20,
-  25,
-  30,
-  35,
-  40,
-  45,
-  50,
-  55,
-  60,
-  65,
-  70,
-  75,
-  80,
-  85,
-  90,
-  95,
-  200
-) * year)
+  .083333, 1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80,
+  85, 90, 95, 200) * year)
 
 deathrates <- c(
-  .4014834,
-  .0583379,
-  .0380348,
-  .0395061,
-  .0347255,
-  .0240849,
-  .0300902,
-  .0357914,
-  .0443123,
-  .0604932,
-  .0466799,
-  .0426199,
-  .0268332,
-  .049361,
-  .0234852,
-  .0988317,
-  .046755,
-  .1638875,
-  .1148753,
-  .3409079,
-  .2239224,
-  .8338688
-) / 365
+  .4014834, .0583379, .0380348, .0395061, .0347255, .0240849, .0300902,
+  .0357914, .0443123, .0604932, .0466799, .0426199, .0268332, .049361,
+  .0234852, .0988317, .046755, .1638875, .1148753, .3409079, .2239224,
+  .8338688) / 365
 
 demography_params <- set_demography(
   demography_params,
   agegroups = ages,
-  timesteps = 1,
-  deathrates = matrix(deathrates, nrow = 1),
-  birthrates = find_birthrates(human_population, ages, deathrates)
+  timesteps = 0,
+  deathrates = matrix(deathrates, nrow = 1)
 )
+```
+
+Let's run the simulations and compare the age distributions in the populations
+in year 5:
 
-# combine the outputs
+```{r}
+# run and combine the outputs
 exp_output <- run_simulation(sim_length, simparams)
 exp_output$run <- 'exponential'
 custom_output <- run_simulation(sim_length, demography_params)
 custom_output$run <- 'custom'
 output <- rbind(exp_output, custom_output)
 
-# calculate yearly prevalence
-yearly_output <- aggregate(
-  output$n_inc_severe_730_3650,
-  by = list(year = floor(output$timestep / year), run = output$run),
-  FUN = sum
-)
+# wrangle outputs
+output_long <- output |>
+  dplyr::filter(timestep == 5 * 365) |>
+  dplyr::select(run, timestep, contains("n_age")) |>
+  tidyr::pivot_longer(cols = -c(run, timestep), names_to = "age_group", values_to = "n", names_prefix = "n_age_") |>
+  tidyr::separate(col = age_group, sep = "_", into = c("age_lower", "age_upper"), convert = TRUE) |>
+  dplyr::mutate(age_mid = (age_lower + (age_upper - age_lower) / 2) / 365,
+                run = factor(run, levels = c("exponential", "custom")))
+
+# Plot the age distributions
+ggplot(output_long, aes(x = age_mid, y = n)) +
+  geom_bar(stat = "identity") +
+  theme_bw() +
+  facet_wrap(~ run)
 
-# Plot the output
-ggplot(yearly_output) + geom_line(
-  aes(
-    x = year,
-    y = x,
-    group = run,
-    color = run
-  )
-) + labs(x = "Year", y = "PfPR2-10 severe")
+```
+
+We can also specify time-varying death rates to capture a dynamic demography
+
+```{r}
+dynamic_demography_params <- simparams
+
+# Set a flat demography
+ages <- round(c(
+  .083333, 1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80,
+  85, 90, 95, 200) * year)
+
+deathrates <- c(
+  .4014834, .0583379, .0380348, .0395061, .0347255, .0240849, .0300902,
+  .0357914, .0443123, .0604932, .0466799, .0426199, .0268332, .049361,
+  .0234852, .0988317, .046755, .1638875, .1148753, .3409079, .2239224,
+  .8338688) / 365
+# Let's increase the death rates in some age groups
+deathrates_changed <- deathrates
+deathrates_changed[3:6] <- deathrates_changed[3:6] * 10
+
+dynamic_demography_params <- set_demography(
+  dynamic_demography_params,
+  agegroups = ages,
+  timesteps = c(0, 2 * 365),
+  deathrates = matrix(c(deathrates, deathrates_changed), nrow = 2, byrow = TRUE)
+)
 
+dynamic_demography_output <- run_simulation(sim_length, dynamic_demography_params)
+
+# wrangle outputs
+dynamic_output_long <- dynamic_demography_output |>
+  dplyr::filter(timestep %in% (1:5 * 365)) |>
+  dplyr::select(timestep, contains("n_age")) |>
+  tidyr::pivot_longer(cols = -timestep, names_to = "age_group", values_to = "n", names_prefix = "n_age_") |>
+  tidyr::separate(col = age_group, sep = "_", into = c("age_lower", "age_upper"), convert = TRUE) |>
+  dplyr::mutate(age_mid = (age_lower + (age_upper - age_lower) / 2) / 365)
+
+# Plot the age distributions each year
+ggplot(dynamic_output_long, aes(x = age_mid, y = n)) +
+  geom_bar(stat = "identity") +
+  theme_bw() +
+  facet_wrap(~ timestep / 365, ncol = 5)
 ```
\ No newline at end of file

From 4670e6a8bd2ea3f9cc143cdcc71ad957d2d410c3 Mon Sep 17 00:00:00 2001
From: Peter Winskill <p.winskill@imperial.ac.uk>
Date: Tue, 26 Jul 2022 11:01:04 +0100
Subject: [PATCH 4/6] Tidy unneeded birthrate functionality

---
 NAMESPACE                 |  1 -
 R/population_parameters.R | 14 +-------------
 R/variables.R             |  2 +-
 3 files changed, 2 insertions(+), 15 deletions(-)

diff --git a/NAMESPACE b/NAMESPACE
index c0d7331a..ea369539 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -4,7 +4,6 @@ export(AL_params)
 export(DHA_PQP_params)
 export(SP_AQ_params)
 export(arab_params)
-export(find_birthrates)
 export(fun_params)
 export(gamb_params)
 export(get_correlation_parameters)
diff --git a/R/population_parameters.R b/R/population_parameters.R
index 13d68264..3b08cc3d 100644
--- a/R/population_parameters.R
+++ b/R/population_parameters.R
@@ -14,13 +14,12 @@ set_demography <- function(
   deathrates
   ) {
   
+  stopifnot(all(timesteps >= 0))
   if(min(timesteps) != 0){
     stop("Must include the baseline demography (timesteps must include 0),
          when setting a custom demography")
   }
-  
   stopifnot(all(agegroups > 0))
-  stopifnot(all(timesteps >= 0))
   stopifnot(all(deathrates > 0 & deathrates < 1))
   stopifnot(length(agegroups) == ncol(deathrates))
   stopifnot(length(timesteps) == nrow(deathrates))
@@ -30,8 +29,6 @@ set_demography <- function(
   parameters$deathrate_agegroups <- agegroups
   parameters$deathrate_timesteps <- timesteps
   parameters$deathrates <- deathrates
-  parameters$birthrates <- find_birthrates(parameters$human_population, agegroups, deathrates[1,])
-  parameters$birthrate_timesteps <- 0
 
   parameters
 }
@@ -63,7 +60,6 @@ get_equilibrium_population <- function(age_high, birthrate, deathrates) {
 #' @param age_high a vector of age groups
 #' @param deathrates vector of deathrates for each age group
 #' @importFrom stats uniroot
-#' @export
 find_birthrates <- function(pops, age_high, deathrates) {
   vnapply(
     pops,
@@ -76,14 +72,6 @@ find_birthrates <- function(pops, age_high, deathrates) {
   )
 }
 
-get_birthrate <- function(parameters, timestep) {
-  if (!parameters$custom_demography) {
-    return(1 / parameters$average_age * get_human_population(parameters, timestep))
-  }
-  last_birthrate <- match_last_timestep(parameters$birthrate_timesteps, timestep)
-  parameters$birthrates[last_birthrate]
-}
-
 get_human_population <- function(parameters, timestep) {
   last_pop <- match_last_timestep(parameters$human_population_timesteps, timestep)
   parameters$human_population[last_pop]
diff --git a/R/variables.R b/R/variables.R
index 95acb018..51197dfb 100644
--- a/R/variables.R
+++ b/R/variables.R
@@ -410,7 +410,7 @@ calculate_initial_ages <- function(parameters) {
   age_width <- diff(c(0, age_high))
   age_low <- age_high - age_width
   n_age <- length(age_high)
-  birthrate <- get_birthrate(parameters, 0)
+  birthrate <- find_birthrates(parameters$human_population, agegroups, deathrates[1,])
   deathrates <- parameters$deathrates[1,]
 
   eq_pop <- get_equilibrium_population(age_high, birthrate, deathrates)

From 73022eda8f971f3f13dd8b009fd59e254121cd6c Mon Sep 17 00:00:00 2001
From: Peter Winskill <p.winskill@imperial.ac.uk>
Date: Tue, 26 Jul 2022 11:04:22 +0100
Subject: [PATCH 5/6] Fix indexing of match_timestep

---
 R/mortality_processes.R          |  2 +-
 R/population_parameters.R        |  2 +-
 R/utils.R                        |  8 --------
 tests/testthat/test-demography.R | 16 ++++++++++++++++
 4 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/R/mortality_processes.R b/R/mortality_processes.R
index 976383c0..5884c928 100644
--- a/R/mortality_processes.R
+++ b/R/mortality_processes.R
@@ -19,7 +19,7 @@ create_mortality_process <- function(variables, events, renderer, parameters) {
       renderer$render('natural_deaths', died$size(), timestep)
     } else {
       age <- get_age(variables$birth$get_values(), timestep)
-      last_deathrate <- match_last_timestep(parameters$deathrate_timesteps, timestep)
+      last_deathrate <- match_timestep(parameters$deathrate_timesteps, timestep)
       deathrates <- rep(1, pop)
       age_groups <- .bincode(age, c(0, parameters$deathrate_agegroups))
       in_range <- !is.na(age_groups)
diff --git a/R/population_parameters.R b/R/population_parameters.R
index 3b08cc3d..89015790 100644
--- a/R/population_parameters.R
+++ b/R/population_parameters.R
@@ -73,6 +73,6 @@ find_birthrates <- function(pops, age_high, deathrates) {
 }
 
 get_human_population <- function(parameters, timestep) {
-  last_pop <- match_last_timestep(parameters$human_population_timesteps, timestep)
+  last_pop <- match_timestep(parameters$human_population_timesteps, timestep)
   parameters$human_population[last_pop]
 }
diff --git a/R/utils.R b/R/utils.R
index f2b3dfd1..343e6309 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -62,14 +62,6 @@ rtexp <- function(n, m, t) { itexp(runif(n), m, t) }
 #'@param t current timestep
 #'@noRd
 match_timestep <- function(ts, t) {
-  min(sum(ts < t) + 1, length(ts))
-}
-
-#'@title Find index of the latest timestep in vector of timesteps
-#'@param ts timesteps, assumed to be sorted and unique
-#'@param t current timestep
-#'@noRd
-match_last_timestep <- function(ts, t) {
   min(sum(ts <= t), length(ts))
 }
 
diff --git a/tests/testthat/test-demography.R b/tests/testthat/test-demography.R
index c471fe90..572706c2 100644
--- a/tests/testthat/test-demography.R
+++ b/tests/testthat/test-demography.R
@@ -55,3 +55,19 @@ test_that('calculate_initial_ages calculates truncated exp custom demographic',
   mockery::expect_args(mock_rtexp, 2, 2, .75, 50 * 365)
   expect_setequal(ages, c(25 * 365, 75 * 365, 30 * 365, 80 * 365))
 })
+
+test_that('match_timestep always gives 0 for constant demography', {
+  expect_equal(match_timestep(c(0), 0), 1)
+  expect_equal(match_timestep(c(0), 1), 1)
+  expect_equal(match_timestep(c(0), 50), 1)
+})
+
+test_that('match_timestep works on the boundaries', {
+  expect_equal(match_timestep(c(0, 50, 100), 0), 1)
+  expect_equal(match_timestep(c(0, 50, 100), 1), 1)
+  expect_equal(match_timestep(c(0, 50, 100), 49), 1)
+  expect_equal(match_timestep(c(0, 50, 100), 50), 2)
+  expect_equal(match_timestep(c(0, 50, 100), 99), 2)
+  expect_equal(match_timestep(c(0, 50, 100), 100), 3)
+  expect_equal(match_timestep(c(0, 50, 100), 101), 3)
+})
\ No newline at end of file

From d973473b5cb54fcbc70d5a7de0e672a5f8dbe3ba Mon Sep 17 00:00:00 2001
From: Peter Winskill <p.winskill@imperial.ac.uk>
Date: Tue, 26 Jul 2022 11:38:22 +0100
Subject: [PATCH 6/6] Remove tidyr dependency from demography vignette

---
 DESCRIPTION                      |  1 -
 R/variables.R                    |  3 +-
 tests/testthat/test-demography.R |  4 +--
 vignettes/Demography.Rmd         | 60 ++++++++++++++++++++------------
 4 files changed, 42 insertions(+), 26 deletions(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index bf6e19eb..e60e26c9 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -64,7 +64,6 @@ Suggests:
   DiagrammeR,
   cowplot,
   ggplot2,
-  tidyr,
   covr,
   mgcv 
 RoxygenNote: 7.2.1
diff --git a/R/variables.R b/R/variables.R
index 51197dfb..b12d1d8d 100644
--- a/R/variables.R
+++ b/R/variables.R
@@ -406,11 +406,12 @@ calculate_initial_ages <- function(parameters) {
     )))
   }
 
+  deathrates <- parameters$deathrates[1, , drop = FALSE]
   age_high <- parameters$deathrate_agegroups
   age_width <- diff(c(0, age_high))
   age_low <- age_high - age_width
   n_age <- length(age_high)
-  birthrate <- find_birthrates(parameters$human_population, agegroups, deathrates[1,])
+  birthrate <- find_birthrates(parameters$human_population, age_high, deathrates)
   deathrates <- parameters$deathrates[1,]
 
   eq_pop <- get_equilibrium_population(age_high, birthrate, deathrates)
diff --git a/tests/testthat/test-demography.R b/tests/testthat/test-demography.R
index 572706c2..3b08dd37 100644
--- a/tests/testthat/test-demography.R
+++ b/tests/testthat/test-demography.R
@@ -31,7 +31,7 @@ test_that('find_birthrates is consistent with get_equilibrium_population', {
 
 
 test_that('calculate_initial_ages calculates truncated exp custom demographic', {
-  parameters <- get_parameters()
+  parameters <- get_parameters(list(human_population = 4))
   ages <- c(50, 100) * 365
   deathrates <- c(.5, .75)
   parameters <- set_demography(
@@ -50,7 +50,7 @@ test_that('calculate_initial_ages calculates truncated exp custom demographic',
     mockery::mock(c(3, 1))
   )
   ages <- calculate_initial_ages(parameters)
-  mockery::expect_args(mock_groups, 1, 2, 100, replace = TRUE, prob = c(3, 1))
+  mockery::expect_args(mock_groups, 1, 2, 4, replace = TRUE, prob = c(3, 1))
   mockery::expect_args(mock_rtexp, 1, 2, .5, 50 * 365)
   mockery::expect_args(mock_rtexp, 2, 2, .75, 50 * 365)
   expect_setequal(ages, c(25 * 365, 75 * 365, 30 * 365, 80 * 365))
diff --git a/vignettes/Demography.Rmd b/vignettes/Demography.Rmd
index 20c0dac9..cfa68a00 100644
--- a/vignettes/Demography.Rmd
+++ b/vignettes/Demography.Rmd
@@ -16,7 +16,6 @@ knitr::opts_chunk$set(
 
 ```{r setup}
 suppressPackageStartupMessages(library(ggplot2))
-library(tidyr)
 library(malariasimulation)
 ```
 
@@ -31,11 +30,14 @@ sim_length <- 5 * year
 human_population <- 1000
 starting_EIR <- 5
 
+age_min <- seq(0, 80, 5) * 365
+age_max <- seq(5, 85, 5) * 365
+
 simparams <- get_parameters(
   list(
     human_population = human_population,
-    age_group_rendering_min_ages = seq(0, 80, 5) * 365,
-    age_group_rendering_max_ages = seq(5, 85, 5) * 365
+    age_group_rendering_min_ages = age_min,
+    age_group_rendering_max_ages = age_max
   )
 )
 
@@ -68,8 +70,7 @@ demography_params <- set_demography(
 )
 ```
 
-Let's run the simulations and compare the age distributions in the populations
-in year 5:
+Let's run the simulations 
 
 ```{r}
 # run and combine the outputs
@@ -77,19 +78,36 @@ exp_output <- run_simulation(sim_length, simparams)
 exp_output$run <- 'exponential'
 custom_output <- run_simulation(sim_length, demography_params)
 custom_output$run <- 'custom'
-output <- rbind(exp_output, custom_output)
+```
+
+and now we can compare the age distributions in the populations
+in year 5:
 
+```{r}
 # wrangle outputs
-output_long <- output |>
-  dplyr::filter(timestep == 5 * 365) |>
-  dplyr::select(run, timestep, contains("n_age")) |>
-  tidyr::pivot_longer(cols = -c(run, timestep), names_to = "age_group", values_to = "n", names_prefix = "n_age_") |>
-  tidyr::separate(col = age_group, sep = "_", into = c("age_lower", "age_upper"), convert = TRUE) |>
-  dplyr::mutate(age_mid = (age_lower + (age_upper - age_lower) / 2) / 365,
-                run = factor(run, levels = c("exponential", "custom")))
+output <- rbind(exp_output, custom_output)
+output <- output[output$timestep == 5 * 365,]
+
+# A function to extract the age variables and convert to long format
+convert_to_long <-  function(age_min, age_max, output){
+  output <- lapply(
+    seq_along(age_min),
+    function(i) {
+      data.frame(
+        age_lower = age_min[[i]],
+        age_upper = age_max[[i]],
+        n = output[,paste0('n_age_', age_min[[i]], '_',age_max[[i]])],
+        age_mid = (age_min[[i]] + (age_min[[i]] - age_max[[i]]) / 2) / 365,
+        run = output$run, 
+        timestep = output$timestep)
+    }
+  )
+  output <- do.call("rbind", output)
+}
+output <- convert_to_long(age_min, age_max, output)
 
 # Plot the age distributions
-ggplot(output_long, aes(x = age_mid, y = n)) +
+ggplot(output, aes(x = age_mid, y = n)) +
   geom_bar(stat = "identity") +
   theme_bw() +
   facet_wrap(~ run)
@@ -98,7 +116,7 @@ ggplot(output_long, aes(x = age_mid, y = n)) +
 
 We can also specify time-varying death rates to capture a dynamic demography
 
-```{r}
+```{r, fig.width=7}
 dynamic_demography_params <- simparams
 
 # Set a flat demography
@@ -125,15 +143,13 @@ dynamic_demography_params <- set_demography(
 dynamic_demography_output <- run_simulation(sim_length, dynamic_demography_params)
 
 # wrangle outputs
-dynamic_output_long <- dynamic_demography_output |>
-  dplyr::filter(timestep %in% (1:5 * 365)) |>
-  dplyr::select(timestep, contains("n_age")) |>
-  tidyr::pivot_longer(cols = -timestep, names_to = "age_group", values_to = "n", names_prefix = "n_age_") |>
-  tidyr::separate(col = age_group, sep = "_", into = c("age_lower", "age_upper"), convert = TRUE) |>
-  dplyr::mutate(age_mid = (age_lower + (age_upper - age_lower) / 2) / 365)
+dynamic_demography_output <- dynamic_demography_output[dynamic_demography_output$timestep %in% (1:5 * 365),]
+dynamic_demography_output$run <- "dynamic"
+
+dynamic_demography_output <- convert_to_long(age_min, age_max, dynamic_demography_output)
 
 # Plot the age distributions each year
-ggplot(dynamic_output_long, aes(x = age_mid, y = n)) +
+ggplot(dynamic_demography_output, aes(x = age_mid, y = n)) +
   geom_bar(stat = "identity") +
   theme_bw() +
   facet_wrap(~ timestep / 365, ncol = 5)