R/method-sampling.R

In INLAvaan: Approximate Bayesian Latent Variable Analysis

#' Draw Samples from the Generative Model
#'
#' Sample model parameters, latent variables, or observed variables from the
#' generative model underlying a fitted INLAvaan model. By default, parameters
#' are drawn from the **posterior** distribution; set `prior = TRUE` to draw
#' from the **prior** instead (useful for prior predictive checks).
#'
#' @details
#' Each row of the output corresponds to a **fresh parameter draw**: a new
#' \eqn{\boldsymbol\theta^{(s)}} is sampled and then propagated through the
#' generative chain to produce one latent vector and one observed vector. This
#' makes `sampling()` ideal for **prior and posterior predictive checks**
#' (e.g., density overlays, test statistic distributions).
#'
#' The generative chain is:
#' \deqn{\boldsymbol\theta^{(s)} \sim \pi(\boldsymbol\theta \mid \mathbf{y})}
#' \deqn{\boldsymbol\eta^{(s)} \sim N((\mathbf{I} - \mathbf{B})^{-1}\boldsymbol\alpha,\,\boldsymbol\Phi)}
#' \deqn{\mathbf{y}^{*(s)} \sim N(\boldsymbol\Lambda\boldsymbol\eta^{(s)} + \boldsymbol\nu,\,\boldsymbol\Theta)}
#'
#' If you need **complete replicate datasets** (many observations from a single
#' parameter draw) — for example, for simulation-based calibration (SBC) — use
#' [simulate()] instead.
#'
#' This is distinct from [predict()], which computes individual-specific
#' factor scores \eqn{\boldsymbol\eta \mid \mathbf{y},\boldsymbol\theta}
#' conditional on observed data.
#'
#' @param object An object of class [INLAvaan] (or `inlavaan_internal`).
#' @param type Character string specifying what to sample:
#'   \describe{
#'     \item{`"lavaan"`}{(Default) The lavaan-side (constrained) model
#'       parameters. Returns an `nsamp` by `npar` matrix.}
#'     \item{`"theta"`}{The INLAvaan-side unconstrained parameters.
#'       Returns an `nsamp` by `npar` matrix.}
#'     \item{`"latent"`}{Latent variables from the model-implied
#'       distribution. Returns an `nsamp` by `nlv` matrix (one draw per
#'       posterior sample, not tied to any individual). For two-level
#'       models the matrix holds the within- *and* between-level latent
#'       variables, the level-2 columns carrying the `.l2` suffix when the
#'       same latent variable also exists at level 1.}
#'     \item{`"observed"`}{Observed variables generated from the full
#'       model. Returns an `nsamp` by `nobs_vars` matrix. For two-level
#'       models each row is a draw from the two-level generative model,
#'       \eqn{\mathbf{y} = \mathbf{y}^B + \mathbf{y}^W}: variables that
#'       live at both levels sum their between- and within-level draws,
#'       and within-only or between-only variables take the single level
#'       available to them.}
#'     \item{`"implied"`}{Model-implied moments. Returns a length-`nsamp`
#'       list, each element a list with `cov` (model-implied covariance
#'       matrix) and, when `meanstructure = TRUE`, `mean` (model-implied
#'       mean vector). For multi-group models each element is itself a
#'       list of groups. For two-level models each element is a list with
#'       a `within` and a `cluster` block, each holding a `cov` and a
#'       `mean`, as [lavaan::lavInspect()] reports them.}
#'     \item{`"all"`}{A named list with elements `lavaan`, `theta`,
#'       `latent`, `observed`, and `implied`.}
#'   }
#' @param nsamp Number of samples to draw.
#' @param samp_copula Logical. When `TRUE` (default), posterior parameter
#'   samples use the copula method with the fitted marginals. When `FALSE`,
#'   samples are drawn from the joint Gaussian (Laplace) approximation.
#'   Ignored when `prior = TRUE`.
#' @param prior Logical. When `TRUE`, parameters are drawn from the prior
#'   distribution and then propagated through the generative model. When
#'   `FALSE` (default), parameters come from the posterior.
#' @param silent Logical. When `TRUE`, suppresses the informational message
#'   about rejected non-PD draws during prior rejection sampling. Default
#'   `FALSE`.
#' @param ... Additional arguments (currently unused).
#'
#' @returns A matrix or named list, depending on `type`.
#'
#' @seealso [simulate()] for generating complete replicate datasets (e.g.,
#'   for SBC); [predict()] for individual-specific factor scores;
#'   [bfit_indices()] for Bayesian fit indices.
#'
#' @example inst/examples/ex-sampling.R
#' @export
setGeneric("sampling", function(object, ...) standardGeneric("sampling"))

