Merge pull request #34 from epiforecasts/26-focus-25A-emergence

Issue 26: add a plot focused on 25A emergence

Author: kaitejohnson

R/fig_zoom_clade_mult_nowcasts.R

@@ -0,0 +1,202 @@
+#' Get a plot of nowcasts for multiple nowcast dates vs observations
+#'
+#' @param model_preds_mult_nowcasts Model predicted observations for multiple
+#'   nowcasts
+#' @param final_eval_data Final sequence counts by clade and location
+#' @param clade_to_zoom Character string indicating the clade to zoom
+#' @param horizon_to_plot Range of integers to filter to for the plot
+#'
+#' @returns ggplot object
+get_plot_model_preds_mult <- function(model_preds_mult_nowcasts,
+                                      final_eval_data,
+                                      clade_to_zoom = "25A",
+                                      horizon_to_plot = c(-6, 0)) {
+  df_filt <- model_preds_mult_nowcasts |>
+    filter(clade == clade_to_zoom) |>
+    mutate(horizon = as.integer(target_date - nowcast_date)) |>
+    filter(horizon <= max(horizon_to_plot), horizon >= min(horizon_to_plot))
+
+  weekly_obs_data <- daily_to_weekly(final_eval_data) |>
+    filter(location %in% unique(df_filt$location))
+  total_seq <- weekly_obs_data |>
+    group_by(date, location) |>
+    summarise(n_seq = sum(sequences))
+  weekly_obs <- left_join(weekly_obs_data, total_seq) |>
+    filter(clades_modeled == clade_to_zoom)
+
+  plot_comps <- plot_components()
+
+  p <- ggplot(df_filt) +
+    geom_line(aes(
+      x = target_date, y = q_0.5, color = model_id,
+      group = nowcast_date
+    )) +
+    geom_ribbon(
+      aes(
+        x = target_date,
+        ymin = q_0.25,
+        ymax = q_0.75, fill = model_id,
+        group = nowcast_date
+      ),
+      alpha = 0.2
+    ) +
+    geom_ribbon(
+      aes(
+        x = target_date,
+        ymin = q_0.025,
+        ymax = q_0.975, fill = model_id,
+        group = nowcast_date
+      ),
+      alpha = 0.1
+    ) +
+    geom_point(
+      data = weekly_obs,
+      aes(x = date, y = sequences / n_seq),
+      color = "#CAB2D6"
+    ) +
+    geom_line(
+      data = weekly_obs,
+      aes(x = date, y = sequences / n_seq),
+      color = "#CAB2D6"
+    ) +
+    facet_grid(vars(model_id), vars(location)) +
+    get_plot_theme(dates = TRUE) +
+    scale_color_manual(
+      name = "Model",
+      values = plot_comps$model_colors
+    ) +
+    scale_fill_manual(
+      name = "Model",
+      values = plot_comps$model_colors
+    ) +
+    xlab("") +
+    ylab("Model predictions across nowcast dates") +
+    guides(
+      color = "none",
+      fill = "none"
+    ) +
+    scale_x_date(
+      limits = c(min(df_filt$target_date), max(df_filt$target_date)),
+      date_breaks = "1 week",
+      date_labels = "%d %b %Y"
+    ) +
+    ggtitle("25A emergence")
+
+  return(p)
+}
+
+#' Get a plot of scores by nowcast date for locations and models
+#'
+#' @param scores Data.frame of scores
+#' @param locs Vector of character strings of locations
+#' @param nowcast_dates Set of nowcast dates to summarise over
+#' @param date_range Range of dates to plot
+#' @param horizon_to_plot horizon days to plot
+#'
+#' @returns ggplot
+get_plot_scores_by_date <- function(scores,
+                                    locs,
+                                    nowcast_dates,
+                                    date_range,
+                                    horizon_to_plot = c(-6, 0)) {
+  scores_avg <- scores |>
+    filter(
+      location %in% locs,
+      nowcast_date %in% nowcast_dates,
+      !is.na(energy_score)
+    ) |>
+    group_by(model, location) |>
+    summarise(energy_score = mean(energy_score, na.rm = TRUE))
+
+  scores_df <- scores |>
+    mutate(horizon = as.numeric(target_date - nowcast_date)) |>
+    filter(
+      location %in% locs,
+      nowcast_date %in% nowcast_dates,
+      !is.na(energy_score),
+      horizon <= max(horizon_to_plot),
+      horizon >= min(horizon_to_plot)
+    ) |>
+    group_by(nowcast_date, location, model) |>
+    summarise(energy_score = mean(energy_score, na.rm = TRUE))
+
+  plot_comps <- plot_components()
+  p <- ggplot(scores_df) +
+    geom_point(aes(
+      x = nowcast_date, y = energy_score,
+      color = model
+    )) +
+    geom_line(aes(
+      x = nowcast_date, y = energy_score,
+      color = model
+    )) +
+    geom_hline(
+      data = scores_avg,
+      aes(yintercept = energy_score, color = model),
+      linetype = "dashed"
+    ) +
+    facet_wrap(~location) +
+    get_plot_theme(dates = TRUE) +
+    scale_color_manual(
+      name = "Model",
+      values = plot_comps$model_colors
+    ) +
+    xlab("") +
+    ylab("Average energy score") +
+    scale_x_date(
+      limits = date_range,
+      date_breaks = "1 week",
+      date_labels = "%d %b %Y"
+    )
+  return(p)
+}
+
+#' Multiplot panel looking at 25A emergence across nowcast dates
+#'
+#' @param grid A
+#' @param underlay B
+#' @param plot_name name of plot
+#' @param output_fp filepath directory
+#'
+#' @returns patchwork
+get_fig_zoom_25A <- function(grid,
+                             underlay,
+                             plot_name,
+                             output_fp = file.path(
+                               "output", "figs",
+                               "zoom_25A", "final"
+                             )) {
+  fig_layout <- "
+  AAA
+  AAA
+  AAA
+  BBB
+  "
+
+  fig_zoom <- grid +
+    underlay +
+    plot_layout(
+      design = fig_layout,
+      axes = "collect",
+      guides = "collect"
+    ) +
+    plot_annotation(
+      tag_levels = "A",
+      tag_suffix = "",
+      tag_sep = "",
+      theme = theme(
+        legend.position = "top",
+        legend.title = element_text(hjust = 0.5),
+        legend.justification = "center",
+        plot.tag = element_text(size = 20)
+      )
+    )
+
+  dir_create(output_fp, recurse = TRUE)
+  ggsave(file.path(output_fp, glue::glue("{plot_name}.png")),
+    plot = fig_zoom,
+    width = 8,
+    height = 11
+  )
+  return(fig_zoom)
+}

