Skip to content

Instantly share code, notes, and snippets.

@njtierney
Created August 8, 2022 07:08
Show Gist options
  • Save njtierney/82e17e1512a511eb3ed537fdd3e6ebfd to your computer and use it in GitHub Desktop.
Save njtierney/82e17e1512a511eb3ed537fdd3e6ebfd to your computer and use it in GitHub Desktop.
library(mgcv)
#> Loading required package: nlme
#> This is mgcv 1.8-40. For overview type 'help("mgcv-package")'.
library(tidyverse)
library(ggplot2)
set.seed(2022-08-08)

# the fake function we would like to mirror
f <- function(age) {
  
  points <- tribble(
    ~x, ~y,
    0, 0.1,
    5, 0.1,
    10, 0.1,
    15, 0.1,
    20, 0.1,
    23, 0.5,
    24, 0.7,
    25, 0.9,
    27, 0.9,
    30, 0.9,
    35, 0.9,
    40, 0.9,
    45, 0.9,
    48, 0.9,
    50, 0.9,
    51, 0.7,
    52, 0.5,
    55, 0.1,
    60, 0.1,
    65, 0.1,
    70, 0.1,
    75, 0.1,
    80, 0.1
  )
  interp <- splinefun(x = points$x,
                   y = points$y,
                   method = "periodic")
  interp(age)
}

g <- function(age_from, age_to) {
  sin(age_from / 10 - 1) * sin(age_to / 10 - 1)
}

h <- function(age_from, age_to) {
  -0.1 * exp(sin(abs(age_from - age_to) / 5)) ^ 2
}

eta <- function(age_from, age_to) {
  1 + f(age_from) + f(age_to) + g(age_from, age_to) + h(age_from, age_to)
}

data <- expand.grid(
  age_from = 1:90,
  age_to = 1:90
) %>%
  mutate(
    # the fake log number of contacts, as if we had the model:
    # smooth(x, y) = f(x) + f(y)
    term_eta = eta(age_from, age_to),
    participants = 500,
    # fake data (ignore number of participants and population offsets for now)
    response = rpois(n(), exp(term_eta)) * participants
  )

# fit with various symmetric smooth types
m <- gam(response ~
           s(I(abs(age_from - age_to))) +
           s(I(abs(age_from - age_to)^2)) +
           s(I(abs(age_from * age_to))) +
           s(I(abs(age_from + age_to))) +
           s(I(pmax(age_from, age_to))) +
           s(I(pmin(age_from, age_to))),
         family = poisson,
         offset = log(participants),
         data = data)

# add terms back into the data frame
data <- data %>%
  mutate(
    term_intercept = m$coefficients[1],
    term_offdiag = predict(m, ., type = "terms", terms = "s(I(abs(age_from - age_to)))"),
    term_offdiag_2 = predict(m, ., type = "terms", terms = "s(I(abs(age_from - age_to)^2))"),
    term_diag_prod = predict(m, ., type = "terms", terms = "s(I(abs(age_from * age_to)))"),
    term_diag_sum = predict(m, ., type = "terms", terms = "s(I(abs(age_from + age_to)))"),
    term_max = predict(m, ., type = "terms", terms = "s(I(pmax(age_from, age_to)))"),
    term_min = predict(m, ., type = "terms", terms = "s(I(pmin(age_from, age_to)))"),
    term_fitted = term_intercept +
      term_offdiag +
      term_offdiag_2 +
      term_diag_prod +
      term_diag_sum +
      term_max +
      term_min
  )


# pivot_longer for plotting
data_terms <- data %>%
  pivot_longer(
    starts_with("term"),
    names_to = "term",
    values_to = "value",
    names_prefix = "term_"
  )

# putting eta and fitted in with the terms is a bit inelegant, but I couldn't
# think of a nicer pattern that wasn't a lot of code.

# plot the truth and fitted model
unique(data_terms$term)
#> [1] "eta"       "intercept" "offdiag"   "offdiag_2" "diag_prod" "diag_sum" 
#> [7] "max"       "min"       "fitted"
data_terms %>%
  filter(term %in% c("eta", "fitted")) %>%
  ggplot(
    aes(
      x = age_from,
      y = age_to,
      group = term,
      fill = exp(value)
    )
  ) +
  facet_wrap(~term) +
  geom_tile() +
  scale_fill_viridis_c() +
  theme_minimal() +
  coord_fixed()