#' @name sampling
#' @rdname sampling
#' @aliases sampling,INLAvaan-method
#' @export
setMethod(
  "sampling",
  "INLAvaan",
  function(
    object,
    type = c("lavaan", "theta", "latent", "observed", "implied", "all"),
    nsamp = 1000L,
    samp_copula = TRUE,
    prior = FALSE,
    silent = FALSE,
    ...
  ) {
    sampling_impl(
      object@external$inlavaan_internal,
      type = type,
      nsamp = nsamp,
      samp_copula = samp_copula,
      prior = prior,
      meanstructure = isTRUE(object@Options$meanstructure),
      silent = silent,
      ...
    )
  }
)

#' @exportS3Method sampling inlavaan_internal
sampling.inlavaan_internal <- function(
  object,
  type = c("lavaan", "theta", "latent", "observed", "implied", "all"),
  nsamp = 1000L,
  samp_copula = TRUE,
  prior = FALSE,
  silent = FALSE,
  ...
) {
  sampling_impl(
    object,
    type = type,
    nsamp = nsamp,
    samp_copula = samp_copula,
    prior = prior,
    silent = silent,
    ...
  )
}

# ---- Internal: draw parameter samples (theta/x) -----------------------------

sample_params_posterior <- function(int, nsamp, samp_copula) {
  method <- if (isTRUE(samp_copula)) int$marginal_method else "sampling"
  sample_params(
    theta_star = int$theta_star,
    Sigma_theta = int$Sigma_theta,
    method = method,
    approx_data = int$approx_data,
    pt = int$partable,
    lavmodel = int$lavmodel,
    nsamp = nsamp,
    R_star = int$R_star
  )
}

# Draw from the prior: independent draws per free parameter, map to x-space.
sample_params_prior <- function(int, nsamp) {
  pt <- int$partable
  lavmodel <- int$lavmodel

  # All unique free parameter indices in the partable
  PTFREEIDX <- which(pt$free > 0L & !duplicated(pt$free))
  m <- length(PTFREEIDX)

  # Draw in the interpretable/natural scale for each parameter
  x_natural <- matrix(NA_real_, nrow = nsamp, ncol = m)

  for (j in seq_len(m)) {
    prior_str <- pt$prior[PTFREEIDX[j]]

    if (is.na(prior_str) || prior_str == "") {
      # Fallback: wide normal for params without an explicit prior
      x_natural[, j] <- stats::rnorm(nsamp, 0, 10) # nocov
    } else if (grepl("^normal", prior_str)) {
      par <- as.numeric(strsplit(
        gsub("normal\\(|\\)", "", prior_str),
        ","
      )[[1]])
      x_natural[, j] <- stats::rnorm(nsamp, par[1], par[2])
    } else if (grepl("^gamma", prior_str)) {
      is_sd <- grepl("\\[sd\\]", prior_str)
      is_prec <- grepl("\\[prec\\]", prior_str)
      par <- as.numeric(strsplit(
        gsub("gamma\\(|\\)|\\[sd\\]|\\[prec\\]", "", prior_str),
        ","
      )[[1]])
      raw <- stats::rgamma(nsamp, shape = par[1], rate = par[2])
      if (is_sd) {
        x_natural[, j] <- raw^2 # SD → variance  # nocov
      } else if (is_prec) {
        x_natural[, j] <- 1 / raw # precision → variance  # nocov
      } else {
        x_natural[, j] <- raw
      }
    } else if (grepl("^beta", prior_str)) {
      # nocov start
      par <- as.numeric(strsplit(
        gsub("beta\\(|\\)", "", prior_str),
        ","
      )[[1]])
      x_natural[, j] <- stats::rbeta(nsamp, par[1], par[2]) * 2 - 1
    } else {
      x_natural[, j] <- stats::rnorm(nsamp, 0, 10)
    } # nocov end
  }

  # Map natural scale → unconstrained theta-space via g()
  theta_samp <- x_natural
  for (j in seq_len(m)) {
    theta_samp[, j] <- vapply(
      x_natural[, j],
      pt$g[[PTFREEIDX[j]]],
      numeric(1)
    )
  }

  # Apply equality constraints if present
  if (lavmodel@ceq.simple.only) {
    # nocov start
    K <- lavmodel@ceq.simple.K
    theta_samp <- t(apply(theta_samp, 1, function(p) as.numeric(K %*% p)))
  } # nocov end

  # Map theta → lavaan x-space (handles covariance = cor * sqrt(var1 * var2))
  x_samp <- t(apply(theta_samp, 1, pars_to_x, pt = pt))

  list(theta_samp = theta_samp, x_samp = x_samp)
}

