Brms custom skew_generalized_t family

So I’ve updated the posterior_epred_skew_generalized_t() function, and now I am able to use some of the convenience functions in brms, however the conditional effects plot looks quite bizarre and I’m not sure if it’s because of how I coded the posterior_epred_skew_generalized_t() function.

This is the new function:

posterior_epred_skew_generalized_t <- function(prep) {
  mu <- brms::get_dpar(prep, "mu")
  sigma <- brms::get_dpar(prep, "sigma")
  lambda <- brms::get_dpar(prep, "lambda")
  p <- brms::get_dpar(prep, "p")
  q <- brms::get_dpar(prep, "q")
  
  ndraws <- prep$ndraws
  ns <- nrow(mu)
  
  # Use apply() to generate the epred matrix
  epred <- apply(cbind(mu, sigma, lambda, p, q), 1, function(x) {
    rskew_generalized_t(ndraws, x[1], x[2], x[3], x[4], x[5])
  })
  
  return(epred)
}

And this is my model:

init_fun <- function() {
  list(mu = rnorm(1, 0, 1),
       sigma = runif(1, 0.1, 1),
       lambda = runif(1, -0.99, 0.99),
       p = runif(1, 1, 10),
       q = runif(1, 1, 10))
}

SPT_rain_model <- brm(bf(SPT ~ rain + (rain | tag)), 
                      data = data[complete.cases(data[,c("rain", "SPT")]),],
                      save_pars = save_pars(all = TRUE),
                      iter = 5000,
                      init = init_fun,
                      prior = c(
                        prior(student_t(3,580, 70), class = Intercept,),
                        prior(student_t(3,0, 20), class = b),
                        prior(student_t(3,0, 20), class = sd),
                        prior(cauchy(0, 25), class = sigma),
                        prior(uniform(-0.99, 0.99), class = lambda, lb = -0.99, ub = 0.99),
                        prior(exponential(.1), class = p, lb = 2),
                        prior(exponential(.1), class = q, lb = 2)
                        ),
                      family = skew_generalized_t, stanvars = stanvars,
                      refresh = 1,
                      backend = "cmdstanr",
                      threads = threading(2),
                      control = list(max_treedepth = 10, adapt_delta = .9))

And here is the conditional effects plot

I tried transposing the epred matrix, thinking maybe that was my issue, but things still look jagged.

posterior_epred_skew_generalized_t <- function(prep) {
  mu <- brms::get_dpar(prep, "mu")
  sigma <- brms::get_dpar(prep, "sigma")
  lambda <- brms::get_dpar(prep, "lambda")
  p <- brms::get_dpar(prep, "p")
  q <- brms::get_dpar(prep, "q")
  
  ndraws <- prep$ndraws
  ns <- nrow(mu)
  
  # Use apply() to generate the epred matrix
  epred <- t(apply(cbind(mu, sigma, lambda, p, q), 1, function(x) {
    rskew_generalized_t(ndraws, x[1], x[2], x[3], x[4], x[5])
  }))
  
  
  return(epred)
}

Does anyone have any advice as to how I can improve this?

Also, I am posting to this thread because these are the people who I know to be interested in this distribution and who might find these functions helpful. Let me know if it’s better that I create a new thread.