Attempt to compute conditional SDs for coxme fits

coxme_condsd.R

library(coxme)
library(Matrix)

#' @param fit a fitted \code{coxme} model
#' @param data the data frame for which to evaluate conditional SDs
#' @param grpvar (character) grouping variable (should be able to
#' extract this from the model, but it's a nuisance)
#' @param rform random-effects term (ditto)
#' @param sd.only return only sqrt(diag(V))? (otherwise return full covariance matrix)
coxme.condsd <- function(fit,
                         data = eval(getCall(fit)$data),
                         grpvar,
                         rform,
                         sd.only = TRUE) {
    ## model.frame() to get a sanitized (NA-free) data set
    vf <- all.vars(fit$formula$fixed)
    ## pick apart random effects terms
    vr <- (fit$formula$random
        |> lapply(function(x) lapply(lme4::findbars(x), all.vars))
        |> unlist(recursive = TRUE)
    )
    mf <- model.frame(reformulate(c(vf, vr)), data)
    ## indicator matrix for grouping variable
    Fmat <- t(Matrix::fac2sparse(mf[[grpvar]]))
    ## Z matrix
    Xr <- model.matrix(rform, mf)
    ## was (special case): Z <- cbind(Fmat, Fmat*mf$ph.ecog)
    Z <- t(Matrix::KhatriRao(t(Xr), t(Fmat)))
    ## remove intercept from RHS of fixed formula
    fform <- update(fit$formula$fixed[-2], ~ 0 + .)
    X <- model.matrix(fform, data = mf)
    ## respect order of vars in cov matrix: RE folowed by FE coefs
    ZX <- cbind(Z,X)
    V <- as.matrix(fit$variance)
    V2 <- ZX %*% V %*% t(ZX)
    if (sd.only) sqrt(diag(V2)) else V2
}

## cc <- c(ranef(fit)$inst[,1], ranef(fit)$inst[,2], coef(fit))
## head(drop(ZX %*% cc))
## head(fit$linear.predictor)
fit <- coxme(Surv(time, status) ~ ph.ecog + age + (1 + ph.ecog|inst), lung)

## this gets conditional SD (including both fixed & random effect uncertainty,
##  but using a plug-in estimate of the RE *variances*) for each observation in
## the original data set
coxme.condsd(fit, grpvar = "inst", rform = ~1 + ph.ecog)

## could also make up newdata with only unique values of the grouping variable,
##  setting fixed effect predictors either to the mean for each group or to some
##  overall baseline ...