# ---- Internal: generate eta from model-implied distribution ------------------

# Cholesky factor of a covariance block. Outside strict mode a non-PD block is
# projected onto the nearest PD matrix. Prior rejection sampling asks for
# strict = TRUE so that the error propagates and the draw is rejected.
chol_cov_block <- function(S, strict = FALSE) {
  if (strict) {
    t(chol(S)) # nocov - error propagates if non-PD
  } else {
    tryCatch(t(chol(S)), error = function(e) t(chol(make_pd(S))))
  }
}

draw_latent_block <- function(glist, strict = FALSE) {
  Psi <- glist$psi
  B <- glist$beta
  alpha <- glist$alpha

  IminB <- if (is.null(B)) diag(nrow(Psi)) else (diag(nrow(B)) - B)
  if (is.null(alpha)) {
    alpha <- rep(0, nrow(Psi))
  }
  IminB_inv <- solve(IminB)

  mu_eta <- as.numeric(IminB_inv %*% alpha)
  Phi <- IminB_inv %*% Psi %*% t(IminB_inv)
  chol_Phi <- chol_cov_block(Phi, strict)

  eta <- mu_eta + as.numeric(chol_Phi %*% stats::rnorm(length(mu_eta)))
  names(eta) <- colnames(Psi)
  eta
}

sample_latent_from_model <- function(x_row, lavmodel, strict = FALSE) {
  GLIST <- get_SEM_param_matrix(x_row, "all", lavmodel)
  nG <- lavmodel@ngroups
  eta_list <- lapply(seq_len(nG), function(g) {
    draw_latent_block(GLIST[[g]], strict = strict)
  })

  if (nG == 1L) eta_list[[1L]] else eta_list
}

# ---- Internal: generate y from model given eta ------------------------------

draw_observed_block <- function(glist, eta, strict = FALSE) {
  Lambda <- glist$lambda
  Theta <- glist$theta
  nu <- glist$nu

  if (is.null(nu)) {
    nu <- rep(0, nrow(Lambda))
  }

  mu_y <- as.numeric(Lambda %*% eta + nu)
  chol_Theta <- chol_cov_block(Theta, strict)
  y <- mu_y + as.numeric(chol_Theta %*% stats::rnorm(length(mu_y)))
  names(y) <- rownames(Lambda)
  y
}

sample_observed_from_model <- function(x_row, eta, lavmodel, strict = FALSE) {
  GLIST <- get_SEM_param_matrix(x_row, "all", lavmodel)
  nG <- lavmodel@ngroups
  y_list <- lapply(seq_len(nG), function(g) {
    eta_g <- if (nG == 1L) eta else eta[[g]]
    draw_observed_block(GLIST[[g]], eta_g, strict = strict)
  })

  if (nG == 1L) y_list[[1L]] else y_list
}

# ---- Internal: compute model-implied moments --------------------------------

