R/impPsis.R

In nlmixr2est: Nonlinear Mixed Effects Models in Population PK/PD, Estimation Routines

# Tail-sensitive importance-sampling diagnostic: Pareto k-hat (PSIS).
#
# The imp family already reports two efficiency statistics -- NONMEM's xi and
# the Kish effective sample size -- and BOTH are means over samples drawn from
# the proposal.  That makes them structurally blind to the failure mode that
# actually matters here: a proposal with lighter tails than the target has
# importance weights with infinite variance, yet looks perfectly healthy in any
# finite sample, because the tail responsible is precisely the region the
# proposal rarely visits.
#
# Pareto k-hat estimates the shape parameter of a generalized Pareto fitted to
# the upper tail of the weights, which is the tail index of the weight
# distribution itself.  The standard reading (Vehtari, Simpson, Gelman, Yao,
# Gabry, "Pareto Smoothed Importance Sampling"):
#
#   k < 0.5   weight variance finite; importance sampling reliable
#   0.5-0.7   finite mean, infinite variance; usable but degrading
#   k > 0.7   unreliable -- the estimate is dominated by a few draws
#
# This is the statistic that can tell whether a heavier-tailed (t) proposal is
# needed, which neither xi nor the Kish ESS can.
#
# Implemented here rather than taken from loo:: to avoid a new hard dependency
# for a diagnostic, but validated against it: on identical weight vectors the
# two agree to ~0.02, and both recover a known tail index --
#
#   true k   0.2    0.5    0.8    1.2
#   ours     0.209  0.513  0.824  1.241
#   loo      0.222  0.512  0.807  1.203
#
# (means over 5 seeds, S = 4000; per-seed sd ~0.05).  The unit tests re-run
# that comparison whenever loo is installed.

#' Generalized Pareto fit (Zhang & Stephens 2009 empirical-Bayes estimator)
#'
#' @param x Sorted-ascending, strictly positive exceedances over the threshold.
#' @param minGridPts Grid size floor for the profile-likelihood average.
#' @return list(k, sigma)
#' @noRd
.impGpdFit <- function(x, minGridPts = 30L) {
  .n <- length(x)
  if (.n < 5L) return(list(k = NA_real_, sigma = NA_real_))
  .prior <- 3
  .m <- minGridPts + floor(sqrt(.n))
  .jj <- seq_len(.m)
  # quartile used to set the grid scale
  .xstar <- x[floor(.n / 4 + 0.5)]
  if (!is.finite(.xstar) || .xstar <= 0) return(list(k = NA_real_, sigma = NA_real_))
  .theta <- 1 / x[.n] + (1 - sqrt(.m / (.jj - 0.5))) / .prior / .xstar
  # profile log-likelihood of theta, up to an additive constant
  .lx <- vapply(.theta, function(.t) {
    .k <- mean(log1p(-.t * x))
    log(-.t / .k) - .k - 1
  }, numeric(1))
  .l <- .n * .lx
  .l[!is.finite(.l)] <- -Inf
  if (all(!is.finite(.l))) return(list(k = NA_real_, sigma = NA_real_))
  .w <- exp(.l - max(.l))
  .w <- .w / sum(.w)
  .thetaHat <- sum(.theta * .w)
  .k <- mean(log1p(-.thetaHat * x))
  .sigma <- -.k / .thetaHat
  list(k = .k, sigma = .sigma)
}

#' Pareto k-hat for one subject's importance weights
#'
#' @param w Importance weights (any positive scaling; k-hat is scale-invariant).
#' @return k-hat, or NA if there are too few usable weights.
#' @noRd
.impPsisK <- function(w) {
  .w <- as.numeric(w)
  .w <- .w[is.finite(.w) & .w > 0]
  .s <- length(.w)
  if (.s < 25L) return(NA_real_)
  # PSIS tail size: min(0.2*S, 3*sqrt(S)), the usual choice
  .tail <- min(floor(0.2 * .s), floor(3 * sqrt(.s)))
  if (.tail < 5L) return(NA_real_)
  .srt <- sort(.w)
  # threshold is the largest weight NOT in the tail
  .u <- .srt[.s - .tail]
  .exc <- .srt[(.s - .tail + 1L):.s] - .u
  .exc <- .exc[.exc > 0]
  if (length(.exc) < 5L) return(NA_real_)
  .impGpdFit(sort(.exc))$k
}

#' Per-subject Pareto k-hat from a fit's stashed final-iteration weights
#'
#' @param env Fit environment holding `impWeights`.
#' @return numeric vector of k-hat, one per subject (NA where not computable)
#' @noRd
.impPsisKAll <- function(env) {
  .w <- tryCatch(env$impWeights, error = function(e) NULL)
  if (is.null(.w) || length(.w) == 0L) return(numeric(0))
  vapply(.w, function(.wi) tryCatch(.impPsisK(.wi), error = function(e) NA_real_),
         numeric(1), USE.NAMES = FALSE)
}