R/fps_weighting.R

In FPScausal: Functional Propensity Score for Causal Inference

#' Estimate functional propensity score weights
#'
#' Computes covariate-balancing weights for a functional treatment using the
#' empirical-likelihood balancing framework of Ciardulli, S. and Fontana, N.
#' (2026). Treatment is represented via FPCA (Karhunen--Loeve expansion);
#' The treatment is first represented via Functional Principal Component
#' Analysis (FPCA) through its Karhunen--Loeve expansion truncated at rank
#' \emph{L}; the resulting FPC scores and observed confounders are balanced
#' by solving the dual of the empirical-likelihood problem via the BFGS
#' quasi-Newton algorithm.  Functional covariates enter the balancing step
#' through their own FPC scores.
#'
#' @param treatment n x T numeric matrix of observed treatment trajectories, or
#'   an \code{fd} object from the \pkg{fda} package.
#' @param treat_grid Numeric vector of length T giving the observation grid
#'   of the treatment.  Required when \code{treatment} is a matrix; inferred
#'   automatically when \code{treatment} is an \code{fd} object.
#' @param treat_domain Numeric vector \code{c(a, b)} specifying the domain of
#'   the treatment.  If \code{NULL} (default) and \code{treatment} is a matrix,
#'   the domain is inferred as \code{c(min(treat_grid), max(treat_grid))}.
#'   Inferred automatically when \code{treatment} is an \code{fd} object.
#' @param domain_name Character string naming the domain variable (default
#'   \code{"s"}).  Used in axis labels and domain-overlap checks.
#' @param nbasis Integer.  Number of B-spline basis functions used for the
#'   treatment FPCA.  If \code{NULL} (default), chosen automatically as
#'   \code{max(10, round(0.6 * length(treat_grid)))}.
#' @param pve Numeric in (0, 1].  Proportion of variance explained threshold
#'   for the treatment FPCA (default 0.95).
#' @param covariates Either (a) an n x p numeric matrix of scalar covariates,
#'   or (b) a named list with elements \code{scalar} (n x p matrix, may be
#'   \code{NULL}) and \code{functional} (a list of matrices or \code{fd}
#'   objects representing functional covariates).
#' @param cov_grids A list of numeric vectors (one per functional covariate)
#'   giving the observation grids.  Required if \code{covariates$functional}
#'   contains matrices; inferred from the domain when \code{NULL}.
#' @param cov_domains A list of numeric vectors \code{c(a, b)} (one per
#'   functional covariate).  If \code{NULL}, inferred from
#'   \code{cov_grids} extremes.
#' @param cov_nbasis A list of integers (or \code{NULL}) for B-spline basis
#'   sizes of functional covariates.  Defaults to auto-selection.
#' @param cov_pve Numeric in (0, 1].  PVE threshold for functional covariate
#'   FPCA (default 0.95).
#' @param normalize Logical.  If \code{TRUE} (default), standardise FPC scores
#'   and confounders before the dual optimisation.
#' @param tol Relative convergence tolerance for the BFGS optimiser
#'   (default 1e-8).
#' @param maxit Maximum number of BFGS iterations (default 1000).
#'
#' @return An object of class \code{"fps_weighting"}, which is a named list
#'   with the following components:
#' \describe{
#'   \item{weights}{Numeric vector of length n.  Positive weights summing to
#'     1.}
#'   \item{fpca_treatment}{List returned by the internal FPCA routine,
#'     containing FPC scores (\code{scr}), eigenfunctions (\code{efn}), mean
#'     function (\code{mean}), eigenvalues (\code{eval}), variance proportions
#'     (\code{varprop}), cumulative PVE (\code{perc}), raw \code{pca.fd}
#'     object (\code{pca_fd}), number of components retained (\code{L}), and
#'     the \code{t_grid} and \code{domain} used.}
#'   \item{fpca_covariates}{List of FPCA results for functional covariates, or
#'     \code{NULL} if none were supplied.}
#'   \item{scalar_covariates}{The n x p scalar covariate matrix used.}
#'   \item{conf_matrix}{Full augmented confounder matrix fed to the optimiser
#'     (scalar covariates column-bound with FPC scores of functional
#'     covariates).}
#'   \item{convergence}{Convergence code from \code{\link[stats]{optim}}
#'     (0 = success).}
#'   \item{domain_name}{The domain name passed via \code{domain_name}.}
#'   \item{call}{The matched call.}
#' }
#'
#' @seealso \code{\link{fps_effect_estimation}}, \code{\link{simulate_fps_data}}
#'
#' @examples
#' \donttest{
#' dat <- simulate_fps_data(n = 2000, setting = "LL", seed = 1)
#'
#' # Scalar covariates only (treat_domain inferred from treat_grid)
#' w <- fps_weighting(
#'   treatment   = dat$X,
#'   treat_grid  = dat$t_grid,
#'   covariates  = dat$C
#' )
#' print(w)
#' plot(w, type = "balance")
#'
#' # Include one functional covariate
#' w2 <- fps_weighting(
#'   treatment   = dat$X,
#'   treat_grid  = dat$t_grid,
#'   treat_domain = c(0, 1),
#'   covariates  = list(scalar = dat$C, functional = list(dat$D)),
#'   cov_grids   = list(dat$t_grid)
#' )
#' plot(w2, type = "balance")
#' }
#'
#' @export
fps_weighting <- function(treatment,
                           treat_grid   = NULL,
                           treat_domain = NULL,
                           domain_name  = "s",
                           nbasis       = NULL,
                           pve          = 0.95,
                           covariates,
                           cov_grids    = NULL,
                           cov_domains  = NULL,
                           cov_nbasis   = NULL,
                           cov_pve      = 0.95,
                           normalize    = TRUE,
                           tol          = 1e-8,
                           maxit        = 1000) {

  cl <- match.call()

  # ---- Handle fd input for treatment ----
  if (inherits(treatment, "fd")) {
    rng <- treatment$basis$rangeval
    if (is.null(treat_domain)) treat_domain <- rng
    if (is.null(treat_grid))
      treat_grid <- seq(treat_domain[1], treat_domain[2], length.out = 51)
    treatment <- t(fda::eval.fd(treat_grid, treatment))
  } else {
    if (is.null(treat_grid)) {
      stop(
        "'treat_grid' must be provided when 'treatment' is not an 'fd' object."
      )
    }
    if (is.null(treat_domain)) {
      treat_domain <- c(min(treat_grid), max(treat_grid))
    }
  }

  if (!is.matrix(treatment)) treatment <- as.matrix(treatment)

  # ---- FPCA on treatment ----
  fpca_treat <- .fps_fpca(treatment, pve = pve, t_grid = treat_grid,
                           domain = treat_domain, nbasis = nbasis)

  # ---- Process covariates ----
  scalar_mat     <- NULL
  fpca_cov_list  <- NULL
  fpc_score_list <- list()

  if (is.matrix(covariates) || is.data.frame(covariates)) {
    scalar_mat <- as.matrix(covariates)
  } else if (is.list(covariates)) {
    if (!is.null(covariates$scalar)) {
      scalar_mat <- as.matrix(covariates$scalar)
    }
    if (!is.null(covariates$functional)) {
      func_covs     <- covariates$functional
      n_func        <- length(func_covs)
      fpca_cov_list <- vector("list", n_func)

      for (k in seq_len(n_func)) {
        fc <- func_covs[[k]]

        if (inherits(fc, "fd")) {
          cov_dom_k  <- fc$basis$rangeval
          cov_grid_k <- seq(cov_dom_k[1], cov_dom_k[2], length.out = 51)
          # Pass fd directly; .fps_fpca() handles it without Data2fd
        } else {
          cov_grid_k <- if (!is.null(cov_grids)) cov_grids[[k]] else treat_grid
          cov_dom_k  <- if (!is.null(cov_domains)) {
            cov_domains[[k]]
          } else {
            c(min(cov_grid_k), max(cov_grid_k))
          }
          fc <- as.matrix(fc)
        }

        nb_k <- if (!is.null(cov_nbasis)) cov_nbasis[[k]] else NULL
        fpca_cov_list[[k]] <- .fps_fpca(fc, pve = cov_pve,
                                         t_grid = cov_grid_k,
                                         domain = cov_dom_k,
                                         nbasis = nb_k)

        scores_k <- fpca_cov_list[[k]]$scr

        # Assign column names: Func_<CovName>_FPCk or Func_CovK_FPCk
        cov_name <- names(func_covs)[k]
        if (!is.null(cov_name) && nzchar(cov_name)) {
          prefix <- paste0("Func_", cov_name, "_FPC")
        } else {
          prefix <- paste0("Func_Cov", k, "_FPC")
        }
        colnames(scores_k) <- paste0(prefix, seq_len(ncol(scores_k)))
        fpc_score_list[[k]] <- scores_k
      }
    }
  } else {
    stop(paste0(
      "'covariates' must be a numeric matrix or a list with 'scalar' ",
      "and/or 'functional' elements."
    ))
  }

  # Build augmented confounder matrix
  conf_parts <- Filter(Negate(is.null), c(list(scalar_mat), fpc_score_list))
  if (length(conf_parts) == 0) {
    stop("At least one covariate (scalar or functional) must be provided.")
  }
  conf_matrix <- as.matrix(do.call(cbind, conf_parts))

  # ---- Dual optimisation ----
  opt_res <- .fps_compute_weights(fpca_treat$scr, conf_matrix,
                                   normalize = normalize,
                                   tol = tol, maxit = maxit)

  structure(
    list(
      weights           = opt_res$weights,
      fpca_treatment    = fpca_treat,
      fpca_covariates   = fpca_cov_list,
      scalar_covariates = scalar_mat,
      conf_matrix       = conf_matrix,
      convergence       = opt_res$convergence,
      domain_name       = domain_name,
      call              = cl
    ),
    class = "fps_weighting"
  )
}