compute_implied_moments <- function(x_row, lavmodel, meanstructure = FALSE) {
  GLIST <- get_SEM_param_matrix(x_row, "all", lavmodel)
  nG <- lavmodel@ngroups
  out_list <- vector("list", nG)

  for (g in seq_len(nG)) {
    glist <- GLIST[[g]]
    Lambda <- glist$lambda
    Psi <- glist$psi
    Theta <- glist$theta
    B <- glist$beta
    alpha <- glist$alpha
    nu <- glist$nu

    IminB <- if (is.null(B)) diag(nrow(Psi)) else (diag(nrow(B)) - B)
    IminB_inv <- solve(IminB)

    # Sigma_y = Lambda (I-B)^{-1} Psi [(I-B)^{-1}]' Lambda' + Theta
    front <- Lambda %*% IminB_inv
    Sigma_y <- front %*% Psi %*% t(front) + Theta
    rownames(Sigma_y) <- colnames(Sigma_y) <- rownames(Lambda)

    res <- list(cov = Sigma_y)

    if (meanstructure) {
      # nocov start
      if (is.null(alpha)) {
        alpha <- rep(0, nrow(Psi))
      }
      if (is.null(nu)) {
        nu <- rep(0, nrow(Lambda))
      }
      mu_y <- as.numeric(Lambda %*% IminB_inv %*% alpha + nu)
      names(mu_y) <- rownames(Lambda)
      res$mean <- mu_y
    } # nocov end

    out_list[[g]] <- res
  }

  if (nG == 1L) out_list[[1L]] else out_list # nocov (else = multigroup)
}

# ---- Internal: two-level generative draws ------------------------------------
#
# A two-level fit stores nblocks = ngroups * nlevels sets of model matrices,
# block (g - 1) * nlevels + l holding level l of group g. Slicing the GLIST by
# group alone, as get_SEM_param_matrix() does, would return the within block
# only, so the helpers below address the blocks themselves through the number
# of matrices per block recorded in lavmodel@nmat.

get_block_param_matrix <- function(x_row, lavmodel) {
  lavmodel_x <- lavaan::lav_model_set_parameters(lavmodel, x_row)
  offset <- cumsum(c(0L, lavmodel_x@nmat))

  lapply(seq_len(lavmodel_x@nblocks), function(b) {
    mm <- seq_len(lavmodel_x@nmat[b]) + offset[b]
    glist <- Map(
      function(mat, dn) {
        rownames(mat) <- dn[[1]]
        colnames(mat) <- dn[[2]]
        mat
      },
      lavmodel_x@GLIST[mm],
      lavmodel_x@dimNames[mm]
    )
    names(glist) <- names(lavmodel_x@GLIST)[mm]
    glist
  })
}

# Latent variable names across levels, flattened into a single vector. A
# level-2 name takes the ".l2" suffix of the coefficient labels, but only when
# the same latent variable also exists at level 1.
ml_latent_names <- function(name_list) {
  out <- name_list
  for (l in seq_along(name_list)[-1L]) {
    seen <- unlist(name_list[seq_len(l - 1L)], use.names = FALSE)
    dup <- out[[l]] %in% seen
    out[[l]][dup] <- paste0(out[[l]][dup], ".l", l)
  }
  unlist(out, use.names = FALSE)
}

# One draw of the latent and (optionally) observed vectors from the two-level
# generative model. An observed variable that lives at both levels is the sum
# of its between- and within-level draws, and a variable that lives at one
# level only takes that level's draw.
sample_generative_ml <- function(
  x_row,
  lavmodel,
  lavdata,
  need_obs = TRUE,
  strict = FALSE
) {
  GLIST <- get_block_param_matrix(x_row, lavmodel)
  nG <- lavmodel@ngroups
  nlevels <- lavdata@nlevels
  eta_list <- vector("list", nG)
  y_list <- vector("list", nG)

  for (g in seq_len(nG)) {
    ov_names <- lavdata@ov.names[[g]]
    y_g <- rep(0, length(ov_names))
    names(y_g) <- ov_names
    eta_g <- vector("list", nlevels)

    for (l in seq_len(nlevels)) {
      glist <- GLIST[[(g - 1) * nlevels + l]]
      eta_g[[l]] <- draw_latent_block(glist, strict = strict)
      if (need_obs) {
        y_l <- draw_observed_block(glist, eta_g[[l]], strict = strict)
        y_g[names(y_l)] <- y_g[names(y_l)] + y_l
      }
    }

    eta_list[[g]] <- stats::setNames(
      unlist(eta_g, use.names = FALSE),
      ml_latent_names(lapply(eta_g, names))
    )
    y_list[[g]] <- y_g
  }

  if (nG > 1L) {
    # nocov start
    suffix <- function(v, g) stats::setNames(v, paste0(names(v), ".g", g))
    eta_list <- Map(suffix, eta_list, seq_len(nG))
    y_list <- Map(suffix, y_list, seq_len(nG))
  } # nocov end

  list(
    latent = unlist(eta_list),
    observed = if (need_obs) unlist(y_list) else NULL
  )
}

