From 5636839f75b05466bef1ebca04d64011ceef78c0 Mon Sep 17 00:00:00 2001
From: Giovanni Charles <giovanni.charles@gmail.com>
Date: Wed, 24 Jan 2024 16:12:45 +0000
Subject: [PATCH] Extend grid to all statistics:

 * parameterise grid outputs with string vector
 * fix incidence grid rendering for p_ stats
---
 R/parameters.R                    |  4 +--
 R/render.R                        | 56 ++++++++++++++++++++++++++-----
 tests/testthat/test-grid-render.R | 19 +++++++++++
 3 files changed, 68 insertions(+), 11 deletions(-)

diff --git a/R/parameters.R b/R/parameters.R
index 70ac2e86..03d6d886 100644
--- a/R/parameters.R
+++ b/R/parameters.R
@@ -207,7 +207,7 @@
 #' * severe_incidence_rendering_min_ages - the minimum ages for severe incidence
 #' outputs; default = turned off
 #' * severe_incidence_rendering_max_ages - the corresponding max ages; default = turned off
-#' * render_grid - whether to render prevalence and incidence values for a grid of year wide age groups between 0 and 100; default = turned off
+#' * render_grid - which statistics to render values for in a grid of year wide age groups between 0 and 100; e.g. 'p_detect_' ; default = turned off
 #'
 #' miscellaneous:
 #'
@@ -395,7 +395,7 @@ get_parameters <- function(overrides = list()) {
     severe_incidence_rendering_max_ages = numeric(0),
     age_group_rendering_min_ages = numeric(0),
     age_group_rendering_max_ages = numeric(0),
-    render_grid = FALSE,
+    render_grid = NULL,
     # misc
     custom_demography = FALSE,
     human_population = 100,
diff --git a/R/render.R b/R/render.R
index 326eda14..d1ed2679 100644
--- a/R/render.R
+++ b/R/render.R
@@ -56,7 +56,7 @@ create_prevelance_renderer <- function(
       )
     }
 
-    if (parameters$render_grid) {
+    if (!is.null(parameters$render_grid)) {
       grid_renderer(
         birth,
         renderer,
@@ -64,14 +64,31 @@ create_prevelance_renderer <- function(
         'n',
         timestep
       )
+    }
+
+    if ('n_detect_' %in% parameters$render_grid) {
       grid_renderer(
         birth,
         renderer,
         detected,
-        'n_detect',
+        'n_detect_',
         timestep
       )
     }
+    if ('p_detect_' %in% parameters$render_grid) {
+      weights <- rep(0, parameters$human_population)
+      weights[clinically_detected$to_vector()] <- 1
+      weights[asymptomatic$to_vector()] <- prob
+      weights <- weights[detected$to_vector()]
+      grid_renderer(
+        birth,
+        renderer,
+        detected,
+        'p_detect_',
+        timestep,
+        weights
+      )
+    }
   }
 }
 
@@ -120,15 +137,27 @@ incidence_renderer <- function(
     )
   }
 
-  if (render_grid) {
+  n_prefix = paste0('n_', prefix)
+  if (n_prefix %in% render_grid) {
     grid_renderer(
       birth,
       renderer,
       target,
-      'n_inc',
+      n_prefix,
       timestep
     )
   }
+  p_prefix = paste0('p_', prefix)
+  if (p_prefix %in% render_grid) {
+    grid_renderer(
+      birth,
+      renderer,
+      source_pop,
+      p_prefix,
+      timestep,
+      prob
+    )
+  }
 }
 
 create_variable_mean_renderer_process <- function(
@@ -217,19 +246,20 @@ grid_renderer <- function(
   renderer,
   target,
   prefix,
-  timestep
+  timestep,
+  weights = NULL
   ) {
-  counts <- grid_count(birth, target, timestep)
+  counts <- grid_count(birth, target, timestep, weights)
   for (i in seq_along(counts)) {
     renderer$render(
-      paste0('grid_', prefix, '_', i),
+      paste0('grid_', prefix, i),
       counts[[i]],
       timestep
     )
   }
 }
 
-grid_count <- function(birth, selected, timestep) {
+grid_count <- function(birth, selected, timestep, weights = NULL) {
   if (is.null(selected)) {
     selected_births <- birth$get_values()
   } else {
@@ -238,8 +268,16 @@ grid_count <- function(birth, selected, timestep) {
   age <- floor(get_age(selected_births, timestep) / 365)
   age[age < 0] <- NA
   age[age > 100] <- NA
-  non_zero <- table(age)
   counts <- rep(0, 101)
+  if (!is.null(weights)) {
+    if (length(weights) == 0) {
+      return(counts)
+    }
+    non_zero <- aggregate(x = weights, by = list(age = age), FUN = sum)
+    counts[non_zero$age + 1] <- non_zero$x
+    return(counts)
+  }
+  non_zero <- table(age)
   counts[as.numeric(names(non_zero)) + 1] <- non_zero
   counts
 }
diff --git a/tests/testthat/test-grid-render.R b/tests/testthat/test-grid-render.R
index 1e4bba27..bce2f1cd 100644
--- a/tests/testthat/test-grid-render.R
+++ b/tests/testthat/test-grid-render.R
@@ -23,6 +23,25 @@ test_that('grid_count counts internally correctly', {
   expect_equal(grid_count(birth, selected, timestep), expected)
 })
 
+test_that('grid_count does weighted counts correctly', {
+  timestep <- 5 * 365
+  birth <- individual::IntegerVariable$new(timestep - c(1, 1, 3, 4) * 365 - 1)
+  expected <- rep(0, 101)
+  weights <- c(.1, .2, .3, .4)
+  expected[c(1, 3, 4) + 1] <- c(.3, .3, .4)
+  expect_equal(grid_count(birth, NULL, timestep, weights), expected)
+})
+
+test_that('grid_count works for empty case', {
+  timestep <- 5 * 365
+  birth <- individual::IntegerVariable$new(timestep - c(1, 1, 3, 4) * 365 - 1)
+  expected <- rep(0, 101)
+  weights <- numeric(0)
+  selected <- individual::Bitset$new(4)
+  expect_equal(grid_count(birth, selected, timestep, weights), rep(0, 101))
+})
+
+
 test_that('grid_count counts at the boundaries correctly', {
   timestep <- 5 * 365
   birth <- individual::IntegerVariable$new(timestep - c(1, 2, 3, 4) * 365)