R/compute_weights.R

In FPScausal: Functional Propensity Score for Causal Inference

#' Internal dual formulation weight estimator
#'
#' Estimates functional propensity score (FPS) weights by maximising the
#' empirical likelihood subject to covariate-balancing constraints and solving
#' the resulting dual problem.  The dual reduces to minimising
#' log(sum_i exp(-theta' g_i)) over the unconstrained parameter vector theta,
#' where g_i = [A_i, C_i, vec(A_i C_i')] stacks the balancing moments.
#' The optimisation is solved via the BFGS quasi-Newton algorithm and uses the
#' log-sum-exp trick for numerical stability. 
#'
#' @param treat_scores n x L numeric matrix of FPC scores for the treatment.
#' @param conf_matrix n x p numeric matrix of confounders (scalar + FPC scores
#'   of functional covariates).
#' @param normalize Logical. If TRUE (default), standardise A and C before
#'   building g_i.
#' @param tol Relative convergence tolerance for \code{optim} (default 1e-8).
#' @param maxit Maximum number of BFGS iterations (default 1000).
#' @importFrom stats optim var
#' @return A named list:
#'   \item{weights}{Numeric vector of length n. Positive, sums to 1.}
#'   \item{theta}{Optimal dual parameter vector.}
#'   \item{convergence}{Convergence code from \code{optim} (0 = success).}
#'   \item{value}{Optimal dual objective value.}
#' @keywords internal
.fps_compute_weights <- function(treat_scores, conf_matrix,
                                  normalize = TRUE,
                                  tol = 1e-8,
                                  maxit = 1000) {

  n <- nrow(treat_scores)

  if (normalize) {
    treat_scores <- scale(treat_scores)
    conf_chol    <- chol(var(conf_matrix))
    conf_std     <- conf_matrix %*% solve(conf_chol)
    conf_matrix  <- scale(conf_std, center = TRUE, scale = TRUE)
  }

  g_mat <- do.call(rbind, lapply(seq_len(n), function(i) {
    A_i <- treat_scores[i, ]
    C_i <- conf_matrix[i, ]
    c(A_i, C_i, as.vector(outer(A_i, C_i)))
  }))

  objective <- function(theta) {
    linpred <- as.vector(-g_mat %*% theta)
    m       <- max(linpred)
    log(sum(exp(linpred - m))) + m
  }

  gradient <- function(theta) {
    linpred <- as.vector(-g_mat %*% theta)
    m       <- max(linpred)
    w       <- exp(linpred - m)
    w       <- w / sum(w)
    -colSums(sweep(g_mat, 1, w, "*"))
  }

  theta_init <- rep(0, ncol(g_mat))

  opt <- stats::optim(
    par     = theta_init,
    fn      = objective,
    gr      = gradient,
    method  = "BFGS",
    control = list(reltol = tol, maxit = maxit)
  )

  if (opt$convergence != 0) {
    warning(
      "Dual optimisation may not have converged (code ", opt$convergence, "). ",
      "Consider increasing maxit or checking for multicollinearity in confounders."
    )
  }

  linpred <- as.vector(-g_mat %*% opt$par)
  m       <- max(linpred)
  w       <- exp(linpred - m)
  w       <- w / sum(w)

  list(
    weights     = w,
    theta       = opt$par,
    convergence = opt$convergence,
    value       = opt$value
  )
}