R/simulate_fps.R

In FPScausal: Functional Propensity Score for Causal Inference

#' Simulate functional propensity score data
#'
#' Generates a synthetic dataset for testing and illustrating the FPScausal
#' workflow. 
#'
#' **Treatment** X_i(t) is built from 6 Fourier eigenfunctions with eigenvalues
#' (16, 12, 8, 4, 1, 0.5). **Scalar confounders** C_i are 3-dimensional vectors
#' whose relationship to X's FPC scores is either linear or quadratic. An
#' optional **functional covariate** D_i(t) is generated from 4 Fourier
#' components. The **scalar outcome** is
#' Y_i = 1 + integral(beta(t) * X_i(t)) + g(C_i) + N(0,25),
#' and the **functional outcome** is
#' Y_i(t) = mu0(t) + integral(mu(s,t) * X_i(s) ds) + h(D_i) + GP_error.
#'
#' The four settings ("LL", "LN", "NL", "NN") vary whether the
#' treatment-confounder ("L"inear / "N"onlinear) and confounder-outcome
#' ("L"inear / "N"onlinear) relationships are linear or quadratic.
#'
#' @param n Integer. Number of subjects (default 200).
#' @param setting Character. One of `"LL"`, `"LN"`, `"NL"`, `"NN"`, where the
#'   first letter controls the treatment-confounder relationship and the second
#'   controls the confounder-outcome relationship. Default `"LL"`.
#' @param outcome_type Character. Either `"scalar"` or `"functional"`.
#'   Default `"scalar"`.
#' @param p_scalar Integer. Number of scalar confounders. Default 3.
#' @param include_functional_cov Logical. If `TRUE` (default), include one
#'   functional covariate D(t) in the returned list.
#' @param domain Numeric vector \code{c(a, b)} giving the time domain of the
#'   functional objects.  Default \code{c(0, 1)}, matching the paper's
#'   simulation study. Change this to use a different time range (e.g.
#'   \code{c(50, 70)} for age in years).  The 51 evaluation points are always
#'   equally spaced within \code{domain}.
#' @param seed Integer or NULL. Random seed for reproducibility.
#'
#' @return A named list with:
#' \describe{
#'   \item{X}{n x 51 matrix. Observed treatment trajectories on [0,1].}
#'   \item{Y}{If `outcome_type = "scalar"`: numeric vector of length n.
#'            If `outcome_type = "functional"`: n x 51 matrix.}
#'   \item{C}{n x p_scalar matrix. Scalar confounders.}
#'   \item{D}{n x 51 matrix. Functional covariate (if `include_functional_cov = TRUE`,
#'            else `NULL`).}
#'   \item{t_grid}{Numeric vector of 51 equally-spaced points on [0,1].}
#'   \item{true_beta}{True causal effect. For scalar outcome: numeric vector
#'     of length 51. For functional outcome: 51 x 51 matrix mu(s,t).}
#'   \item{setting}{The `setting` argument used.}
#'   \item{outcome_type}{The `outcome_type` argument used.}
#' }
#'
#' @examples
#' dat <- simulate_fps_data(n = 100, setting = "LL", outcome_type = "scalar",
#'                          seed = 42)
#' str(dat)
#'
#' @export
simulate_fps_data <- function(n = 200,
                               setting = c("LL", "LN", "NL", "NN"),
                               outcome_type = c("scalar", "functional"),
                               p_scalar = 3,
                               include_functional_cov = TRUE,
                               domain = c(0, 1),
                               seed = NULL) {

  setting      <- match.arg(setting)
  outcome_type <- match.arg(outcome_type)

  if (!is.null(seed)) set.seed(seed)

  linear_C <- substr(setting, 1, 1) == "L"
  linear_Y <- substr(setting, 2, 2) == "L"

  tgrid <- seq(domain[1], domain[2], length.out = 51L)

  # ---- Latent scores ----
  p_z <- max(6L, p_scalar)
  Z   <- MASS::mvrnorm(n, mu = rep(0, p_z), Sigma = diag(p_z))

  # FPC scores for treatment (6 components)
  A        <- matrix(NA_real_, n, 6)
  A[, 1]   <- sqrt(16) * Z[, 1]
  A[, 2]   <- sqrt(12) * Z[, 2]
  A[, 3]   <- sqrt( 8) * Z[, 3]
  A[, 4]   <- sqrt( 4) * Z[, 4]
  A[, 5]   <- sqrt( 1) * Z[, 5]
  A[, 6]   <- sqrt(0.5) * Z[, 6]

  # Treatment matrix X
  X <- outer(A[, 1], sqrt(2) * sin(2 * pi * tgrid)) +
       outer(A[, 2], sqrt(2) * cos(2 * pi * tgrid)) +
       outer(A[, 3], sqrt(2) * sin(4 * pi * tgrid)) +
       outer(A[, 4], sqrt(2) * cos(4 * pi * tgrid)) +
       outer(A[, 5], sqrt(2) * sin(6 * pi * tgrid)) +
       outer(A[, 6], sqrt(2) * cos(6 * pi * tgrid))

  # ---- Scalar confounders ----
  C <- matrix(NA_real_, n, p_scalar)
  if (linear_C) {
    C[, 1] <- Z[, 1] + stats::rnorm(n, 0, 1)
    if (p_scalar > 1) {
      for (j in 2:p_scalar) C[, j] <- 0.2 * Z[, j] + stats::rnorm(n, 0, 0.5)
    }
  } else {
    C[, 1] <- (Z[, 1] + 0.5)^2 + stats::rnorm(n, 0, 1)
    if (p_scalar > 1) {
      for (j in 2:p_scalar) C[, j] <- 0.2 * Z[, j] + stats::rnorm(n, 0, 0.5)
    }
  }

  # ---- Functional covariate D (optional) ----
  D <- NULL
  LAMBDA <- NULL
  if (include_functional_cov || outcome_type == "functional") {
    LAMBDA <- matrix(NA_real_, n, 4)
    if (linear_C) {
      LAMBDA[, 1] <- Z[, 1] + 4 * stats::rnorm(n)
      LAMBDA[, 2] <- Z[, 2] + 2 * stats::rnorm(n)
      LAMBDA[, 3] <- Z[, 3] +     stats::rnorm(n)
      LAMBDA[, 4] <- Z[, 4] +     stats::rnorm(n) / sqrt(2)
    } else {
      LAMBDA[, 1] <- (Z[, 1] + 0.5)^2 + 4 * stats::rnorm(n)
      LAMBDA[, 2] <- Z[, 2] + 2 * stats::rnorm(n)
      LAMBDA[, 3] <- Z[, 3] +     stats::rnorm(n)
      LAMBDA[, 4] <- Z[, 4] +     stats::rnorm(n) / sqrt(2)
    }
    D <- outer(LAMBDA[, 1], sqrt(2) * sin(2 * pi * tgrid)) +
         outer(LAMBDA[, 2], sqrt(2) * cos(2 * pi * tgrid)) +
         outer(LAMBDA[, 3], sqrt(2) * sin(4 * pi * tgrid)) +
         outer(LAMBDA[, 4], sqrt(2) * cos(4 * pi * tgrid))
    if (!include_functional_cov) D <- NULL
  }

  # ---- True causal effect ----
  if (outcome_type == "scalar") {
    true_beta <- 2 * sqrt(2) * sin(2 * pi * tgrid) +
                     sqrt(2) * cos(2 * pi * tgrid) +
                 0.5 * sqrt(2) * sin(4 * pi * tgrid) +
                 0.5 * sqrt(2) * cos(4 * pi * tgrid)
  } else {
    sg <- tgrid
    tg <- tgrid
    true_beta <- outer(sg, tg, function(s, t) {
      2 * sqrt(2) * sin(2 * pi * s) * cos(2 * pi * t) +
      2 * sqrt(2) * sin(2 * pi * t) * cos(2 * pi * s) +
          sqrt(2) * cos(4 * pi * t) * sin(4 * pi * s) +
          sqrt(2) * cos(4 * pi * s) * sin(4 * pi * t)
    })
  }

  # ---- Outcome ----
  if (outcome_type == "scalar") {
    # integral(beta(t) * X_i(t) dt) approximated as mean(beta * X_i)
    dt         <- tgrid[2] - tgrid[1]
    int_betaX  <- as.vector(X %*% true_beta) * dt
    if (linear_Y) {
      g_C <- 2 * C[, 1]
    } else {
      g_C <- 2 * C[, 1] + C[, 2]^2
    }
    Y <- 1 + int_betaX + g_C + stats::rnorm(n, 0, 5)

  } else {
    # Functional outcome
    # mu0(t) = 1 + cos(2*pi*t)
    mu0 <- 1 + cos(2 * pi * tgrid)

    # integral(mu(s,t) * X_i(s) ds) -> n x 51
    dt       <- tgrid[2] - tgrid[1]
    int_muX  <- X %*% true_beta * dt   # n x 51

    if (linear_Y) {
      h_D <- outer(LAMBDA[, 1], rep(2, length(tgrid)))
    } else {
      h_D <- outer(LAMBDA[, 1], rep(2, length(tgrid))) +
             matrix(LAMBDA[, 2]^2, n, length(tgrid))
    }

    # GP errors: squared-exponential kernel; length scale = 10% of domain width
    l_sq    <- (0.1 * (domain[2] - domain[1]))^2
    cov_mat <- exp(-0.5 * outer(tgrid, tgrid, "-")^2 / l_sq)
    errors  <- MASS::mvrnorm(n, mu = rep(0, length(tgrid)), Sigma = cov_mat)

    Y <- matrix(mu0, n, length(tgrid), byrow = TRUE) +
         int_muX + h_D + errors
  }

  list(
    X            = X,
    Y            = Y,
    C            = C,
    D            = D,
    t_grid       = tgrid,
    true_beta    = true_beta,
    setting      = setting,
    outcome_type = outcome_type
  )
}