conmat-partial-preds-scale-vals.R

# Set-up ---------

library(conmat)
library(tidyverse)
library(gratia)
library(patchwork)

set.seed(2022 - 12 - 19)

# Obtain data from conmat ----------

# Extract setting-wise contact data
polymod_contact_data <- get_polymod_setting_data()
polymod_survey_data <- get_polymod_population()

# Fit the model
polymod_setting_models <- fit_setting_contacts(
  contact_data_list = polymod_contact_data,
  population = polymod_survey_data
)

p_home <- polymod_setting_models$home
p_work <- polymod_setting_models$work
p_school <- polymod_setting_models$school
p_other <- polymod_setting_models$other

# Home setting-specific -------------

#%% Partial dependency plots ------
# Of each covariate
pdp_home <- draw(p_home, residuals = TRUE) +
  plot_annotation(title = "Home setting")

pdp_home

ggsave(pdp_home,
       filename = "pdp_home.png",
       device = "png",
       height = 15, width = 20, units = "cm")

#%% Model summary -----
summary(p_home)

# new data
## TODO
## Wrap this up into a function that generates an age grid data frame
## with all the terms needed to fit a conmat model
## (from `fit_single_contact_model.R`)
library(tidyverse)
age_grid <- expand.grid(
  age_from = 1:90,
  age_to = 1:90
) |>
  as_tibble() |>
  # prepare the age data so it has all the right column names
  # that are used inside of `fit_single_contact_model()`
  conmat::add_symmetrical_features() |>
  # this ^^^ does the same as the commented part below:
  # dplyr::mutate(
  #   gam_age_offdiag = abs(age_from - age_to),
  #   gam_age_offdiag_2 = abs(age_from - age_to)^2,
  #   gam_age_diag_prod = abs(age_from * age_to),
  #   gam_age_diag_sum = abs(age_from + age_to),
  #   gam_age_pmax = pmax(age_from, age_to),
  #   gam_age_pmin = pmin(age_from, age_to)
  # ) |>
  # This is to add the school_probability and work_probability columns
  # that are used inside fit_single_contact_model() when fitting the model.
  conmat::add_modelling_features()

age_grid

library(mgcv)
age_grid

age_predictions <- age_grid |>
  mutate(
    ## TODO
    ## wrap up this "predict" code into a function that takes a
    ## "terms" argument, and we provide, e.g., "s(gam_age_offdiag)"
    ## and then create a data frame of predictions for each
    ## of the terms in conmat
    offdiag = predict(object = p_home,
                      newdata = age_grid,
                      type = "terms",
                      terms = "s(gam_age_offdiag)")
  ) |>
  relocate(
    gam_age_offdiag,
    offdiag
  )

## TODO
## You will need to pivot_longer the terms, similar to what Nick G wrote here:
## https://gist.github.com/njtierney/4ebf9536f92492b02398fa9e0f16576a

age_predictions

ggplot(age_predictions,
       aes(x = age_from,
           y = age_to,
           fill = offdiag)) +
  geom_tile() +
  scale_fill_viridis_c() +
  theme_minimal() +
  coord_fixed()

# tasks
#
# https://gist.github.com/njtierney/4ebf9536f92492b02398fa9e0f16576a

#%% Smooth estimates -----
sm_home <- smooth_estimates(p_home)

sm_home

dat_home <- sm_home %>%
  pivot_longer(
    cols = starts_with("gam_age"),
    names_prefix = "gam_age_",
    values_to = "x_value"
  ) %>%
  filter(!is.na(x_value)) %>%
  select(!c(`.type`, `.by`, `.smooth`)) %>%
  rename(estimate = `.estimate`, se = `.se`) %>%
  relocate(name)

## NOTE: It appears that just providing "terms" gives us quite
## different values to what we get if we specify each term individually above
conmat_terms <- predict(object = p_home,
                        newdata = age_grid,
                        type = "terms") |>
  as_tibble()

new_age_predictions <- bind_cols(
  age_grid,
  conmat_terms
)

new_age_predictions

ggplot(new_age_predictions,
       aes(x = age_from,
           y = age_to,
           fill = gam_age_offdiag)) +
  geom_tile() +
  scale_fill_viridis_c() +
  theme_minimal() +
  coord_fixed()


#TODO: Better to rename these covariates, especially "estimate" and "x_value"

# Save the smooth estimates
write_csv(dat_home, "./output/smooth-values-home.csv")