# plot the parts of the model
data_terms %>%
  filter(!(term %in% c("eta", "fitted"))) %>%
  ggplot(
    aes(
      x = age_from,
      y = age_to,
      group = term,
      fill = value
    )
  ) +
  facet_wrap(~term) +
  geom_tile() +
  scale_fill_viridis_c() +
  theme_minimal() +
  coord_fixed()

# plot the 1D smooths
par(mfrow = c(2, 3))
plot(m)

par(mfrow = c(1, 1))

Created on 2022-08-08 by the reprex package (v2.0.1)

Session info
sessioninfo::session_info()
#> ─ Session info ───────────────────────────────────────────────────────────────
#>  setting  value
#>  version  R version 4.2.0 (2022-04-22)
#>  os       macOS Monterey 12.3.1
#>  system   aarch64, darwin20
#>  ui       X11
#>  language (EN)
#>  collate  en_AU.UTF-8
#>  ctype    en_AU.UTF-8
#>  tz       Australia/Perth
#>  date     2022-08-08
#>  pandoc   2.18 @ /Applications/RStudio.app/Contents/MacOS/quarto/bin/tools/ (via rmarkdown)
#> 
#> ─ Packages ───────────────────────────────────────────────────────────────────
#>  package     * version    date (UTC) lib source
#>  assertthat    0.2.1      2019-03-21 [1] CRAN (R 4.2.0)
#>  backports     1.4.1      2021-12-13 [1] CRAN (R 4.2.0)
#>  broom         0.8.0      2022-04-13 [1] CRAN (R 4.2.0)
#>  cellranger    1.1.0      2016-07-27 [1] CRAN (R 4.2.0)
#>  cli           3.3.0.9000 2022-06-15 [1] Github (r-lib/cli@31a5db5)
#>  colorspace    2.0-3      2022-02-21 [1] CRAN (R 4.2.0)
#>  crayon        1.5.1      2022-03-26 [1] CRAN (R 4.2.0)
#>  curl          4.3.2      2021-06-23 [1] CRAN (R 4.2.0)
#>  DBI           1.1.2      2021-12-20 [1] CRAN (R 4.2.0)
#>  dbplyr        2.1.1      2021-04-06 [1] CRAN (R 4.2.0)
#>  digest        0.6.29     2021-12-01 [1] CRAN (R 4.2.0)
#>  dplyr       * 1.0.9      2022-04-28 [1] CRAN (R 4.2.0)
#>  ellipsis      0.3.2      2021-04-29 [1] CRAN (R 4.2.0)
#>  evaluate      0.15       2022-02-18 [1] CRAN (R 4.2.0)
#>  fansi         1.0.3      2022-03-24 [1] CRAN (R 4.2.0)
#>  farver        2.1.0      2021-02-28 [1] CRAN (R 4.2.0)
#>  fastmap       1.1.0      2021-01-25 [1] CRAN (R 4.2.0)
#>  forcats     * 0.5.1      2021-01-27 [1] CRAN (R 4.2.0)
#>  fs            1.5.2      2021-12-08 [1] CRAN (R 4.2.0)
#>  generics      0.1.2      2022-01-31 [1] CRAN (R 4.2.0)
#>  ggplot2     * 3.3.6      2022-05-03 [1] CRAN (R 4.2.0)
#>  glue          1.6.2      2022-02-24 [1] CRAN (R 4.2.0)
#>  gtable        0.3.0      2019-03-25 [1] CRAN (R 4.2.0)
#>  haven         2.5.0      2022-04-15 [1] CRAN (R 4.2.0)
#>  highr         0.9        2021-04-16 [1] CRAN (R 4.2.0)
#>  hms           1.1.1      2021-09-26 [1] CRAN (R 4.2.0)
#>  htmltools     0.5.2      2021-08-25 [1] CRAN (R 4.2.0)
#>  httr          1.4.3      2022-05-04 [1] CRAN (R 4.2.0)
#>  jsonlite      1.8.0      2022-02-22 [1] CRAN (R 4.2.0)
#>  knitr         1.39       2022-04-26 [1] CRAN (R 4.2.0)
#>  labeling      0.4.2      2020-10-20 [1] CRAN (R 4.2.0)
#>  lattice       0.20-45    2021-09-22 [1] CRAN (R 4.2.0)
#>  lifecycle     1.0.1      2021-09-24 [1] CRAN (R 4.2.0)
#>  lubridate     1.8.0      2021-10-07 [1] CRAN (R 4.2.0)
#>  magrittr      2.0.3      2022-03-30 [1] CRAN (R 4.2.0)
#>  Matrix        1.4-1      2022-03-23 [1] CRAN (R 4.2.0)
#>  mgcv        * 1.8-40     2022-03-29 [1] CRAN (R 4.2.0)
#>  mime          0.12       2021-09-28 [1] CRAN (R 4.2.0)
#>  modelr        0.1.8      2020-05-19 [1] CRAN (R 4.2.0)
#>  munsell       0.5.0      2018-06-12 [1] CRAN (R 4.2.0)
#>  nlme        * 3.1-157    2022-03-25 [1] CRAN (R 4.2.0)
#>  pillar        1.7.0      2022-02-01 [1] CRAN (R 4.2.0)
#>  pkgconfig     2.0.3      2019-09-22 [1] CRAN (R 4.2.0)
#>  purrr       * 0.3.4      2020-04-17 [1] CRAN (R 4.2.0)
#>  R.cache       0.15.0     2021-04-30 [1] CRAN (R 4.2.0)
#>  R.methodsS3   1.8.1      2020-08-26 [1] CRAN (R 4.2.0)
#>  R.oo          1.24.0     2020-08-26 [1] CRAN (R 4.2.0)
#>  R.utils       2.11.0     2021-09-26 [1] CRAN (R 4.2.0)
#>  R6            2.5.1      2021-08-19 [1] CRAN (R 4.2.0)
#>  readr       * 2.1.2      2022-01-30 [1] CRAN (R 4.2.0)
#>  readxl        1.4.0      2022-03-28 [1] CRAN (R 4.2.0)
#>  reprex        2.0.1      2021-08-05 [1] CRAN (R 4.2.0)
#>  rlang         1.0.4      2022-07-12 [1] CRAN (R 4.2.0)
#>  rmarkdown     2.14       2022-04-25 [1] CRAN (R 4.2.0)
#>  rstudioapi    0.13       2020-11-12 [1] CRAN (R 4.2.0)
#>  rvest         1.0.2      2021-10-16 [1] CRAN (R 4.2.0)
#>  scales        1.2.0      2022-04-13 [1] CRAN (R 4.2.0)
#>  sessioninfo   1.2.2      2021-12-06 [1] CRAN (R 4.2.0)
#>  stringi       1.7.6      2021-11-29 [1] CRAN (R 4.2.0)
#>  stringr     * 1.4.0      2019-02-10 [1] CRAN (R 4.2.0)
#>  styler        1.7.0      2022-03-13 [1] CRAN (R 4.2.0)
#>  tibble      * 3.1.7      2022-05-03 [1] CRAN (R 4.2.0)
#>  tidyr       * 1.2.0      2022-02-01 [1] CRAN (R 4.2.0)
#>  tidyselect    1.1.2      2022-02-21 [1] CRAN (R 4.2.0)
#>  tidyverse   * 1.3.1      2021-04-15 [1] CRAN (R 4.2.0)
#>  tzdb          0.3.0      2022-03-28 [1] CRAN (R 4.2.0)
#>  utf8          1.2.2      2021-07-24 [1] CRAN (R 4.2.0)
#>  vctrs         0.4.1      2022-04-13 [1] CRAN (R 4.2.0)
#>  viridisLite   0.4.0      2021-04-13 [1] CRAN (R 4.2.0)
#>  withr         2.5.0      2022-03-03 [1] CRAN (R 4.2.0)
#>  xfun          0.31       2022-05-10 [1] CRAN (R 4.2.0)
#>  xml2          1.3.3      2021-11-30 [1] CRAN (R 4.2.0)
#>  yaml          2.3.5      2022-02-21 [1] CRAN (R 4.2.0)
#> 
#>  [1] /Library/Frameworks/R.framework/Versions/4.2-arm64/Resources/library
#> 
#> ──────────────────────────────────────────────────────────────────────────────
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment