R/plsi.lr.auto.constraint.R

In EPLSIM: Partial Linear Single Index Models for Environmental Mixture Analysis

#' Partial linear single index regression with automatic smoothness selection
#'
#' Same model as \code{plsi.lr.v2()}, but the link function is estimated with a
#' penalized regression spline (\code{mgcv::gam}, REML smoothing-parameter
#' selection) instead of a fixed-df natural spline. This removes the need to
#' hand-tune \code{spline.num}: you supply an upper bound on basis complexity
#' (\code{k}), and REML shrinks unneeded wiggliness toward a straight line
#' automatically. Effective degrees of freedom actually used by the fitted
#' link function are reported in \code{si.fun.edf}.
#'
#' The single-index coefficient \eqn{\beta} is constrained to \eqn{\|\beta\| = 1}
#' and \eqn{\beta_1 > 0} for identifiability. Rather than optimizing over
#' \eqn{\beta} directly and projecting/normalizing after the fact, the search
#' is reparameterized in terms of \eqn{\sigma = (\sigma_1, \dots, \sigma_{p-1})}
#' via \eqn{\beta(\sigma) = \{(1 - \|\sigma\|^2)^{1/2}, \sigma_1, \dots,
#' \sigma_{p-1}\}^T}. Both constraints hold automatically for any
#' \eqn{\sigma} with \eqn{\|\sigma\| < 1}, so \code{optim()} runs unconstrained
#' over \eqn{\sigma} and \eqn{\beta_1 > 0}, \eqn{\|\beta\| = 1} never need to be
#' imposed after the fact. Standard errors for \eqn{\beta} are obtained from
#' the Hessian in \eqn{\sigma}-space via the delta method.
#'
#' The anchor exposure -- the one assigned \eqn{\beta_1}, the coordinate
#' constrained \eqn{> 0} -- is chosen as the candidate exposure with the
#' largest \emph{absolute} correlation with \code{Y.name} (not the largest
#' signed/positive correlation), so an exposure strongly but negatively
#' associated with the outcome can still be selected as the anchor. Which
#' exposure was selected, and its correlation with \code{Y.name}, is reported
#' in \code{attr(si.coefficient, "note")} and in \code{anchor.note} of the
#' returned list.
#'
#' @param data A data set including all needed variables
#' @param Y.name Variable name for scalar outcome
#' @param X.name Variable name vector for exposures
#' @param Z.name Variable name vector for confounders
#' @param k Upper bound on basis dimension for the smooth link function
#'   (analogous to \code{spline.num} in \code{plsi.lr.v2}, but not a fixed df,
#'   REML penalizes down from this ceiling). Default 10.
#' @param bs Smooth basis type passed to \code{mgcv::s()}; default \code{"cr"}
#'   (cubic regression spline).
#' @param initial.random.num Number of random initials for the single-index
#'   direction search
#' @param seed A single integer used to seed the RNG so results (random initials
#'   and downstream optim() convergence) are reproducible. Set to NULL to skip
#'   seeding and use whatever RNG state is already active in the calling
#'   environment. Default is 2026.
#' @return A list of model estimation and prediction results, structured
#'   analogously to \code{plsi.lr.v2()}'s output.
#' @importFrom stats qnorm
#' @importFrom mgcv gam
#'
#' @examples
#' \dontrun{
#' # example to run plsi.lr.auto.constraint
#' data(nhanes.new)
#' data <- nhanes.new
#'
#' Y.name <- "log.triglyceride"
#' X.name <- c("X1_trans.b.carotene", "X2_retinol", "X3_g.tocopherol", "X4_a.tocopherol",
#'             "X5_PCB99", "X6_PCB156", "X7_PCB206",
#'             "X8_3.3.4.4.5.pncb", "X9_1.2.3.4.7.8.hxcdf", "X10_2.3.4.6.7.8.hxcdf")
#' Z.name <- c("AGE.c", "SEX.Female", "RACE.NH.Black",
#'            "RACE.MexicanAmerican", "RACE.OtherRace", "RACE.Hispanic" )
#'
#' k <- 10
#' bs <- "cr"
#' initial.random.num <- 1
#' seed = 2026
#'
#' model_lr_auto <- plsi.lr.auto.constraint(data = data,
#'                       Y.name = Y.name, X.name = X.name, Z.name = Z.name,
#'                       k = k, bs = bs, initial.random.num = initial.random.num, seed = seed)
#' }
#' @keywords partial linear single index
#' @keywords linear regression
#' @keywords automatic smoothness
#' @author Yuyan Wang
#' @export
plsi.lr.auto.constraint <- function(data, Y.name, X.name, Z.name,
                                    k = 10, bs = "cr", initial.random.num = 5,
                                    seed = 2026) {

  stopifnot(initial.random.num >= 1)
  if (length(X.name) < 2) {
    stop("plsi.lr.auto.constraint() requires at least 2 exposures in X.name (the sigma ",
         "reparameterization needs p - 1 >= 1 free coordinates).")
  }
  if (!is.null(seed)) set.seed(seed)
  missing_cols <- setdiff(c(Y.name, X.name, Z.name), colnames(data))
  if (length(missing_cols) > 0) {
    stop("Column(s) not found in data: ", paste(missing_cols, collapse = ", "))
  }
  if (anyNA(data[, c(Y.name, X.name, Z.name)])) {
    stop("Missing values found in Y/X/Z columns; please handle NAs before calling plsi.lr.auto.constraint().")
  }
  if (!requireNamespace("mgcv", quietly = TRUE)) {
    stop("Package 'mgcv' is required for plsi.lr.auto.constraint(). Please install it.")
  }

  cor_linear <- stats::cor(data[, c(Y.name, X.name)])
  cor_linear <- as.data.frame(cor_linear[-1, 1])
  colnames(cor_linear)[1] <- "cor_with_Y"
  cor_linear$exposure <- rownames(cor_linear)
  # Sort by |correlation| with Y, not signed correlation. After this sort,
  # X.name[1] becomes the "anchor" exposure -- the one whose single-index
  # coefficient is beta_1, constrained > 0 by the beta(sigma) reparameterization
  # below. Sorting on the signed correlation (decreasing = TRUE) would only
  # ever pick the most *positively* correlated exposure as the anchor, which
  # biases the anchor choice against exposures that are strongly but
  # *negatively* associated with Y.name. Sorting on |correlation| instead picks
  # whichever exposure is most strongly associated with the outcome in either
  # direction, and its sign is absorbed into beta_1 > 0 by construction.
  cor_linear <- cor_linear[order(abs(cor_linear$cor_with_Y), decreasing = TRUE), ]
  X.name <- cor_linear$exposure

  y <- data[, Y.name]; x <- as.matrix(data[, X.name]); z_df <- data[, Z.name, drop = FALSE]
  n <- nrow(data); x_length <- length(X.name); z_length <- length(Z.name)

  m0 <- stats::glm(y ~ x + as.matrix(z_df), family = "gaussian")

  # beta(sigma) reparameterization for the ||beta|| = 1, beta_1 > 0 constraint
  # ---------------------------------------------------------------------------
  # beta_from_sigma(): maps a free sigma in R^(p-1) to the constrained beta in
  # R^p. beta_1 = sqrt(1 - ||sigma||^2) is >= 0 by construction (>0 whenever
  # ||sigma|| < 1, which fn()/optim() enforce), and ||beta||^2 = beta_1^2 +
  # ||sigma||^2 = 1 always. Returns NULL for infeasible sigma. The 1e-8 buffer
  # (instead of the exact boundary ||sigma|| = 1) keeps beta_1 strictly > 0,
  # avoiding a literal division-by-zero later in the delta-method Jacobian
  # (which divides by beta_1) if optim() ever lands exactly on the boundary.
  beta_from_sigma <- function(sigma) {
    ss <- sum(sigma^2)
    if (ss >= 1 - 1e-8) return(NULL)
    c(sqrt(1 - ss), sigma)
  }
  # sigma_from_beta(): inverse map, used only to build starting values from
  # candidate beta directions (linear-model coefficients, random draws). The
  # sign/normalization here is just to produce a valid (||sigma|| < 1)
  # starting point for optim(); it plays no role in the constraint itself,
  # which beta_from_sigma() enforces for every sigma optim() proposes.
  sigma_from_beta <- function(beta) {
    beta <- beta * sign(beta[1]) / sqrt(sum(beta^2))
    beta[-1]
  }

  initial_table <- as.data.frame(matrix(NA, nrow = (1 + initial.random.num), ncol = x_length))
  colnames(initial_table) <- X.name
  rownames(initial_table) <- c("linear", paste0("random_", 1:initial.random.num))
  initial_table[1, ] <- m0$coefficients[2:(1 + x_length)]
  initial_table[2:(1 + initial.random.num), ] <- stats::runif(initial.random.num * x_length, -1, 1)
  for (i in 1:nrow(initial_table)) {
    initial_table[i, ] <- initial_table[i, ] * sign(initial_table[i, 1]) / sqrt(sum(initial_table[i, ]^2))
  }
  # Starting values for the sigma-space search: drop beta_1 (it's determined
  # by sigma, not free) from each normalized candidate direction above.
  initial_sigma_table <- initial_table[, -1, drop = FALSE]
  colnames(initial_sigma_table) <- paste0("sigma_", seq_len(x_length - 1))

  # REML-penalized smooth link, refit for each candidate direction beta.
  # No spline.num / degrees of freedom to hand-tune: 'k' is just a ceiling,
  # and REML shrinks the smooth toward linear on its own.
  fit_gam <- function(beta_0) {
    u0 <- as.vector(x %*% as.vector(beta_0))
    dat_gam <- cbind(data.frame(y = y, u0 = u0), z_df)
    form <- stats::as.formula(paste0("y ~ s(u0, bs='", bs, "', k=", k, ") + ",
                                     paste(Z.name, collapse = " + ")))
    mgcv::gam(form, data = dat_gam, method = "REML")
  }

  # fn() is now a function of sigma (length p - 1), not beta (length p).
  # beta_from_sigma() returns NULL when sigma is infeasible under the
  # ||beta|| = 1 constraint, and a failed gam() fit is also possible for a
  # feasible sigma. Both cases are scored with a large-but-finite penalty
  # (neg_penalty) rather than -Inf: optimHess() later differentiates fn() by
  # finite differences at the selected optimum, and a finite-difference step
  # that lands on a -Inf evaluation produces a non-finite derivative and
  # aborts with "non-finite finite-difference value" -- even though the
  # optimum itself is a perfectly good, feasible point. A large finite
  # penalty is just as effective at steering Nelder-Mead away from bad
  # candidates during the search, but keeps every evaluation (and therefore
  # every finite difference) numerically well-defined.
  # neg_penalty is scaled off m0's logLik so it stays far below the range of
  # any real candidate fit regardless of the data's own log-likelihood scale.
  ref_loglik <- as.numeric(stats::logLik(m0))
  neg_penalty <- ref_loglik - 1e6 * (abs(ref_loglik) + 1)

  fn <- function(sigma) {
    utils::flush.console()
    beta_0 <- beta_from_sigma(sigma)
    if (is.null(beta_0)) return(neg_penalty)
    ll <- tryCatch(as.numeric(stats::logLik(fit_gam(beta_0))),
                   error = function(e) {
                     warning("fit_gam() failed for a candidate beta and was scored as a large ",
                             "penalty: ", conditionMessage(e), call. = FALSE)
                     neg_penalty
                   })
    if (!is.finite(ll)) neg_penalty else ll
  }

  fit_results <- vector("list", nrow(initial_sigma_table))
  crit_table <- data.frame(logLik = NA_real_, AIC = NA_real_, BIC = NA_real_)[rep(1, nrow(initial_sigma_table)), ]
  rownames(crit_table) <- rownames(initial_sigma_table)

  for (i in 1:nrow(initial_sigma_table)) {
    sigma_0 <- as.numeric(initial_sigma_table[i, ])
    # Nelder-Mead instead of L-BFGS-B: the sphere constraint ||sigma|| < 1 has
    # no natural per-coordinate box, and fn() returning -Inf outside it is
    # handled gracefully by a derivative-free simplex search without needing
    # gradient information near the (infeasible) boundary.
    m_temp <- stats::optim(par = sigma_0, fn = fn, method = "Nelder-Mead",
                           control = list(fnscale = -1, maxit = 100))
    fit_results[[i]] <- m_temp
    g <- fit_gam(beta_from_sigma(m_temp$par))
    crit_table$logLik[i] <- as.numeric(stats::logLik(g))
    crit_table$AIC[i] <- stats::AIC(g)
    crit_table$BIC[i] <- stats::BIC(g)
  }

  order_ob <- order(crit_table$BIC)[1]
  m_selected <- fit_results[[order_ob]]
  sigma_est <- m_selected$par
  # beta_1 > 0 and ||beta|| = 1 hold by construction here -- no post-hoc
  # normalization or sign flip needed (contrast with the old beta_BeforeNorm
  # approach, which had to renormalize and flip sign(beta_BeforeNorm[1])).
  beta_est <- beta_from_sigma(sigma_est)

  # Standard errors for beta: fn() re-fits a gam() at every evaluation, so no
  # analytic gradient/Hessian is available from optim() directly (hessian=TRUE
  # was skipped inside the search loop to avoid paying that cost for every
  # candidate direction). Recover the Hessian numerically at the selected
  # optimum only, the same way plsi.lr.v2() derives beta_sigma from optim()'s
  # Hessian -- but now in sigma-space (dimension p - 1), since sigma is the
  # free parameter optim() actually searched over.
  #
  # fn() itself is now finite everywhere (see neg_penalty above), which is the
  # main defense against optimHess()'s "non-finite finite-difference value"
  # error. This tryCatch is a second line of defense: if sigma_est still sits
  # close enough to the ||sigma|| = 1 boundary (equivalently, the anchor
  # exposure's beta_1 close to 0) that finite differencing is numerically
  # unstable for some other reason, fall back to a deliberately conservative
  # (large-variance) covariance instead of aborting the whole function --
  # consistent with how the eigenvalue floor below already treats Hessian
  # instability as something to degrade gracefully from, not error out on.
  hess <- tryCatch(
    stats::optimHess(par = sigma_est, fn = fn, control = list(fnscale = -1)),
    error = function(e) {
      warning("optimHess() could not compute a finite-difference Hessian at the selected optimum ",
              "(", conditionMessage(e), "). This typically means sigma_est is very close to the ",
              "||sigma|| = 1 boundary -- i.e. the anchor exposure's beta_1 is close to 0. Falling ",
              "back to a conservative (deliberately large-variance) covariance for beta; consider ",
              "a different anchor exposure (see the |correlation|-based X.name sort above) or ",
              "increasing initial.random.num for a more interior optimum.", call. = FALSE)
      NULL
    }
  )

  if (is.null(hess)) {
    cov_sigma <- diag(1e8, length(sigma_est))  # deliberately large fallback variance
  } else {
    info_matrix <- -hess

    eig_decomp <- eigen(info_matrix, symmetric = TRUE)
    eig_vals <- eig_decomp$values
    if (any(eig_vals <= 0)) {
      warning("Observed information matrix is not positive definite at the selected optimum ",
              "(min eigenvalue = ", signif(min(eig_vals), 3), "). Standard errors are computed from ",
              "an eigenvalue-floored version of the information matrix and should be treated as ",
              "conservative approximations; consider increasing initial.random.num for a more ",
              "reliable optimum.")
      # optimHess() differentiates fn(), which itself refits a penalized gam()
      # (an inner iterative optimization) at every evaluation -- small
      # discontinuities from that inner fit landing in slightly different places
      # for nearby sigma vectors can corrupt the finite-difference Hessian's
      # eigenvalues even when the underlying likelihood surface is well-behaved.
      # Floor tiny/negative eigenvalues to a small positive value (relative to
      # the largest eigenvalue) so the matrix is positive definite and
      # invertible, rather than silently producing NaN/nonsensical SEs from
      # ginv() on an indefinite matrix -- this makes the affected SEs
      # conservative (larger) rather than wrong.
      floor_val <- max(eig_vals) * 1e-6
      eig_vals <- pmax(eig_vals, floor_val)
      info_matrix <- eig_decomp$vectors %*% diag(eig_vals, nrow = length(eig_vals)) %*% t(eig_decomp$vectors)
    }
    cov_sigma <- suppressWarnings(MASS::ginv(info_matrix))
  }

  # Delta method: cov_sigma lives in sigma-space (dimension p - 1), but
  # si.coefficient reports beta (dimension p). Map the covariance across
  # beta = beta(sigma) via the Jacobian J = d(beta)/d(sigma):
  #   d(beta_1)/d(sigma_j)     = -sigma_j / beta_1   (row 1)
  #   d(beta_{j+1})/d(sigma_j) = 1, and 0 for k != j   (rows 2..p, = I)
  # so Cov(beta) ~= J %*% Cov(sigma) %*% t(J). Note this Jacobian divides by
  # beta_1, so it (and therefore beta_sigma) will legitimately blow up if
  # beta_1 is close to 0 -- that's the delta method correctly reflecting that
  # the sigma parameterization is ill-conditioned near that boundary, not a
  # bug; see the anchor-exposure note on si.coefficient if beta_sigma[1]
  # looks unreasonably large.
  jac <- rbind(-sigma_est / beta_est[1], diag(length(sigma_est)))
  cov_beta <- jac %*% cov_sigma %*% t(jac)
  beta_sigma <- suppressWarnings(sqrt(diag(cov_beta)))

  final_model <- fit_gam(beta_est)
  s <- summary(final_model)

  si.coefficient <- data.frame(Estimate = as.numeric(beta_est),
                               `Std.Error` = as.numeric(beta_sigma),
                               row.names = X.name, check.names = FALSE)
  si.coefficient$`t value` <- si.coefficient$Estimate / si.coefficient$`Std.Error`
  si.coefficient$`Pr(>|t|)` <- ifelse(2 * stats::pnorm(-abs(si.coefficient$`t value`)) < 0.0001, "<.0001",
                                      format(round(2 * stats::pnorm(-abs(si.coefficient$`t value`)), 4), nsmall = 4))
  si.coefficient$Lower.95CI <- si.coefficient$Estimate + stats::qnorm(0.025) * si.coefficient$`Std.Error`
  si.coefficient$Upper.95CI <- si.coefficient$Estimate + stats::qnorm(0.975) * si.coefficient$`Std.Error`
  si.coefficient$`Contribution proportion` <- format(round((si.coefficient$Estimate)^2, 3), nsmall = 3)
  # Sort by Estimate, descending (largest coefficient first) -- matches plsi.lr.v2()
  si.coefficient <- si.coefficient[order(si.coefficient$Estimate, decreasing = TRUE), ]

  # Note on how the anchor exposure (beta_1, constrained > 0) was chosen.
  # X.name[1] / cor_linear$exposure[1] is the exposure with the largest
  # |correlation| with Y.name (see the X.name sort above); attach that as a
  # readable note on si.coefficient so it's visible alongside the coefficient
  # table rather than only inferable from which row has the smallest index.
  anchor.exposure <- X.name[1]
  anchor.cor <- cor_linear$cor_with_Y[cor_linear$exposure == anchor.exposure]
  anchor.note <- paste0(
    "First/anchor exposure (single-index coefficient constrained > 0 for identifiability, ",
    "||beta|| = 1): '", anchor.exposure, "', selected as the candidate exposure with the ",
    "largest absolute correlation with ", Y.name, " (r = ", signif(anchor.cor, 3), ")."
  )
  attr(si.coefficient, "note") <- anchor.note

  confounder.coefficient <- as.data.frame(s$p.table[Z.name, , drop = FALSE])
  confounder.coefficient$Lower.95CI <- confounder.coefficient$Estimate + stats::qnorm(0.025) * confounder.coefficient$`Std. Error`
  confounder.coefficient$Upper.95CI <- confounder.coefficient$Estimate + stats::qnorm(0.975) * confounder.coefficient$`Std. Error`
  confounder.coefficient$`Pr(>|t|)` <- ifelse(confounder.coefficient$`Pr(>|t|)` < 0.0001, "<.0001",
                                              format(round(confounder.coefficient$`Pr(>|t|)`, 4), nsmall = 4))

  single_index_estimated <- as.vector(x %*% beta_est)

  # Confounder-free single-index model for prediction -------------------------
  # `final_model` still has AGE.c / SEX.Female / etc. in its formula, so
  # predict() on it requires those columns in newdata. That's fine for the
  # in-sample diagnostics above, but breaks any downstream use that only has
  # an index value and no confounder values to supply -- e.g. plotting the
  # response curve for a single exposure via beta[exposure] * x_value.
  # Mirror plsi.lr.v2()'s approach: residualize y on the confounders using
  # final_model's parametric coefficients, then fit a second, confounder-free
  # smooth of the residualized outcome on the single index alone. Predict from
  # THIS model (si.fun.model) for any new/custom single-index values.
  confounder_coefs <- stats::coef(final_model)[Z.name]
  y_single_index <- y - as.vector(as.matrix(z_df) %*% confounder_coefs)

  dat_si_fun <- data.frame(single_index_estimated = single_index_estimated,
                           y_single_index = y_single_index)
  form_si <- stats::as.formula(paste0("y_single_index ~ s(single_index_estimated, bs='", bs, "', k=", k, ")"))
  m2 <- mgcv::gam(form_si, data = dat_si_fun, method = "REML")

  # NOTE: ciTools::add_ci() does not support "gam" objects from mgcv (it
  # covers lm/glm/lmerMod/glmerMod) -- that mismatch is what produced the
  # "object 'AGE.c' not found" error upstream. Use predict.gam(se.fit=TRUE)
  # directly and build the Wald interval by hand instead. This also sidesteps
  # the predict(type="terms") column-naming fragility entirely, since m2 has
  # only one smooth term.
  ord <- order(single_index_estimated)
  pred2 <- stats::predict(m2, newdata = dat_si_fun[ord, "single_index_estimated", drop = FALSE],
                          type = "link", se.fit = TRUE)
  si.fun <- data.frame(single_index_estimated = single_index_estimated[ord],
                       y_single_index = y_single_index[ord],
                       fit = pred2$fit,
                       se  = pred2$se.fit)
  si.fun$lwr <- si.fun$fit - stats::qnorm(0.975) * si.fun$se
  si.fun$upr <- si.fun$fit + stats::qnorm(0.975) * si.fun$se

  return(list(original.data = list(y = y, x = x, z = z_df),
              original.par = list(k = k, bs = bs, initial.random.num = initial.random.num, seed = seed),
              si.coefficient = si.coefficient,          # carries attr(., "note") on anchor-exposure selection
              anchor.note = anchor.note,                # same note, surfaced directly for convenience
              model.statistics = crit_table[order_ob, ],
              initial.table = initial_table,
              initial.sigma.table = initial_sigma_table,
              confounder.coefficient = confounder.coefficient,
              si.fun.edf = s$edf[1],          # effective df actually used -- this is the "auto-tuned" answer
              si.fun = si.fun,
              si.fun.model = m2,              # confounder-free: predict() only needs single_index_estimated
              full.model = final_model))      # full fit incl. confounders, kept for reference/diagnostics
}