# Model-implied moments of a two-level model, one within/between pair per
# group, named and ordered as lavInspect(object, "implied") reports them.
compute_implied_moments_ml <- function(x_row, lavmodel, lavdata) {
  lavmodel_x <- lavaan::lav_model_set_parameters(lavmodel, x_row)
  implied <- lavaan::lav_model_implied(lavmodel_x)
  nG <- lavmodel@ngroups
  nlevels <- lavdata@nlevels
  out_list <- vector("list", nG)

  for (g in seq_len(nG)) {
    Lp <- lavdata@Lp[[g]]
    blocks <- (g - 1) * nlevels + seq_len(nlevels)
    res <- vector("list", nlevels)

    for (l in seq_len(nlevels)) {
      ov_names <- Lp$ov.names[Lp$ov.idx[[l]]]
      Sigma_y <- implied$cov[[blocks[l]]]
      dimnames(Sigma_y) <- list(ov_names, ov_names)
      # A two-level model always carries a mean structure (the within-level
      # means are zero), so both blocks report a mean vector.
      mu_y <- as.numeric(implied$mean[[blocks[l]]])
      names(mu_y) <- ov_names
      res[[l]] <- list(cov = Sigma_y, mean = mu_y)
    }

    names(res) <- lavdata@block.label[blocks]
    out_list[[g]] <- res
  }

  if (nG == 1L) out_list[[1L]] else out_list # nocov (else = multigroup)
}

# The two-level counterpart of the generative steps in sampling_impl(): the
# implied moments, the latent draws and the observed draws all span both
# levels.
sampling_generative_ml <- function(int, samp, type, nsamp) {
  lavmodel <- int$lavmodel
  lavdata <- int$lavdata

  implied_list <- NULL
  if (type == "implied" || type == "all") {
    implied_list <- lapply(seq_len(nsamp), function(i) {
      compute_implied_moments_ml(samp$x_samp[i, ], lavmodel, lavdata)
    })
    if (type == "implied") {
      return(implied_list)
    }
  }

  need_obs <- type %in% c("observed", "all")
  draws <- lapply(seq_len(nsamp), function(i) {
    sample_generative_ml(
      samp$x_samp[i, ],
      lavmodel,
      lavdata,
      need_obs = need_obs
    )
  })

  eta_mat <- do.call(rbind, lapply(draws, `[[`, "latent"))
  if (type == "latent") {
    return(eta_mat)
  }

  y_mat <- do.call(rbind, lapply(draws, `[[`, "observed"))
  if (type == "observed") {
    return(y_mat)
  }

  # type == "all"
  list(
    lavaan = samp$x_samp,
    theta = samp$theta_samp,
    latent = eta_mat,
    observed = y_mat,
    implied = implied_list
  )
}

# ---- Internal: prior generative sampling with reject-and-redraw --------------
#
# When prior = TRUE and we need latent/observed draws, parameter vectors that
# produce non-positive-definite model-implied covariance matrices are rejected
# and redrawn.  This preserves the exact prior distribution rather than silently
# projecting non-PD matrices to PD space via make_pd().

