R/fps_effect_estimation.R

In FPScausal: Functional Propensity Score for Causal Inference

#' Estimate causal effect of a functional treatment
#'
#' Given the FPS weights produced by \code{\link{fps_weighting}}, estimates the
#' causal effect function \eqn{\hat\mu(t)} (scalar/binary outcome) or the
#' causal effect surface \eqn{\hat\mu(s,t)} (functional outcome) via weighted
#' least squares.  Optional bootstrap inference is available.
#'
#' \strong{Scalar and binary outcomes.}
#' The treatment FPC scores \eqn{A_i} are regressed on the outcome using
#' \code{lm} (scalar or binary, linear probability model) with the FPS
#' weights.  The estimated effect function is then reconstructed as
#' \eqn{\hat\mu(t) = \sum_k \hat\mu_k \phi_k(t)}.
#'
#' \strong{Functional outcome.}
#' For each outcome FPC component \eqn{j}, the regression
#' \eqn{c_{ij} \sim A_i} is solved with the FPS weights.  The causal surface
#' is reconstructed as
#' \eqn{\hat\mu(s,t) = \Phi_X \hat{B} \Phi_Y^\top}
#' where \eqn{\hat{B}} collects the regression coefficients.
#'
#' \strong{Bootstrap CIs.} Scalar/binary: residual bootstrap, B resamples.
#' Functional: pairs bootstrap, B resamples.  Pointwise reflected-percentile
#' confidence intervals are returned.
#'
#' @param outcome Numeric vector (scalar/binary, length n) or n x T matrix
#'   (functional outcome).
#' @param fps_object Object of class \code{"fps_weighting"} returned by
#'   \code{\link{fps_weighting}}.
#' @param outcome_t_grid Numeric vector.  Observation grid for functional
#'   outcome.  Required when \code{outcome} is a matrix.
#' @param outcome_domain Numeric \code{c(a, b)}.  Domain for functional outcome.
#'   Required when \code{outcome} is a matrix.
#' @param outcome_domain_name Character.  Name of the outcome domain (default
#'   \code{"s"}).
#' @param outcome_nbasis Integer or \code{NULL}.  B-spline basis size for the
#'   outcome FPCA.  Auto-selected if \code{NULL}.
#' @param outcome_pve Numeric in (0, 1].  PVE threshold for outcome FPCA
#'   (default 0.95).
#' @param treatment_pve Numeric or \code{NULL}.  If not \code{NULL}, re-runs
#'   FPCA on the treatment with this PVE threshold for the outcome estimation
#'   step (allowing L* != L).  Default \code{NULL} (reuses
#'   \code{fps_object$fpca_treatment}).
#' @param bootstrap Logical.  If \code{TRUE}, compute bootstrap confidence
#'   intervals (default \code{FALSE}).
#' @param B Integer.  Number of bootstrap resamples (default 1000).
#' @param alpha Numeric.  Significance level for bootstrap CIs (default 0.05).
#' @param true_beta Optional.  Numeric vector (scalar/binary) or matrix
#'   (functional) containing the true causal effect, used for visual comparison
#'   in plots and error metrics in \code{summary}.
#' @param seed Integer or \code{NULL}.  Random seed for bootstrap reproducibility.
#'
#' @return An object of class \code{"fps_effect_estimation"}, a named list with:
#' \describe{
#'   \item{outcome_type}{Character: \code{"scalar"}, \code{"binary"}, or
#'     \code{"functional"}.}
#'   \item{beta}{Estimated causal effect, evaluated on \code{t_grid}
#'     (numeric vector for scalar/binary) or on the
#'     \code{t_grid} x \code{outcome_t_grid} grid (matrix for functional).}
#'   \item{beta_unweighted}{Same as \code{beta} but from unweighted regression.}
#'   \item{fpca_treatment}{FPCA of the treatment used in estimation.}
#'   \item{fpca_outcome}{\code{NULL} for scalar/binary; FPCA list for
#'     functional outcome.}
#'   \item{ci_lower, ci_upper}{\code{NULL} if \code{bootstrap = FALSE};
#'     otherwise lower and upper bootstrap CI bounds (same shape as
#'     \code{beta}).}
#'   \item{alpha}{Significance level used.}
#'   \item{t_grid}{Treatment domain grid.}
#'   \item{outcome_t_grid}{\code{NULL} for scalar/binary; outcome grid for
#'     functional.}
#'   \item{domain_name}{Treatment domain name.}
#'   \item{outcome_domain_name}{Outcome domain name.}
#'   \item{true_beta}{Passed through unchanged.}
#'   \item{fps_object}{The input \code{fps_weighting} object.}
#'   \item{call}{The matched call.}
#' }
#'
#' @seealso \code{\link{fps_weighting}}, \code{\link{simulate_fps_data}}
#'
#' @examples
#' \donttest{
#' dat <- simulate_fps_data(n = 2000, setting = "LL", seed = 1)
#'
#' w <- fps_weighting(dat$X, dat$t_grid, c(0, 1), covariates = dat$C)
#'
#' # Scalar outcome, no bootstrap
#' eff <- fps_effect_estimation(dat$Y, w, true_beta = dat$true_beta)
#' plot(eff, type = "effect")
#' plot(eff, type = "comparison")
#'
#' # With bootstrap (small B for illustration)
#' eff_boot <- fps_effect_estimation(dat$Y, w, bootstrap = TRUE, B = 100,
#'                                   true_beta = dat$true_beta, seed = 42)
#' plot(eff_boot, type = "significance")
#' }
#'
#' @export
fps_effect_estimation <- function(outcome,
                                   fps_object,
                                   outcome_t_grid       = NULL,
                                   outcome_domain       = NULL,
                                   outcome_domain_name  = "t",
                                   outcome_nbasis       = NULL,
                                   outcome_pve          = 0.95,
                                   treatment_pve        = NULL,
                                   bootstrap            = FALSE,
                                   B                    = 1000,
                                   alpha                = 0.05,
                                   true_beta            = NULL,
                                   seed                 = NULL) {

  cl <- match.call()
  .check_fps_weighting(fps_object)

  outcome_type <- .detect_outcome_type(outcome)

  # ---- Extract grid/domain from fd outcome if needed ----
  if (inherits(outcome, "fd")) {
    rng_o <- outcome$basis$rangeval
    if (is.null(outcome_t_grid))
      outcome_t_grid <- seq(rng_o[1], rng_o[2], length.out = 100)
    if (is.null(outcome_domain))
      outcome_domain <- rng_o
  }

  # ---- Validate functional outcome args ----
  if (outcome_type == "functional") {
    if (is.null(outcome_t_grid)) {
      stop("'outcome_t_grid' must be provided for a functional outcome.")
    }
    if (is.null(outcome_domain)) {
      outcome_domain <- c(min(outcome_t_grid), max(outcome_t_grid))
    }
    .check_domains_overlap(
      fps_object$fpca_treatment$domain, fps_object$domain_name,
      outcome_domain, outcome_domain_name
    )
  }

  # ---- Treatment FPCA for estimation step ----
  if (!is.null(treatment_pve)) {
    # Re-run FPCA with different PVE for outcome estimation step (L* != L)
    fpca_treat  <- .fps_fpca(
      X      = fps_object$fpca_treatment$pca_fd$scores %*%
                 t(fps_object$fpca_treatment$efn) +
                 matrix(fps_object$fpca_treatment$mean,
                        nrow = nrow(fps_object$fpca_treatment$scr),
                        ncol = length(fps_object$fpca_treatment$t_grid),
                        byrow = TRUE),
      pve    = treatment_pve,
      t_grid = fps_object$fpca_treatment$t_grid,
      domain = fps_object$fpca_treatment$domain
    )
  } else {
    fpca_treat <- fps_object$fpca_treatment
  }

  w   <- fps_object$weights
  A   <- fpca_treat$scr     # n x L*
  efX <- fpca_treat$efn     # T x L*

  # ---- Outcome FPCA (functional case) ----
  fpca_out <- NULL
  if (outcome_type == "functional") {
    fpca_out <- .fps_fpca(outcome,
                           pve    = outcome_pve,
                           t_grid = outcome_t_grid,
                           domain = outcome_domain,
                           nbasis = outcome_nbasis)
  }

  # ---- Weighted regression ----
  beta_w <- .estimate_beta(outcome_type, outcome, A, efX, w, fpca_out)
  beta_u <- .estimate_beta(outcome_type, outcome, A, efX,
                            rep(1 / nrow(A), nrow(A)), fpca_out)

  # ---- Analytical SE (scalar/binary only; used for CIs without bootstrap) ----
  analytical_se <- NULL
  if (outcome_type != "functional") {
    fit_w_se  <- .fit_weighted_model(outcome_type, outcome, A, w)
    vc        <- stats::vcov(fit_w_se)[-1, -1, drop = FALSE]
    var_beta  <- diag(efX %*% vc %*% t(efX))
    analytical_se <- sqrt(pmax(var_beta, 0))
  }

  # ---- Bootstrap ----
  ci_lower <- NULL
  ci_upper <- NULL

  if (bootstrap) {
    if (!is.null(seed)) set.seed(seed)
    n <- nrow(A)
    pb <- progress::progress_bar$new(
      format  = "  Bootstrap [:bar] :percent  ETA :eta",
      total   = B,
      clear   = FALSE
    )

    if (outcome_type == "functional") {
      nS     <- nrow(efX)
      nT_out <- nrow(fpca_out$efn)
      boot_arr <- array(NA_real_, dim = c(B, nS, nT_out))
      C_scores <- fpca_out$scr  # n x Ly

      L_star <- ncol(A)
      Ly     <- ncol(C_scores)
      if (is.null(Ly) || Ly == 0) Ly <- 1L  # guard for 1-column case

      for (b in seq_len(B)) {
        idx  <- sample(n, replace = TRUE)
        A_b  <- A[idx, , drop = FALSE]
        C_b  <- as.matrix(C_scores[idx, , drop = FALSE])
        w_b  <- as.vector(w[idx])

        # Direct WLS avoids lm formula quirks (rank-deficient bootstrap samples
        # cause lm to drop predictors and shorten coef() silently).
        Xb    <- cbind(1, A_b)              # n x (L*+1)
        sw    <- sqrt(w_b)
        XtWX  <- crossprod(Xb * sw)         # (L*+1) x (L*+1)
        XtWY  <- t(Xb) %*% (C_b * w_b)     # (L*+1) x Ly
        B_full <- tryCatch(
          solve(XtWX, XtWY),
          error = function(e) MASS::ginv(XtWX) %*% XtWY
        )
        B_hat_b <- matrix(B_full[-1, ], nrow = L_star, ncol = Ly)
        boot_arr[b, , ] <- efX %*% B_hat_b %*% t(fpca_out$efn)
        pb$tick()
      }

      ci_lower <- matrix(NA_real_, nS, nT_out)
      ci_upper <- matrix(NA_real_, nS, nT_out)
      for (i in seq_len(nS)) {
        for (j in seq_len(nT_out)) {
          ci_band        <- .reflected_ci(beta_w[i, j], boot_arr[, i, j], alpha)
          ci_lower[i, j] <- ci_band["lwr"]
          ci_upper[i, j] <- ci_band["upr"]
        }
      }
    } else {
      # Scalar / binary: residual bootstrap
      fit_obs       <- .fit_weighted_model(outcome_type, outcome, A, w)
      fitted_obs    <- stats::fitted(fit_obs)
      res_obs       <- outcome - fitted_obs
      coef_obs      <- stats::coef(fit_obs)[-1]
      beta_obs_vec  <- as.vector(efX %*% coef_obs)

      boot_mat <- matrix(NA_real_, B, length(beta_obs_vec))
      for (b in seq_len(B)) {
        res_b     <- sample(res_obs, replace = TRUE)
        y_b       <- fitted_obs + res_b
        fit_b     <- .fit_weighted_model(outcome_type, y_b, A, w)
        coef_b    <- stats::coef(fit_b)[-1]
        boot_mat[b, ] <- as.vector(efX %*% coef_b)
        pb$tick()
      }
      ci_lower <- vapply(seq_along(beta_obs_vec), function(t_idx) {
        .reflected_ci(beta_obs_vec[t_idx], boot_mat[, t_idx], alpha)["lwr"]
      }, numeric(1))
      ci_upper <- vapply(seq_along(beta_obs_vec), function(t_idx) {
        .reflected_ci(beta_obs_vec[t_idx], boot_mat[, t_idx], alpha)["upr"]
      }, numeric(1))
    }
  }

  structure(
    list(
      outcome_type        = outcome_type,
      beta                = beta_w,
      beta_unweighted     = beta_u,
      fpca_treatment      = fpca_treat,
      fpca_outcome        = fpca_out,
      ci_lower            = ci_lower,
      ci_upper            = ci_upper,
      analytical_se       = analytical_se,
      alpha               = alpha,
      t_grid              = fpca_treat$t_grid,
      outcome_t_grid      = outcome_t_grid,
      domain_name         = fps_object$domain_name,
      outcome_domain_name = outcome_domain_name,
      true_beta           = true_beta,
      fps_object          = fps_object,
      call                = cl
    ),
    class = "fps_effect_estimation"
  )
}

