tests/testthat/test-bootstrap.R

In TrialEmulation: Causal Analysis of Observational Time-to-Event Data

test_that("Same weights recalculated if we use the same weight model coefficients and original dataset", {
  skip_on_cran()
  set.seed(887)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))


  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model")),
      adjustment_terms = ~x2
    ) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")


  result <- weight_func_bootstrap(
    object = trial_pp, remodel = FALSE, quiet = TRUE, boot_idx = unique(trial_pp@data@data$id),
    new_coef_sw_d0 = trial_pp@switch_weights@fitted$d0@summary$tidy$estimate,
    new_coef_sw_d1 = trial_pp@switch_weights@fitted$d1@summary$tidy$estimate,
    new_coef_sw_n0 = trial_pp@switch_weights@fitted$n0@summary$tidy$estimate,
    new_coef_sw_n1 = trial_pp@switch_weights@fitted$n1@summary$tidy$estimate,
    new_coef_c_d0 = trial_pp@censor_weights@fitted$d0@summary$tidy$estimate,
    new_coef_c_d1 = trial_pp@censor_weights@fitted$d1@summary$tidy$estimate,
    new_coef_c_n1 = trial_pp@censor_weights@fitted$n1@summary$tidy$estimate,
    new_coef_c_n0 = trial_pp@censor_weights@fitted$n0@summary$tidy$estimate
  )

  expect_equal(result$data$weight, trial_pp@outcome_data@data$weight)
})

test_that("Same weights refitted if we use original dataset", {
  skip_on_cran()
  set.seed(222)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  result <- weight_func_bootstrap(
    object = trial_pp,
    remodel = TRUE,
    quiet = TRUE,
    boot_idx = unique(trial_pp@data@data$id)
  )

  expect_equal(result$data$weight, trial_pp@outcome_data@data$weight)
})