sampling_prior_generative <- function(
  int,
  type,
  nsamp,
  meanstructure = FALSE,
  silent = FALSE
) {
  pt <- int$partable
  xnames <- pt$names[pt$free > 0 & !duplicated(pt$free)]
  lavmodel <- int$lavmodel
  lavdata <- int$lavdata
  nG <- lavmodel@ngroups
  two_level <- is_multilevel(lavdata)

  # For 'implied' alone, no Cholesky decomposition is needed — just sample
  # parameters and compute the moments directly (no rejection required).
  if (type == "implied") {
    samp <- sample_params_prior(int, nsamp)
    colnames(samp$x_samp) <- xnames
    return(lapply(seq_len(nsamp), function(i) {
      if (two_level) {
        compute_implied_moments_ml(samp$x_samp[i, ], lavmodel, lavdata)
      } else {
        compute_implied_moments(samp$x_samp[i, ], lavmodel, meanstructure)
      }
    }))
  }

  need_obs <- type %in% c("observed", "all")
  need_implied <- type == "all"

  # Pre-compute dimensions from a single draw
  samp0 <- sample_params_prior(int, 1L)
  if (two_level) {
    draw0 <- sample_generative_ml(samp0$x_samp[1, ], lavmodel, lavdata)
    eta_cn <- names(draw0$latent)
    y_cn <- names(draw0$observed)
  } else {
    GLIST0 <- get_SEM_param_matrix(samp0$x_samp[1, ], "all", lavmodel)
    nlv <- ncol(GLIST0[[1]]$psi)
    nobs <- nrow(GLIST0[[1]]$lambda)
    lv_names <- colnames(GLIST0[[1]]$psi)
    ov_names <- rownames(GLIST0[[1]]$lambda)

    # Column names for output matrices
    if (nG == 1L) {
      eta_cn <- lv_names
      y_cn <- ov_names
    } else {
      eta_cn <- paste0(rep(lv_names, nG), ".g", rep(seq_len(nG), each = nlv)) # nocov
      y_cn <- paste0(rep(ov_names, nG), ".g", rep(seq_len(nG), each = nobs)) # nocov
    }
  }

  # Pre-allocate storage
  npar <- length(xnames)
  x_mat <- matrix(NA_real_, nsamp, npar)
  theta_mat <- matrix(NA_real_, nsamp, npar)
  eta_mat <- matrix(NA_real_, nsamp, length(eta_cn))
  y_mat <- if (need_obs) matrix(NA_real_, nsamp, length(y_cn)) else NULL
  colnames(x_mat) <- xnames
  colnames(theta_mat) <- xnames
  colnames(eta_mat) <- eta_cn
  if (need_obs) {
    colnames(y_mat) <- y_cn
  }
  implied_list <- if (need_implied) vector("list", nsamp) else NULL

  max_attempts <- nsamp * 20L # tolerate up to ~95% rejection rate
  collected <- 0L
  attempts <- 0L

  while (collected < nsamp && attempts < max_attempts) {
    # Draw a batch of parameters (draw what we still need)
    batch_n <- nsamp - collected
    samp_batch <- sample_params_prior(int, batch_n)

    for (i in seq_len(batch_n)) {
      attempts <- attempts + 1L
      if (attempts > max_attempts) {
        # nocov
        break
      }

      x1 <- samp_batch$x_samp[i, ]

      # Try the generative draw (strict = TRUE: no make_pd fallback). The
      # two-level draw already returns both levels in one flat vector.
      if (two_level) {
        draw1 <- tryCatch(
          sample_generative_ml(
            x1,
            lavmodel,
            lavdata,
            need_obs = need_obs,
            strict = TRUE
          ),
          error = function(e) NULL
        )
        if (is.null(draw1)) {
          next
        } # nocov
        eta1 <- draw1$latent
        y1 <- draw1$observed
      } else {
        eta1 <- tryCatch(
          sample_latent_from_model(x1, lavmodel, strict = TRUE),
          error = function(e) NULL
        )
        if (is.null(eta1)) {
          # nocov
          next
        }

        # Try observed draw if needed
        if (need_obs) {
          y1 <- tryCatch(
            sample_observed_from_model(x1, eta1, lavmodel, strict = TRUE),
            error = function(e) NULL
          )
          if (is.null(y1)) next # nocov
        }
      }

      # Valid draw -- store it
      collected <- collected + 1L
      x_mat[collected, ] <- x1
      theta_mat[collected, ] <- samp_batch$theta_samp[i, ]

      if (two_level || nG == 1L) {
        eta_mat[collected, ] <- eta1
        if (need_obs) y_mat[collected, ] <- y1
      } else {
        # nocov start
        eta_mat[collected, ] <- unlist(eta1)
        if (need_obs) y_mat[collected, ] <- unlist(y1)
      } # nocov end

      if (need_implied) {
        implied_list[[collected]] <- if (two_level) {
          compute_implied_moments_ml(x1, lavmodel, lavdata) # nocov
        } else {
          compute_implied_moments(x1, lavmodel, meanstructure)
        }
      }

      if (collected >= nsamp) break
    }
  }

  rejected <- attempts - collected
  if (rejected > 0L && !isTRUE(silent)) {
    # nocov start
    rej_pct <- round(100 * rejected / attempts, 1)
    cli_inform(
      "Prior sampling: {rejected} of {attempts} draw{?s} ({rej_pct}%) rejected (non-PD model-implied covariance)."
    )
  } # nocov end

  if (collected < nsamp) {
    # nocov start
    cli_warn(c(
      "Prior rejection sampling fell short of the requested sample size.",
      "i" = "Only {collected} of {nsamp} samples obtained after {attempts} attempts.",
      "i" = "Consider using more informative priors."
    ))
    if (collected == 0L) {
      cli_abort("No valid prior draws obtained. Priors may be too vague.")
    }
    # Trim to valid rows
    x_mat <- x_mat[seq_len(collected), , drop = FALSE]
    theta_mat <- theta_mat[seq_len(collected), , drop = FALSE]
    eta_mat <- eta_mat[seq_len(collected), , drop = FALSE]
    if (need_obs) {
      y_mat <- y_mat[seq_len(collected), , drop = FALSE]
    }
    if (need_implied) implied_list <- implied_list[seq_len(collected)]
  } # nocov end

  if (type == "latent") {
    return(eta_mat)
  }
  if (type == "observed") {
    return(y_mat)
  }

  # type == "all"
  list(
    lavaan = x_mat,
    theta = theta_mat,
    latent = eta_mat,
    observed = y_mat,
    implied = implied_list
  )
}