# ---- Internal helpers for outcome estimation ----

#' Fit weighted regression model for a given outcome type
#' @keywords internal
.fit_weighted_model <- function(outcome_type, outcome, A, w) {
  df <- as.data.frame(A)
  names(df) <- paste0("A", seq_len(ncol(A)))
  df$outcome <- outcome
  fmla <- stats::as.formula(paste("outcome ~", paste(names(df)[seq_len(ncol(A))], collapse = " + ")))
  # Both scalar and binary use WLS (linear probability model for binary).
  # glm with FPS weights is numerically unstable.
  stats::lm(fmla, data = df, weights = w * length(w))
}

#' Estimate beta (effect function or surface)
#' @keywords internal
.estimate_beta <- function(outcome_type, outcome, A, efX, w, fpca_out = NULL) {
  if (outcome_type == "functional") {
    C_scores <- fpca_out$scr   # n x Ly
    fit      <- stats::lm(C_scores ~ A, weights = w * nrow(A))
    cf       <- stats::coef(fit)
    if (is.vector(cf)) {
      B_hat <- matrix(cf[-1], ncol = 1)
    } else {
      B_hat <- cf[-1, , drop = FALSE]   # Lx x Ly
    }
    efX %*% B_hat %*% t(fpca_out$efn)  # Tx x Ty
  } else {
    fit      <- .fit_weighted_model(outcome_type, outcome, A, w)
    coefs    <- stats::coef(fit)[-1]
    as.vector(efX %*% coefs)
  }
}