Brms & pk models?

So there are few things to really break the shackles brms imposes :-) Not sure it is better for your use case (I don’t think I understand your aims fully), but hope it is an inspiration.

The first trick is to “turn off” any default brms processing by specifying an empty family, i.e. a custom distribution that always return 0.

The second trick is to use a stanvar for the likelihood block to supply our own processing.

A simple example combining the two:

test_df <- data.frame(res = rlnorm(100), patient = sample(1:10, 100, replace = TRUE), x = rnorm(100))

empty_single_param <- function() {
  custom_family(
    "empty_single_param", dpars = c("mu"),
    links = c("identity"), 
    lb = 0, # We can put constraints that we expect for the data, but that's optional
    type = "real")
}

sv_empty_single_param <- stanvar(scode = "
      real empty_single_param_lpdf(real y, real mu) {
        return 0;
      }
    ", block = "functions")


sv_likelihood <- stanvar(scode = "// Arbitrary computation here", block = "likelihood", position = "end")

make_stancode(res ~ x + (1 | patient), family = empty_single_param, stanvars = sv_empty_single_param + sv_likelihood,
              data = test_df)

gives us:

// generated with brms 2.15.1
functions {
  
      real empty_single_param_lpdf(real y, real mu) {
        return 0;
      }
    
}
data {
  int<lower=1> N;  // total number of observations
  vector[N] Y;  // response variable
  int<lower=1> K;  // number of population-level effects
  matrix[N, K] X;  // population-level design matrix
  // data for group-level effects of ID 1
  int<lower=1> N_1;  // number of grouping levels
  int<lower=1> M_1;  // number of coefficients per level
  int<lower=1> J_1[N];  // grouping indicator per observation
  // group-level predictor values
  vector[N] Z_1_1;
  int prior_only;  // should the likelihood be ignored?
}
transformed data {
  int Kc = K - 1;
  matrix[N, Kc] Xc;  // centered version of X without an intercept
  vector[Kc] means_X;  // column means of X before centering
  for (i in 2:K) {
    means_X[i - 1] = mean(X[, i]);
    Xc[, i - 1] = X[, i] - means_X[i - 1];
  }
}
parameters {
  vector[Kc] b;  // population-level effects
  real Intercept;  // temporary intercept for centered predictors
  vector<lower=0>[M_1] sd_1;  // group-level standard deviations
  vector[N_1] z_1[M_1];  // standardized group-level effects
}
transformed parameters {
  vector[N_1] r_1_1;  // actual group-level effects
  r_1_1 = (sd_1[1] * (z_1[1]));
}
model {
  // likelihood including constants
  if (!prior_only) {
    // initialize linear predictor term
    vector[N] mu = Intercept + Xc * b;
    for (n in 1:N) {
      // add more terms to the linear predictor
      mu[n] += r_1_1[J_1[n]] * Z_1_1[n];
    }
    for (n in 1:N) {
      target += empty_single_param_lpdf(Y[n] | mu[n]);
    }
    // Arbitrary computation here
  }
  // priors including constants
  target += student_t_lpdf(Intercept | 3, 0.8, 2.5);
  target += student_t_lpdf(sd_1 | 3, 0, 2.5)
    - 1 * student_t_lccdf(0 | 3, 0, 2.5);
  target += std_normal_lpdf(z_1[1]);
}
generated quantities {
  // actual population-level intercept
  real b_Intercept = Intercept - dot_product(means_X, b);
}

The important thing to note that at the point our “arbitrary computation” in the likelihood block is executed, the mu vector is already computed and in scope, so we can supply any code we want that goes from mu to incrementing the target, including any looping behaviour etc.

All of the rest is just working on top of this template to make brms prepare you the linear predictors you need in a format you can use for your custom code.

In your case you probably won’t be happy with just predicting a single parameter. There are many ways to have brms prepare you more predictors. One is to use a multi-parameter custom (and empty) family, with the limitation that one of the parameters must be called “mu” e.g.:

empty_multi_param <- function() {
  custom_family(
    "empty_multi_param", dpars = c("mu", "lke", "lV"), 
    links = c("identity", "identity", "identity"), 
    lb = 0, # We can put constraints that we expect for the data, but that's optional
    type = "real")
}

sv_empty_multi_param <- stanvar(scode = "
      real empty_multi_param_lpdf(real y, real mu, real lke, real lV) {
        return 0;
      }
    ", block = "functions")


make_stancode(bf(res ~ x + (1 | patient), lke ~ (1 | patient), lV ~ x), family = empty_multi_param, stanvars = sv_empty_multi_param + sv_likelihood,
              data = test_df)

Now we have three vectors (mu, lke, and lV) ready and computed at the point our custom likelihood is exectued. You can also use the vint or vreal addition terms to put additional data here.

Alternatively, you could use multivariate formula with fake data:

test_df_mv <- test_df 
test_df_mv$lke <- 0 # The number is irrelevant and will be ignored
test_df_mv$lV <- 0 # The number is irrelevant and will be ignored

make_stancode(bf(res ~ x + (1 | patient)) + bf(lke ~ (1 | patient)) + bf(lV ~ x) + set_rescor(FALSE), family = empty_single_param, stanvars = sv_empty_single_param + sv_likelihood,
              data = test_df_mv)

where we once again have three vectors prepared, but they are called mu_res, mu_lke and mu_lV.

Finally, one can use stanvar to also add additional data you would need. And note that at this point, we can even completely break the core brms abstraction that a single row in the data frame corresponds to a single observation. We can for example make each row in a dataframe correspond to a patient, add a dummy (ignored) response and the pass a n_patients x n_time matrix as additional data, e.g.:

patient_df <- data.frame(patient = 1:10, x = rnorm(10), .dummy = 0)

n_time <- 5
obs_matrix <- matrix(rnorm(n_time * nrow(patient_df)), nrow = nrow(patient_df), ncol = n_time)


sv_n_time <- stanvar(x = n_time, name = "n_time", scode = "int n_time;")
sv_obs_matrix <- stanvar(x = obs_matrix, name = "obs_matrix", scode = "matrix[N, n_time] obs_matrix;")

sv_likelihood_2 <- stanvar(scode = "
  for(n in 1:N) {
     target += normal_lpdf(obs_matrix[n,] | solve_my_ode(mu[n], lke[n], lV[n]), 1);
  }
                           ", block = "likelihood", position = "end")

brm(bf(.dummy ~ x + (1 | patient), lke ~ (1 | patient), lV ~ x), family = empty_multi_param, stanvars = sv_empty_multi_param + sv_n_time + sv_obs_matrix + sv_likelihood_2,
              data = patient_df)

In all of the above approaches, posterior_linpred will work just fine to get your predictions of the parameters, but you’ll need to write your own code for posterior_epred or posterior_predict.

Does that make sense?

2 Likes