# ---- Main workhorse ----------------------------------------------------------

sampling_impl <- function(
  int,
  type = c("lavaan", "theta", "latent", "observed", "implied", "all"),
  nsamp = 1000L,
  samp_copula = TRUE,
  prior = FALSE,
  meanstructure = FALSE,
  silent = FALSE,
  ...
) {
  type <- match.arg(type)

  # For prior sampling with generative draws, use reject-and-redraw to preserve
  # the exact prior (no silent PD projection).
  if (isTRUE(prior) && type %in% c("latent", "observed", "implied", "all")) {
    return(sampling_prior_generative(
      int,
      type,
      nsamp,
      meanstructure,
      silent = silent
    ))
  }

  # Step 1: draw parameters
  if (isTRUE(prior)) {
    samp <- sample_params_prior(int, nsamp)
  } else {
    samp <- sample_params_posterior(int, nsamp, samp_copula)
  }

  pt <- int$partable
  xnames <- pt$names[pt$free > 0 & !duplicated(pt$free)]
  lavmodel <- int$lavmodel

  colnames(samp$x_samp) <- xnames
  colnames(samp$theta_samp) <- xnames

  # Early return for parameter-only types
  if (type == "lavaan") {
    return(samp$x_samp)
  }
  if (type == "theta") {
    return(samp$theta_samp)
  }

  # A two-level fit generates a within-level and a between-level quantity for
  # each of the types below, so it takes the two-level generative path.
  if (is_multilevel(int$lavdata)) {
    return(sampling_generative_ml(int, samp, type, nsamp))
  }

  # Compute model-implied moments if requested
  if (type == "implied" || type == "all") {
    implied_list <- lapply(seq_len(nsamp), function(i) {
      compute_implied_moments(samp$x_samp[i, ], lavmodel, meanstructure)
    })
    if (type == "implied") return(implied_list)
  }

  # Pre-compute dimensions from the first draw
  nG <- lavmodel@ngroups
  GLIST0 <- get_SEM_param_matrix(samp$x_samp[1, ], "all", lavmodel)
  nlv <- ncol(GLIST0[[1]]$psi)
  nobs <- nrow(GLIST0[[1]]$lambda)
  lv_names <- colnames(GLIST0[[1]]$psi)
  ov_names <- rownames(GLIST0[[1]]$lambda)

  # Step 2: draw latent variables from model-implied distribution
  if (nG == 1L) {
    # matrix(..., byrow = TRUE) rather than t(vapply()): with a single latent
    # variable vapply() returns a length-nsamp vector and t() would produce a
    # 1 x nsamp row matrix, corrupting the sample/variable orientation.
    eta_mat <- matrix(
      vapply(
        seq_len(nsamp),
        function(i) {
          sample_latent_from_model(samp$x_samp[i, ], lavmodel)
        },
        numeric(nlv)
      ),
      nrow = nsamp,
      ncol = nlv,
      byrow = TRUE
    )
    colnames(eta_mat) <- lv_names
  } else {
    # nocov start
    eta_list <- lapply(seq_len(nsamp), function(i) {
      sample_latent_from_model(samp$x_samp[i, ], lavmodel)
    })
    eta_mat <- do.call(
      rbind,
      lapply(eta_list, function(el) {
        unlist(el)
      })
    )
    colnames(eta_mat) <- paste0(
      rep(lv_names, nG),
      ".g",
      rep(seq_len(nG), each = nlv)
    )
  } # nocov end

  if (type == "latent") {
    return(eta_mat)
  }

  # Step 3: draw observed variables from model given eta
  if (nG == 1L) {
    # matrix(..., byrow = TRUE) rather than t(vapply()): guards the single
    # observed variable case the same way as the latent draws above.
    y_mat <- matrix(
      vapply(
        seq_len(nsamp),
        function(i) {
          sample_observed_from_model(samp$x_samp[i, ], eta_mat[i, ], lavmodel)
        },
        numeric(nobs)
      ),
      nrow = nsamp,
      ncol = nobs,
      byrow = TRUE
    )
    colnames(y_mat) <- ov_names
  } else {
    # nocov start
    y_mat <- t(vapply(
      seq_len(nsamp),
      function(i) {
        eta_per_group <- split(eta_mat[i, ], rep(seq_len(nG), each = nlv))
        eta_per_group <- lapply(eta_per_group, unname)
        unlist(sample_observed_from_model(
          samp$x_samp[i, ],
          eta_per_group,
          lavmodel
        ))
      },
      numeric(nobs * nG)
    ))
    colnames(y_mat) <- paste0(
      rep(ov_names, nG),
      ".g",
      rep(seq_len(nG), each = nobs)
    )
  } # nocov end

  # Without a mean structure, sample_observed_from_model() centres the
  # draws at zero (nu does not exist); add the saturated (sample) means of
  # the fitted data so replicates live on the data scale
  if (!isTRUE(lavmodel@meanstructure) && !isTRUE(prior) && nG == 1L) {
    y_mat <- sweep(y_mat, 2L, colMeans(int$lavdata@X[[1L]], na.rm = TRUE), "+")
    # under the marginalised likelihood the saturated means have posterior
    # N(ybar, Sigma/n); propagate that uncertainty into each replicate, as
    # estimated-nu draws do automatically when a mean structure exists
    if (marginalised_means_active(lavmodel)) {
      n_fit <- nrow(int$lavdata@X[[1L]])
      for (i in seq_len(nrow(y_mat))) {
        Sg <- compute_implied_moments(samp$x_samp[i, ], lavmodel)$cov
        ch <- tryCatch(chol(Sg), error = function(e) NULL) # nocov
        if (!is.null(ch)) {
          y_mat[i, ] <- y_mat[i, ] +
            as.numeric(crossprod(ch, rnorm(ncol(y_mat)))) / sqrt(n_fit)
        }
      }
    }
  }

  if (type == "observed") {
    return(y_mat)
  }

  # type == "all"
  list(
    lavaan = samp$x_samp,
    theta = samp$theta_samp,
    latent = eta_mat,
    observed = y_mat,
    implied = implied_list
  )
}