Extracting residual covariance structure in brms

Modeling
1

Hi everyone,

I’m working with the Ovary dataset from the nlme package and fitting models using brms as follows:

brm_mlm <- brm(follicles ~ sin_time + cos_time+(sin_time||Mare),
               data=ovary_df,cores=4)

brm_ar <- brm(follicles ~ sin_time + cos_time+(sin_time||Mare)+ar(gr=Mare,cov=TRUE),
               data=ovary_df,cores=4)

These are inspired by the models below from Mixed-Effects Models in S and S-PLUS (Pinheiro & Bates, 2000):

fm1Ovar.lme <- lme( follicles ~ sin(2*pi*Time) + cos(2*pi*Time), 
data = Ovary, random = pdDiag(~sin(2*pi*Time)) )
fm2Ovar.lme <- update( fm1Ovar.lme, correlation = corAR1() )
fm5Ovar.lme <- update(fm1Ovar.lme, corr = corARMA(p = 1, q = 1))

In the book, they examine the autocorrelation function (ACF) of the standardized residuals from the first model, observe significant lag-1 autocorrelation, and proceed to fit an AR(1) model (and later an ARMA(1,1)) to address it.

They then plot the ACF of normalized residuals (standardized residuals pre-multiplied by the inverse square root of the estimated residual covariance matrix):

plot( ACF(fm5Ovar.lme, maxLag = 10, resType = "n"),
+ alpha = 0.01 )

to obtain normalized residuals behaving like white noise.

I’m trying to replicate this with my brm_ar model. Here is its summary output:

 brm_ar
 Family: gaussian 
  Links: mu = identity; sigma = identity 
Formula: follicles ~ sin_time + cos_time + (sin_time || Mare) + ar(gr = Mare, cov = TRUE) 
   Data: ovary_df (Number of observations: 308) 
  Draws: 4 chains, each with iter = 2000; warmup = 1000; thin = 1;
         total post-warmup draws = 4000

Correlation Structures:
      Estimate Est.Error l-95% CI u-95% CI Rhat Bulk_ESS Tail_ESS
ar[1]     0.58      0.07     0.46     0.72 1.00     1353     1868

Multilevel Hyperparameters:
~Mare (Number of levels: 11) 
              Estimate Est.Error l-95% CI u-95% CI Rhat Bulk_ESS Tail_ESS
sd(Intercept)     3.14      0.93     1.75     5.24 1.00      994     1976
sd(sin_time)      2.63      1.05     0.67     5.00 1.00     1010     1230

Regression Coefficients:
          Estimate Est.Error l-95% CI u-95% CI Rhat Bulk_ESS Tail_ESS
Intercept    12.25      1.04    10.13    14.28 1.00     1024     1522
sin_time     -0.60      0.55    -1.64     0.49 1.00     2905     2934
cos_time     -1.23      0.58    -2.37    -0.10 1.00     2351     2581

Further Distributional Parameters:
      Estimate Est.Error l-95% CI u-95% CI Rhat Bulk_ESS Tail_ESS
sigma     2.92      0.13     2.69     3.18 1.00     2293     2336

Draws were sampled using sampling(NUTS). For each parameter, Bulk_ESS
and Tail_ESS are effective sample size measures, and Rhat is the potential
scale reduction factor on split chains (at convergence, Rhat = 1).

For the moment, I have tried to manually construct the covariance matrix using the following functions:


make_ar_matrix <- function(n,rho,sigma_ar_2){
  M <- matrix(0,n,n)
  for (i in 1:n){
    for(j in 1:n){
      M[i,j] = sigma_ar_2*rho^(abs(i-j))
    }
  }
  return(M)
}

inverse_sqrt_matrix <- function(A) {
  eig <- eigen(A)
  V <- eig$vectors
  D <- eig$values
  
  # Check positive definiteness
  if (any(D <= 0)) stop("Matrix is not positive definite")
  
  D_inv_sqrt <- diag(1 / sqrt(D))
  A_inv_sqrt <- V %*% D_inv_sqrt %*% t(V)
  
  return(A_inv_sqrt)
}

normalize_resid_ar <- function(df,rho_ar,sigma_ar_2){
  count_mare_obs <- df %>% count(Mare) %>% pull(n)
  ar_mat_mare <- lapply(count_mare_obs , function(x) make_ar_matrix(x,rho_ar,sigma_ar_2))
  ar_inv_sqrt <- lapply(ar_mat_mare, inverse_sqrt_matrix)
  group_split_df <- df %>% group_split(Mare)
  normalized_res <- lapply(1:length(count_mare_obs ), function(i){
    ar_inv_sqrt[[i]]%*%group_split_df[[i]]$resid
  })
  df <- df %>% mutate(resid_nrm=do.call(rbind,normalized_res))
}

Assuming the residuals follow an AR(1) structure, we expect:

\text{Cov}(\epsilon_{ij},\epsilon{ik}) = \sigma^2\rho^{|j-k|}
Where \rho is ar[1] and \sigma^2 is the residual variance (sigma^2 = 2.92^2 from the summary above).

My question:
Is there a built-in or recommended way to extract:

  1. The residual covariance matrix directly from a brms model with ar(...) terms, or
  2. The normalized residuals (i.e., transformed standardized residuals that account for the autocorrelation structure)?

Or must I always manually compute them for both AR(1) and ARMA(1,1) structures as shown above?

Thanks in advance!