_targets.R

@@ -83,6 +83,7 @@ plot_targets <- list(
   fig_nowcast_targets,
   fig_pred_plus_data_targets,
   fig_overall_targets,
+  fig_zoom_25A_targets,
   fig_submission_heatmaps_targets
 )
 

man/get_fig_zoom_25A.Rd

@@ -31,6 +31,17 @@ clean_data_targets <- list(
       type = "evaluation"
     )
   ),
+  tar_target(
+    name = clean_variant_data_for_eval_mult_nowcasts,
+    command = get_clean_variant_data(
+      variant_data_for_eval_mult_nowcasts,
+      clade_list,
+      location_data,
+      nowcast_date_range_to_zoom,
+      seq_col_name = "oracle_value",
+      type = "evaluation"
+    )
+  ),
   tar_target(
     name = clean_variant_data_eval_all,
     command = get_clean_variant_data(

man/get_plot_model_preds_mult.Rd

@@ -55,6 +55,14 @@ config_targets <- list(
     name = nowcast_date_for_vis,
     command = as.character(ymd("2025-02-19"))
   ),
+  tar_target(
+    name = nowcast_date_range_to_zoom,
+    command = seq(
+      from = ymd("2025-01-15"),
+      to = ymd("2025-03-19"),
+      by = "week"
+    )
+  ),
   tar_target(
     name = quantiles_to_vis,
     command = c(0.05, 0.25, 0.5, 0.75, 0.975)

man/get_plot_scores_by_date.Rd

@@ -0,0 +1,30 @@
+fig_zoom_25A_targets <- list(
+  tar_target(
+    name = plot_model_preds_mult_nowcasts,
+    command = get_plot_model_preds_mult(
+      model_preds_mult_nowcasts = df_mult_nowcasts,
+      final_eval_data = clean_variant_data_final_all_states,
+      clade_to_zoom = "25A"
+    )
+  ),
+  tar_target(
+    name = plot_score_underlay,
+    command = get_plot_scores_by_date(
+      scores = su_scores,
+      locs = states_for_vis,
+      nowcast_dates = nowcast_date_range_to_zoom,
+      date_range = c(
+        min(nowcast_date_range_to_zoom) - days(6),
+        max(nowcast_date_range_to_zoom)
+      )
+    )
+  ),
+  tar_target(
+    name = fig_zoom_25A,
+    command = get_fig_zoom_25A(
+      plot_model_preds_mult_nowcasts,
+      plot_score_underlay,
+      plot_name = "fig_zoom_25A"
+    )
+  )
+)

targets/clean_data_targets.R

@@ -22,6 +22,14 @@ load_data_targets <- list(
       states = states_for_vis
     )
   ),
+  tar_target(
+    name = variant_data_for_eval_mult_nowcasts,
+    command = get_oracle_output(
+      hub_path = hub_path,
+      nowcast_dates = nowcast_date_range_to_zoom,
+      states = states_for_vis
+      )
+    ),
   # Variant data for evaluation (all dates and locations)
   tar_target(
     name = variant_data_eval_all,
@@ -76,6 +84,15 @@ load_data_targets <- list(
       bucket_name = nowcast_bucket_name
     )
   ),
+  # Model outputs for the selected dates
+  tar_target(
+    name = model_outputs_mult_nowcasts,
+    command = extract_nowcasts(
+      nowcast_dates = nowcast_date_range_to_zoom,
+      states = states_for_vis,
+      bucket_name = nowcast_bucket_name
+      )
+    ),
   # All model outputs for heatmap (all dates and locations)
   tar_target(
     name = all_model_outputs_for_heatmap,