Incorporate prediction capabilities

[csmith701]

Currently model supports estimation runs. Need to expand to allow for prediction run.

• Ability for user to specify estimation vs. prediction run. @pearsonca thoughts on whether this should be an input directly to imugap or part of imugap_options?
• Means of saving data input for estimation run (it is not in the stanfit object)
• Reformat data for prediction run (use script I sent to Matt H as starting point)

[csmith701]

@pearsonca Here is the script I made for Matt:

library(dplyr)
library(data.table)
library(rstan)
library(imuGAP)

# Path to output from original imugap estimation run
est_path <- "your_path"

# Birth year corresponding to the oldest cohort
oldest <- 2008

sm <- stan_model(file = "inst/stan/impute_school_coverage_process_v5.stan")
res_stan <- readRDS(est_path)

# Read in input files Matt sent
loc_ids <- read.csv("SampleLocationIDs.csv")
obs <- read.csv("SampleObservations.csv")
obs_pop <- read.csv("SampleObsPopulation.csv")

# Get the original stan data, based on above inputs
observations <- obs %>% select(sample_n, positive, id)
obs_populations <- obs_pop %>%
  select(obs_id, location, cohort, age, dose, weight) %>%
  mutate(cohort = cohort-min(cohort) + 1)
locations <-  loc_ids %>% select(id, parent_id)
dose_schedule <- c(1, 4)

# check location argument
loc_info <- check_locations(locations)
n_cnty <- loc_info[layer == 2, .N]
n_schl <- loc_info[layer == 3, .N]

# check observations argument
obs <- check_observations(observations)[, .(id, positive, sample_n)]

# check obs_populations - confirm wts locations
wts <- check_obs_population(
  obs_populations, obs, loc_info
)

bsp <- splines::bs(
  seq_len(wts[, diff(range(cohort)) + 1L]),
  df = 5, intercept = TRUE
)

doses <- matrix(0, ncol = length(dose_schedule), nrow = max(wts$age))
for (i in seq_along(dose_schedule)) {
  doses[(dose_schedule[i] + 1):nrow(doses), i] <- 1
}

# prepare dat_stan
dat_stan <- list(
  n_yr = max(wts$age),
  n_cohort = max(wts$cohort),
  n_sch = n_schl,
  n_doses = length(dose_schedule),
  dose_sched = doses,
  k_bs = ncol(bsp),
  bs = bsp,
  n_obs = nrow(obs),
  y_obs = obs$positive,
  y_smp = obs$sample_n,
  n_weights = nrow(wts),
  obs_to_weights_bounds = unique(wts$range_start),
  weights_school = wts$location,
  weights_cohort = wts$cohort,
  weights_life_year = wts$age,
  weights_dose = wts$dose,
  weights = wts$weight,
  n_cnty = n_cnty,
  cnty_bounds = loc_info[layer == 3, unique(layer_bound)],
  predict_mode = 0
)

# Create weights for prediction runs
tract_ids <- loc_ids$id[loc_ids$type == "CensusTract"]
cohorts <- sort(unique(obs_pop$cohort))
cohorts_norm <- cohorts - min(cohorts) + 1


wts_pred <- data.frame(wts_school = rep(tract_ids, each = length(cohorts)), # Each census tract listed once per cohort
                       wts_cohort = rep(cohorts_norm, length(tract_ids))) %>% # 1:n_cohorts for each census tract
  mutate(wts_life_year = 2026 - (wts_cohort + min(cohorts) - 1), # Current age for each cohort
         wts_dose = ifelse(wts_life_year <= 4, 1, 2), # 1 dose coverage if under 4, 2 dose otherwise
         wts = 1) # Always talking about just one cohort and census tract so all weights are 1

# Update stan data object for prediction run
dat_stan_pred <- dat_stan
dat_stan_pred$n_yr <- max(wts_pred$wts_life_year)
dat_stan_pred$n_obs <- nrow(wts_pred)
dat_stan_pred$y_obs <- rep(1, nrow(wts_pred)) # don't get used for the non-inference version
dat_stan_pred$y_smp <- rep(1, nrow(wts_pred)) # don't get used for the non-inference version
dat_stan_pred$n_weights <- nrow(wts_pred)
dat_stan_pred$obs_to_weights_bounds <- 1:nrow(wts_pred)
dat_stan_pred$weights_school <- wts_pred$wts_school
dat_stan_pred$weights_cohort <- wts_pred$wts_cohort
dat_stan_pred$weights_life_year <- wts_pred$wts_life_year
dat_stan_pred$weights_dose <- wts_pred$wts_dose
dat_stan_pred$weights <- wts_pred$wts
dat_stan_pred$predict_mode <- 1 # Indicate that this is a prediction run

# Make predictions via generated quantities
# The predicted values will be called p_gen in the resulting stanfit object
preds <- gqs(sm, data = dat_stan_pred, as.matrix(res_stan))

[csmith701]

The example inputs this code is referencing:

SampleLocationIDs.csv
SampleObservations.csv
SampleObsPopulation.csv

[pearsonca]

rough plan:

• update entry points to sampling, extract, and predict
• predict should receive a population object (+ fit from sampling)
• stretch: provide a convenience builder for population
• double check the stan internals

[pearsonca]

now just updating package data to reflect soup-to-nuts.

[pearsonca]

Also: once satisfied with additional package data, update @examples for predict.imugap_fit to use the example data.