attempt to impose arbitrary equality constraints on correlations in a glmmTMB fit

constrained_cor.R

## simulate example (NOT satisfying constraints, but large and close enough that we
##  probably won't hit singular fits etc. etc.)

library(glmmTMB)
set.seed(1001)
nx <- 3
N <- 1000
ng <- 20
dd <- replicate(nx, rnorm(N), simplify = FALSE) |> as.data.frame() |> setNames(paste0("x", 1:nx))
dd$g <- factor(sample(1:ng, replace = TRUE, size = N))
thetavec <- rnorm(10)
dd$y <- simulate_new( ~ 1 + us(1 + x1 + x2 + x3 | g),
             newdata = dd,
             newparams = list(beta = 0,
                              theta = thetavec))[[1]]

## unconstrained fit

fit0 <- glmmTMB(y ~ 1 + us(1 + x1 + x2 + x3 | g), data = dd)

## begin modular fitting

m0 <- glmmTMB(y ~ 1 + us(1 + x1 + x2 + x3 | g), data = dd, doFit = FALSE)
m1 <- fitTMB(m0, doOptim = FALSE)
## now we have an *unconstrained* objective function
m1$fn()

## take a model and a set of constraints, return some useful stuff
## (overly complicated?)
##' @param model a partially constructed glmmTMB model (as above - results of fitTMB(., doOptim = FALSE))
##' @param sdeqpos standard deviations constrained to be identical (could/should do this with `map` instead?)
##' @param coreqpos two-column matrix of correlation-matrix positions constrained to be identical (in lower triangle)
##' @param n dimension of correlation matrix (for convenience so we don't have to back-calculate it)
mkwrapper <- function(model, sdeqpos = c(1,3), coreqpos = matrix(c(2, 1, 3, 1, 3, 2), byrow = TRUE, ncol = 2), n = 4) {
    ## function to take an unconstrained `theta` parameter vec (SD and cor parameters) and return a full-length parameter vec
    tranfun <- function(p) {
        ## construct SD vector (fill in constrained values from first element of p)
        sdvec <- rep(NA_real_, n)
        sdvec[ sdeqpos] <- p[1]
        sdvec[-sdeqpos] <- p[2:(n-length(sdeqpos))]
        ## construct cor vector (parameters are in terms of *correlations*, not on the usual
        ##  (glmm)TMB scaled-Cholesky scale. First element is the value for all of the constrained positions;
        ##  remaining parameters are the correlations to fill into the unconstrained position, in lower-triangle/columnwise
        ##  order
        corpars <- tail(p, n*(n-1)/2 - nrow(coreqpos) + 1)
        cormat <- diag(n)
        cormat[coreqpos] <- NA
        coreqpos_fixed <- which( is.na(cormat[lower.tri(cormat)]))
        coreqpos_est <-   which(!is.na(cormat[lower.tri(cormat)]))
        cormat[lower.tri(cormat)][coreqpos_fixed] <- corpars[1]
        cormat[lower.tri(cormat)][coreqpos_est] <- corpars[-1]
        
        ## convert to scaled-Cholesky scale (fill in all-NA, e.g. if we pick cor parameters to give a non-pos-def matrix)
        corvec <- try(put_cor(cormat[lower.tri(cormat)], input_val = "vec"), silent = TRUE)
        if (inherits(corvec, "try-error")) corvec[] <- NA
        c(sdvec, put_cor(cormat[lower.tri(cormat)], input_val = "vec"))
    }
    
    ## machinery to convert the full unconstrained param vec -> constrained param vec
    npar <- length(model$par)
    nconstr <- length(sdeqpos)-1 + nrow(coreqpos) - 1
    nthetapar <- n*(n+1)/2
    get_par <- function(p) {
        c(head(p, npar-nthetapar), tranfun(tail(p, nthetapar - nconstr)))
    }
        
    ## back-transforming the parameters is ugly ... pick all-zero for starting parameters
    ##  (may break for some cases)
    newobj <- list(par = rep(0, length(model$par) - (length(sdeqpos)-1) - (nrow(coreqpos)-1)),
                   fn = function(p) model$fn(get_par(p)), gr = function(p) model$gr(get_par(p)))
    list(newobj = newobj, tranfun = tranfun)
}

## full model should have 1 (fixed) + 1 (dispersion) + 4 + (4*3)/2 = 12 parameters
## constrained model in this case should have 12 - (2-1) - (3-1) = 9 parameters
## covariance matrix should have 7 parameters
m1C <- mkwrapper(m1)
tr1 <- m1C$tranfun(c(1000, 1, 1, -0.2, rep(0.1, 3)))
get_cor(tr1[-(1:4)], return_val = "mat")
nob <- m1C$newobj
length(m1C$tranfun(nob$par))
nob$fn(nob$par)

## finally, fit the constrained model
with(m1C$newobj, nlminb(par, fn))
cc <- with(m1C$newobj, optim(par, fn, method = "BFGS"))
trpar <- m1C$tranfun(cc$par[-(1:2)])
## constraints are satisfied
trpar[1:4]
get_cor(trpar[-(1:4)], return_val = "mat")

## might be able to use finalizeTMB() at this point? Would have to construct a fake 'fit' object
##  with the transformed parameters ...