inst/examples/test_functional_outcome.R

In FPScausal: Functional Propensity Score for Causal Inference

# =============================================================================
# FPScausal - Full pipeline test: FUNCTIONAL outcome
# =============================================================================

library(FPScausal)

cat("============================================================\n")
cat("FPScausal - Functional Outcome Pipeline Test\n")
cat("============================================================\n\n")

# ---- 1. Simulate data ----
cat("Step 1: Simulate data (n=200, setting='LL', outcome='functional')\n")

dat <- simulate_fps_data(
  n                      = 2000,
  setting                = "LL",
  outcome_type           = "functional",
  include_functional_cov = TRUE,
  seed                   = 99
)

cat("  Treatment X  :", nrow(dat$X), "x", ncol(dat$X), "\n")
cat("  Outcome Y    :", nrow(dat$Y), "x", ncol(dat$Y), "\n")
cat("  Confounders  :", nrow(dat$C), "x", ncol(dat$C), "\n")
cat("  Functional D :", nrow(dat$D), "x", ncol(dat$D), "\n")
cat("  True mu(s,t) :", nrow(dat$true_beta), "x", ncol(dat$true_beta), "\n\n")

# ---- 2. Estimate weights ----
cat("Step 2: Estimate FPS weights\n")

w_fn <- fps_weighting(
  treatment   = dat$X,
  t_grid      = dat$t_grid,
  domain      = c(0, 1),
  domain_name = "s",
  pve         = 0.95,
  covariates  = list(
    scalar     = dat$C,
    functional = list(dat$D)
  ),
  cov_t_grids = list(dat$t_grid),
  cov_domains = list(c(0, 1))
)

print(w_fn)
print(plot(w_fn, type = "balance"))
print(plot(w_fn, type = "weights"))
print(plot(w_fn, type = "fpca_treatment"))

# ---- 3. Effect estimation (no bootstrap) ----
cat("\nStep 3: Estimate causal effect surface (no bootstrap)\n")

eff_fn <- fps_effect_estimation(
  outcome             = dat$Y,
  fps_object          = w_fn,
  outcome_t_grid      = dat$t_grid,
  outcome_domain      = c(0, 1),
  outcome_domain_name = "t",
  outcome_pve         = 0.95,
  true_beta           = dat$true_beta
)

print(eff_fn)
summary(eff_fn)
print(plot(eff_fn, type = "effect"))
print(plot(eff_fn, type = "comparison"))
print(plot(eff_fn, type = "fpca_treatment"))
print(plot(eff_fn, type = "fpca_outcome"))

# ---- 4. Effect estimation with bootstrap ----
cat("\nStep 4: Bootstrap CIs (B=200)\n")

eff_fn_boot <- fps_effect_estimation(
  outcome             = dat$Y,
  fps_object          = w_fn,
  outcome_t_grid      = dat$t_grid,
  outcome_domain      = c(0, 1),
  outcome_domain_name = "t",
  outcome_pve         = 0.95,
  bootstrap           = TRUE,
  B                   = 200,
  alpha               = 0.05,
  true_beta           = dat$true_beta,
  seed                = 42
)

print(plot(eff_fn_boot, type = "significance"))

# Slices at fixed treatment times s -- all three panels side by side
cat("  Slices fixing treatment time s = 0.2, 0.5, 0.8\n")
print(
  plot(eff_fn_boot, type = "bootstrap_slice",
       point        = c(0.2, 0.5, 0.8),
       which_domain = "treatment")
)

# Slices at fixed outcome times t -- all three panels side by side
cat("  Slices fixing outcome time t = 0.2, 0.5, 0.8\n")
print(
  plot(eff_fn_boot, type = "bootstrap_slice",
       point        = c(0.2, 0.5, 0.8),
       which_domain = "outcome")
)

# ---- 5. All four settings ----
cat("\nStep 5: All four simulation settings\n")

settings <- c("LL", "LN", "NL", "NN")

beta_results_fn <- lapply(settings, function(s) {
  cat(sprintf("  Setting %s ... ", s))
  d   <- simulate_fps_data(2000, setting = s, outcome_type = "functional",
                            include_functional_cov = TRUE, seed = 2)
  w   <- fps_weighting(d$X, d$t_grid, c(0, 1), domain_name = "s",
                        covariates  = list(scalar = d$C,
                                           functional = list(d$D)),
                        cov_t_grids = list(d$t_grid),
                        cov_domains = list(c(0, 1)))
  eff <- fps_effect_estimation(d$Y, w,
                                outcome_t_grid      = d$t_grid,
                                outcome_domain      = c(0, 1),
                                outcome_domain_name = "t",
                                true_beta           = d$true_beta)
  m_w <- FPScausal:::.compute_error_metrics(eff$beta,            d$true_beta)
  m_u <- FPScausal:::.compute_error_metrics(eff$beta_unweighted, d$true_beta)
  cat(sprintf("ISE_w=%.4f  ISE_u=%.4f\n", m_w["ISE"], m_u["ISE"]))
  list(eff = eff, setting = s, d = d)
})

# Combined heatmap plot: 2x2 grid, weighted surface per setting
library(ggplot2)
library(patchwork)

heatmaps <- lapply(beta_results_fn, function(r) {
  tg  <- r$d$t_grid
  df  <- expand.grid(s = tg, t = tg)
  df$effect <- as.vector(r$eff$beta)
  ggplot(df, aes(x = s, y = t, fill = effect)) +
    geom_tile() +
    scale_fill_gradient2(low = "#2166AC", mid = "white", high = "#D62728",
                          midpoint = 0, name = NULL) +
    labs(x = "s", y = "t",
         title = paste0("Setting ", r$setting, " (weighted)")) +
    theme_bw(base_size = 9) +
    theme(plot.title      = element_text(hjust = 0.5, face = "bold"),
          legend.position = "none")
})
print(wrap_plots(heatmaps, ncol = 2))

cat("\n============================================================\n")
cat("Functional outcome test PASSED.\n")
cat("============================================================\n")