Following up on several discussions of simplex adjustments and ragged arrays of simplexes

The inv ilr is at transforms/simplex_transforms/stan/transforms/ILR/ILR_functions.stan at main · bob-carpenter/transforms · GitHub.

  1. Do things on the log-scale. Also swap in the inv_ILR simplex transform if you want. Note that the inv_ILR transform has one-less parameter than the softmax.
  matrix log_matrix_vector_product(matrix log_A, vector log_x) {
    int N = rows(log_A);
    vector[N] result;

    for (i in 1:n) {
        vector[cols(A)] sum_vector = log_A[i] + log_x;
        result[i] = log_sum_exp(sum_vector);
      }
    
    return result;
  }

  vector simplex_log_constrain_lp(vector y) {
    int N = rows(y);
    real r = log_sum_exp(y);
    vector[N] log_x = y - r;
  
    target += std_normal_lpdf(r - log(N));
    target += log_x[N];
    
    return log_x;
  }
real log_dirichlet_lpdf(vector log_theta, vector alpha) {
  int N = rows(log_theta);
  if (N != rows(alpha))
    reject("Input must contain same number of elements as alpha");
  return dot_product(alpha, log_theta) - log_theta[N]
         + lgamma(sum(alpha)) - sum(lgamma(alpha));
}
 real multinomial_log_lpmf(array int y, vector log_theta) {
    int N = sum(y);
    int K = num_elements(log_theta);
    real lp = lgamma(N + 1);

    for (k in 1:K)
      lp += log_theta[k] * y[k] - lgamma(y[k] + 1);
    
  return lp;
  }
}

data {
  int<lower=1> n;           // number of election units i
  int<lower=2> R;           // number of candidates in election 1
  int<lower=2> C;           // number of candidates in election 2
  array[n] vector<upper=0>[R] log_p_r;       // marginal log proportion for row r in unit i
  array[n, C] int<lower=0> N_c; // marginal count for column c in unit i
  real<lower=0> lambda_1;    // gamma hyperprior shape parameter
  real<lower=0> lambda_2;    // gamma hyperprior scale parameter
}
parameters {
  matrix<lower=0>[R, C] alpha_r;    

  array[n, R] vector<lower=0>[C] gamma_r;
  array[n] vector<lower=0>[C] gamma_c;
}
transformed parameters {
  array[n] vector[C] log_theta_c;
  array[n] matrix[R, C] log_beta_mat;

  for (i in 1:n) {
   // you overwrite theta below
   // theta_c[i] = simplex_log_constrain_lp(gamma_c[i]);
    for (r in 1:R) {
       beta_mat[i, r] = simplex_log_constrain_lp(gamma_r[i, r])';
    }
   log_theta_c[i] = log_matrix_vector_product(beta_mat[i]', log_p_r[i]);
  }
}
model {
  to_vector(alpha_r) ~ gamma(lambda_1, lambda_2);
  for (i in 1:n) {
    for (r in 1:R) {
      log_beta_mat[i, r] ~ log_dirichlet(alpha_r[r,]);
    }
    N_c[i] ~ multinomial_log(log_theta_c[i]);
  }
}
  1. Try the dirichlet_multinomial but since you have p_r it’s a weighted dirichlet multinomial:

We want the marginal distribution after integrating out ( Z ):

P(x \mid n, \alpha, p) = \int \text{Mult}(x \mid n, Z \cdot p) \cdot \prod_{r=1}^R \text{Dir}(Z_{:,r} \mid \alpha_r) \, dZ

I believe this integration results in the following weighted Dirichlet-multinomial distribution:

P(x \mid n, \alpha, p) = \frac{n!}{x_1! \cdots x_C!} \prod_{c=1}^C \left( \frac{\prod_{r=1}^R \Gamma(x_c + \alpha_{c,r})}{\Gamma\left(\sum_{r=1}^R (x_c + \alpha_{c,r})\right)} \cdot \frac{\Gamma\left(\sum_{r=1}^R \alpha_{c,r} \right)}{\prod_{r=1}^R \Gamma(\alpha_{c,r})} \cdot p_r^{x_c} \right)

That’s enough for me now, you can get the matrix beta back out in the generated quantities.

functions {
 real log_multivariate_beta (vector alpha) {
   int N = num_elements(alpha);
   return sum(lgamma(alpha)) - N * lgamma(sum(alpha));
}

real weighted_dirichlet_multinomial_lpmf (array int y, vector alpha, vector p_r) {
   int N = sum(y);
   int K = num_elements(y);
   real lp = lgamma(N + 1);
  
   for (k in 1:K) {
     lp += y[k] * log(p_r[k]) + log_multivariate_beta(y + alpha) - log_multivariate_beta(alpha) - lgamma(y[k] + 1);
   }
   return lp;
 }
}
data {
  int<lower=1> n;           // number of election units i
  int<lower=2> R;           // number of candidates in election 1
  int<lower=2> C;           // number of candidates in election 2
  array[n] simplex[R] p_r;       // marginal proportion for row r in unit i
  array[n, C] int<lower=0> N_c; // marginal count for column c in unit i
  real<lower=0> lambda_1;    // gamma hyperprior shape parameter
  real<lower=0> lambda_2;    // gamma hyperprior scale parameter
}
parameters {
  matrix<lower=0>[R, C] alpha_r;    
}
model {
  to_vector(alpha_r) ~ gamma(lambda_1, lambda_2);

  for (i in 1:n) 
    N_c[i] ~ weighted_dirichlet_multinomial(alpha_r[i], p_r[i]);
}
1 Like