test_that("Same weights refitted if we use example bootstrap sample", {
  skip_on_cran()
  set.seed(194)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  boot_idx <- sort(sample(unique(trial_pp@data@data$id), replace = TRUE))

  extract_idx <- lapply(boot_idx, function(i) which(data_censored$id == i))
  bootstrap_sample <- data_censored[unlist(extract_idx), ]
  bootstrap_sample$id <- rep(seq_along(extract_idx), lengths(extract_idx))


  trial_pp_boot <- trial_sequence(estimand = "PP") |>
    set_data(
      data = bootstrap_sample,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  result <- weight_func_bootstrap(object = trial_pp, remodel = TRUE, quiet = TRUE, boot_idx = boot_idx)
  result$data <- result$data[order(id, trial_period, followup_time)]

  expect_equal(trial_pp_boot@switch_weights@fitted$d0@summary$tidy$estimate,
    result$switch_models$switch_d0$summary$estimate,
    tolerance = 1e-7
  )
  expect_equal(trial_pp_boot@switch_weights@fitted$d1@summary$tidy$estimate,
    result$switch_models$switch_d1$summary$estimate,
    tolerance = 1e-7
  )
  expect_equal(trial_pp_boot@switch_weights@fitted$n0@summary$tidy$estimate,
    result$switch_models$switch_n0$summary$estimate,
    tolerance = 1e-7
  )
  expect_equal(trial_pp_boot@switch_weights@fitted$n1@summary$tidy$estimate,
    result$switch_models$switch_n1$summary$estimate,
    tolerance = 1e-7
  )

  expect_equal(trial_pp_boot@censor_weights@fitted$d0@summary$tidy$estimate,
    result$censor_models$cens_d0$summary$estimate,
    tolerance = 1e-7
  )
  expect_equal(trial_pp_boot@censor_weights@fitted$d1@summary$tidy$estimate,
    result$censor_models$cens_d1$summary$estimate,
    tolerance = 1e-7
  )
  expect_equal(trial_pp_boot@censor_weights@fitted$n0@summary$tidy$estimate,
    result$censor_models$cens_n0$summary$estimate,
    tolerance = 1e-7
  )
  expect_equal(trial_pp_boot@censor_weights@fitted$n1@summary$tidy$estimate,
    result$censor_models$cens_n1$summary$estimate,
    tolerance = 1e-7
  )

  test_data <- trial_pp_boot@outcome_data@data
  test_data$id <- boot_idx[test_data$id]
  test_data <- test_data[!duplicated(test_data), ]
  test_data$weight_boot <- sapply(test_data$id, function(i) sum(boot_idx == i))
  test_data$weight <- test_data$weight * test_data$weight_boot
  test_data <- merge(
    x = trial_pp@outcome_data@data[, c("id", "trial_period", "followup_time")],
    y = test_data,
    all.x = TRUE
  )
  test_data$weight[is.na(test_data$weight)] <- 0
  test_data <- test_data[order(id, trial_period, followup_time)]

  expect_equal(result$data$weight, test_data$weight, tolerance = 1e-8)
})

test_that("Correct weights recalculated if we use new weight model coefficients and bootstrap sample", {
  skip_on_cran()
  set.seed(978)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")
  boot_idx <- sort(sample(unique(trial_pp@data@data$id), replace = TRUE))

  extract_idx <- lapply(boot_idx, function(i) which(data_censored$id == i))
  bootstrap_sample <- data_censored[unlist(extract_idx), ]
  bootstrap_sample$id <- rep(seq_along(extract_idx), lengths(extract_idx))

  trial_pp_boot <- trial_sequence(estimand = "PP") |>
    set_data(
      data = bootstrap_sample,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  result <- weight_func_bootstrap(
    object = trial_pp, remodel = FALSE, quiet = TRUE, boot_idx = boot_idx,
    new_coef_sw_d0 = trial_pp_boot@switch_weights@fitted$d0@summary$tidy$estimate,
    new_coef_sw_d1 = trial_pp_boot@switch_weights@fitted$d1@summary$tidy$estimate,
    new_coef_sw_n0 = trial_pp_boot@switch_weights@fitted$n0@summary$tidy$estimate,
    new_coef_sw_n1 = trial_pp_boot@switch_weights@fitted$n1@summary$tidy$estimate,
    new_coef_c_d0 = trial_pp_boot@censor_weights@fitted$d0@summary$tidy$estimate,
    new_coef_c_d1 = trial_pp_boot@censor_weights@fitted$d1@summary$tidy$estimate,
    new_coef_c_n1 = trial_pp_boot@censor_weights@fitted$n1@summary$tidy$estimate,
    new_coef_c_n0 = trial_pp_boot@censor_weights@fitted$n0@summary$tidy$estimate
  )

  result$data <- result$data[order(id, trial_period, followup_time)]


  test_data <- trial_pp_boot@outcome_data@data
  test_data$id <- boot_idx[test_data$id]
  test_data <- test_data[!duplicated(test_data), ]
  test_data$weight_boot <- sapply(test_data$id, function(i) sum(boot_idx == i))
  test_data$weight <- test_data$weight * test_data$weight_boot
  test_data <- merge(
    x = trial_pp@outcome_data@data[, c("id", "trial_period", "followup_time")],
    y = test_data,
    all.x = TRUE
  )
  test_data$weight[is.na(test_data$weight)] <- 0
  test_data <- test_data[order(id, trial_period, followup_time)]

  expect_equal(result$data$weight, test_data$weight, tolerance = 1e-7)
})

test_that("no bootstrap with ITT", {
  skip_on_cran()
  set.seed(194)
  trial_itt_dir <- withr::local_tempdir("trial_itt", tempdir(TRUE))

  trial_itt <- trial_sequence(estimand = "ITT") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~ age_s + x4,
      denominator = ~ age_s + x4 + x2 + x1,
      pool_models = "both",
      model_fitter = stats_glm_logit(save_path = file.path(trial_itt_dir, "switch_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_itt_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data()

  suppressWarnings(trial_itt <- fit_msm(trial_itt, weight_cols = c("weight")))

  expect_error(predict(trial_itt, predict_times = 1:5, ci_type = "Nonpara. bootstrap"),
    fixed = TRUE,
    regexp = "Bootstrap and Jackknife confidence intervals are only implemented for trial_sequence_PP class."
  )
})

test_that("predict works with bootstrap", {
  skip_on_cran()
  set.seed(194)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  suppressWarnings(result2 <- predict(trial_pp, predict_times = 1:20, ci_type = "Nonpara. bootstrap"))
  expect_snapshot_value(
    as.data.frame(result2),
    style = "json2",
    tolerance = 1e-6
  )
  # TODO check if there is there is some overwriting going on due to data.table, so the results would not be
  # reproducible if we calculate different methods.
})

test_that("predict works with bootstrap with newdata containing ID", {
  skip_on_cran()
  set.seed(194)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  newdata <- trial_pp@outcome_model@fitted@model$model$data
  newdata <- newdata[newdata$trial_period == 0, ]
  suppressWarnings(result2 <- predict(trial_pp, newdata = newdata, predict_times = 1:5, ci_type = "Nonpara. bootstrap"))
  expect_snapshot_value(
    as.data.frame(result2),
    style = "json2",
    tolerance = 1e-6
  )
  # TODO check if there is there is some overwriting going on due to data.table, so the results would not be
  # reproducible if we calculate different methods.
})

test_that("predict works with bootstrap with newdata NOT containing ID", {
  skip_on_cran()
  set.seed(194)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  newdata <- as.data.frame(trial_pp@outcome_model@fitted@model$model$data)
  newdata <- newdata[newdata$trial_period == 0, names(newdata) != "id"]
  suppressWarnings(result2 <- predict(trial_pp, newdata = newdata, predict_times = 1:5, ci_type = "Nonpara. bootstrap"))
  expect_snapshot_value(
    as.data.frame(result2),
    style = "json2",
    tolerance = 1e-6
  )
  # TODO check if there is there is some overwriting going on due to data.table, so the results would not be
  # reproducible if we calculate different methods.
})

test_that("predict works with LEF outcome", {
  skip_on_cran()
  set.seed(194)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  suppressWarnings(result2 <- predict(trial_pp, predict_times = 1:20, ci_type = "LEF outcome"))

  expect_snapshot_value(
    as.data.frame(result2),
    style = "json2",
    tolerance = 1e-6
  )

  # TODO check if there is there is some overwriting going on due to data.table, so the results would not be
  # reproducible if we calculate different methods.
})

test_that("predict works with LEF both", {
  skip_on_cran()
  set.seed(194)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  suppressWarnings(result2 <- predict(trial_pp, predict_times = 1:20, ci_type = "LEF both"))
  expect_snapshot_value(
    as.data.frame(result2),
    style = "json2",
    tolerance = 1e-6
  )

  # TODO check if there is there is some overwriting going on due to data.table, so the results would not be
  # reproducible if we calculate different methods.
})

test_that("predict works with Jackknife Wald", {
  skip_on_cran()
  set.seed(194)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  suppressWarnings(result2 <- predict(trial_pp, predict_times = 1:20, ci_type = "Jackknife Wald"))
  expect_snapshot_value(
    as.data.frame(result2),
    style = "json2",
    tolerance = 1e-6
  )

  # TODO check if there is there is some overwriting going on due to data.table, so the results would not be
  # reproducible if we calculate different methods.
})

test_that("predict works with Jackknife MVN", {
  skip_on_cran()
  set.seed(194)
  trial_pp_dir <- withr::local_tempdir("trial_pp", tempdir(TRUE))

  trial_pp <- trial_sequence(estimand = "PP") |>
    set_data(
      data = data_censored,
      id = "id",
      period = "period",
      treatment = "treatment",
      outcome = "outcome",
      eligible = "eligible"
    ) |>
    set_switch_weight_model(
      numerator = ~age,
      denominator = ~ age + x1 + x3,
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "switch_models"))
    ) |>
    set_censor_weight_model(
      censor_event = "censored",
      numerator = ~x2,
      denominator = ~ x2 + x1,
      pool_models = "none",
      model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "censor_models"))
    ) |>
    calculate_weights() |>
    set_outcome_model(model_fitter = stats_glm_logit(save_path = file.path(trial_pp_dir, "outcome_model"))) |>
    set_expansion_options(
      output = save_to_datatable(),
      chunk_size = 500
    ) |>
    expand_trials() |>
    load_expanded_data() |>
    fit_msm(
      weight_cols = c("weight")
    ) |>
    suppressMatchingWarnings("fitted probabilities")

  suppressWarnings(result2 <- predict(trial_pp, predict_times = 1:20, ci_type = "Jackknife MVN"))
  expect_snapshot_value(
    as.data.frame(result2),
    style = "json2",
    tolerance = 1e-6
  )

  # TODO check if there is there is some overwriting going on due to data.table, so the results would not be
  # reproducible if we calculate different methods.
})