Declare parameter constraints for (non-linear) models

[StaffanBetner]

I have an example (or it's rather something @jackobailey asked about, and I am trying to help out) where it would be useful to be able to declare a parameter vector as a simplex, and put a Dirichlet prior on it. I am sure other constraints would be useful for other people, e.g. ordered parameters.

The idea is to model the propensity to vote for a party with a weighted average of the voting for the same party previous years (in my case on a municipality level, in Jack's case on an individual level).

The formula is

bf(count_2018|trials(pop_2018) ~ a + b*c)+
      lf(a ~ 1)+
      lf(b ~ 1)+
      lf(c ~ 0 + prop_02 + prop_06 + prop_10 + prop_14)+
      set_nl(TRUE)+
      binomial()

and the priors:

prior = set_prior("normal(0, 1.5)", nlpar = "a")+
      set_prior("normal(0, 1.5)", nlpar = "b")+
      set_prior("dirichlet(rep_vector( 0.5, 4 ))", nlpar = "c", class="b")

which generates the following Stan code, with a comment added by me:

// generated with brms 2.15.1
functions {
}
data {
  int<lower=1> N;  // total number of observations
  int Y[N];  // response variable
  int trials[N];  // number of trials
  int<lower=1> K_a;  // number of population-level effects
  matrix[N, K_a] X_a;  // population-level design matrix
  int<lower=1> K_b;  // number of population-level effects
  matrix[N, K_b] X_b;  // population-level design matrix
  int<lower=1> K_c;  // number of population-level effects
  matrix[N, K_c] X_c;  // population-level design matrix
  int prior_only;  // should the likelihood be ignored?
}
transformed data {
}
parameters {
  vector[K_a] b_a;  // population-level effects
  vector[K_b] b_b;  // population-level effects
  vector[K_c] b_c;  // population-level effects // Should be a simplex!!
}
transformed parameters {
}
model {
  // likelihood including constants
  if (!prior_only) {
    // initialize linear predictor term
    vector[N] nlp_a = X_a * b_a;
    // initialize linear predictor term
    vector[N] nlp_b = X_b * b_b;
    // initialize linear predictor term
    vector[N] nlp_c = X_c * b_c;
    // initialize non-linear predictor term
    vector[N] mu;
    for (n in 1:N) {
      // compute non-linear predictor values
      mu[n] = nlp_a[n] + nlp_b[n] * nlp_c[n];
    }
    target += binomial_logit_lpmf(Y | trials, mu);
  }
  // priors including constants
  target += normal_lpdf(b_a | 0, 1.5);
  target += normal_lpdf(b_b | 0, 1.5);
  target += dirichlet_lpdf(b_c | rep_vector( 0.5, 4 ));
}
generated quantities {
}

where vector[K_c] b_c; would in our case be simplex[K_c] b_c;.

One approach might be a setting to override the parameter code block in the stanvar function, or a new argument constraints in brm, working similar to prior, with functions like set_constraint and a reference to the parameters similar to the prior specification.

[jackobailey]

Thanks for looking into this Staffan. Another (maybe simpler) idea might simply be to have brms assert that any parameters modelled using a dirichlet prior are automatically considered to be a simplex and not a vector.

[jackobailey]

Just flagging for anyone who's interested that @StaffanBetner has worked out a way to fit the model and shoehorn it into a brms object in the meantime

https://discourse.mc-stan.org/t/dirichlet-prior-possible-with-stanvars/15812/5?u=jackbailey

[paul-buerkner]

This is now implemented. For an example, see below:

dat <- data.frame(y = rnorm(10), x1 = rnorm(10), x2 = rnorm(10))
make_stancode(y~x1+x2, dat, prior = prior("dirichlet([1,2]')", class = "b"))