Making my stan code, which models sparse two-dimensional data, more efficient

[rubi]

Hi,

I have two-dimensional sparse data to which I fit a hierarchical regression model and was wondering how can I make my stan code more efficient.

Specifically, my data are response observations that come from J individuals treated under T different treatments. Hence it can be organized in a J by T matrix.

The thing is that this matrix is sparse, meaning that I don't necessarily have an observation for each individual j from each treatment t.

To complicate it slightly further, the J individuals belong to two separate groups (labeled "1" and "2") where for each group the regression factors are slightly different.

And to complicate it even slightly further, my data are a mixture of two normals

Specifically, if I denote my response matrix as y (with dimensions J by T):

So if y[j,t] exists, then:

if y[j] is from group "1":

y[j,t] ~ normal(nu[t] + alpha[t] + X[j]*beta, sigma[t]) with probability delta

and 

y[j,t] ~ normal(-1*nu[t] + alpha[t] + X[j]*beta, sigma[t]) with probability 1-delta

and if y[j] is from group "2":

y[j,t] ~ normal(nu[t], sigma[t]) with probability delta

and 

y[j,t] ~ normal(-1*nu[t], sigma[t]) with probability 1-delta

In words, the "1" y[j]'s are affected by the nu[t] and alpha[t] factors of treatment t and beta co-variates of subject j (specified by the design matrix X which has dimensions J subjects by C co-variates). The "2" y[j]'s are only affected by the nu[t] factor. 

Both observations are assumed to have the same variance estimated by sigma[t] of that treatment. The mixture is on the nu[t] factor.

I put a hierarchy on the treatment and therefore have the priors:

sigma[t] ~ inv_gamma(a,b) (where I use empirical a = b = 0.1)

nu[t] ~ lognormal(mu_nu, sigma_nu) (log_normal is due to the mixture)

alpha[t] ~ lognormal(mu_alpha, sigma_alpha)

and then:

mu_nu ~ lognormal(0, 10)

sigma_nu ~ inv_gamma(c, d) (where I use empirical c = d = 0.1)

mu_alpha ~ normal(0, 10)

sigma_alpha ~ inv_gamma(e, f) (where I use empirical e = f = 0.1)

Finally, I use delta = 0.5.

This is my stan code:

data {

  int<lower=1> T; //number of treatments

  int<lower=1> J; //number of subjects

  int<lower=0> C; //number of covariates

  matrix[J,T] y; //response observations

  matrix<lower=0,upper=2>[J,T] map; //this is an indicator map that corresponds to y where: 0 - no observation; 1 - "1" group; 2 - "2" group

  matrix[J,C] X; //covariates

  vector<lower=0>[T] n_1s_treatment; //the number of "1"'s in each treatment

  real<lower=0> a;

  real<lower=0> b;

  real<lower=0> c;

  real<lower=0> d;

  real<lower=0> e;

  real<lower=0> f;

  real<lower=-1,upper=1> delta;

  int<lower=0,upper=1> have_ones; //binary indicator of whether y has "1"'s at all (1/0)

}

parameters {

  real<lower=0> sigma_alpha;

  real mu_alpha;

  real<lower=0> sigma_nu;

  real<lower=0> mu_nu;

  vector[T] alpha;

  vector<lower=0>[T] nu;

  vector<lower=0>[T] sigma;

  vector[C] beta;

}

model {

  if (have_ones == 1) {

    sigma_alpha ~ inv_gamma(e, f);

    mu_alpha ~ normal(0, 10);

  }

  sigma_nu ~ inv_gamma(c, d);

  mu_nu ~ lognormal(0, 10);

  for (t in 1:T) {

    if (n_1s_treatment[t] > 0) {

      alpha[t] ~ normal(mu_alpha, sigma_alpha);

    }

    nu[t] ~ lognormal(mu_nu, sigma_nu);

    sigma[t] ~ inv_gamma(a, b);

    for (j in 1:J) {

      if (map[j,t] > 0) {

        if (map[j,t] == 1) {

          increment_log_prob(log_sum_exp(log(delta) + normal_log(y[j,t], nu[t] + alpha[t] + X[j]*beta, sigma[t]),log1m(delta) + normal_log(y[j,t], -1*nu[t] + alpha[t] + X[j]*beta, sigma[t])));

        }

        else if (map[j,t] == 2) {

          increment_log_prob(log_sum_exp(log(delta) + normal_log(y[j,t], nu[t], sigma[t]), log1m(delta) + normal_log(y[j,t], -1*nu[t], sigma[t])));

        }

      }

    }

  }

}

Obviously vectorizing it would make it more efficient but since my data are sparse, even if I split y to the corresponding two groups and hold an array of vectors y's for each group (a vector per treatment) it seems to me that I'll still need to loop over treatments and subjects 

since I don't see how sparse data can be represented efficiently. But I'm probably wrong and would greatly appreciate being proven so and corrected.

stan code (code.stan) and an example data in RData file (exampleData.RData) are attached.

Thanks a lot,

rubi

[Andrew Gelman]

Hi, instead of coding the matrix as y[j,t], can you do what we do in ARM and code the data as y[i] with index variables jj[i] and tt[i], then you can feed in the data y, jj, tt, and associate data point y[i] with parameters alpha[jj[i]] and beta[tt[i]] (or however you are naming these)?

A

Thanks a lot,

rubi

--
You received this message because you are subscribed to the Google Groups "Stan users mailing list" group.
To unsubscribe from this group and stop receiving emails from it, send an email to stan-users+...@googlegroups.com.
For more options, visit https://groups.google.com/d/optout.

[rubi]

Hi again,

I followed your advice and implemented my stan code similar to this ARM example: 

https://github.com/stan-dev/example-models/blob/master/ARM/Ch.13/earnings_latin_square.stan

But unfortunately there seems to be a problem which I cannot seem to figure out, as there's no convergence to the true values I simulated the data with and there's bad chain mixing. 

I tried to debug this by adding an implementation that loops over the observations rather than sampling vectorially and that seems to work fine. As you'll see below, these two implementations are very similar which makes me think there's something very basic that I'm doing wrong.

Here's the code (also attached):

data {

  int<lower=0> N_one; //total number of observations with label 1

  int<lower=0> N_two; //total number of observations with label 2

  int<lower=0> J_one; //number of subjects with label 1

  int<lower=0> J_two; //number of subjects with label 2

  int<lower=0> C_one; //number of covariates of the label 1 subjects

  int<lower=1> T; //number of treatments

  int<lower=0,upper=J_one> one_observation2subject[N_one]; //y_one_hat[i] -> X[one_observation2subject[i]]

  int<lower=0,upper=J_two> two_observation2subject[N_two];

  int<lower=0,upper=T> one_observation2treatment[N_one]; //y_one_hat[i] -> alpha_treatment[one_observation2treatment[i]], nu_treatment[one_observation2treatment[i]]

  int<lower=0,upper=T> two_observation2treatment[N_two]; //y_two_hat[i] -> nu_treatment[two_observation2treatment[i]]

  vector[N_one] y_one_hat; //one_subject.1.one_treatment.1,one_subject.1.one_treatment.2,...,one_subject.j.one_treatment.t,...,one_subject.J.one_treatment.T

  vector[N_two] y_two_hat; //two_subject.1.two_treatment.1,two_subject.1.two_treatment.2,...,two_subject.j.two_treatment.t,...,two_subject.J.two_treatment.T

  vector[N_one] e_one_hat; //these are measurement errors of the response of label 1

  vector[N_two] e_two_hat; //these are measurement errors of the response of label 2

  matrix[J_one,C_one] X_one; //only subjects of label 1 have covariates

  real<lower=0> a;

  real<lower=0> b;

  real<lower=0> c;

  real<lower=0> d;

  real<lower=0> e;

  real<lower=0> f;

  real<lower=0,upper=1> delta;

}

parameters {

  real<lower=0> sigma_alpha;

  real mu_alpha;

  real<lower=0> sigma_nu;

  real<lower=0> log_mu_nu;

  vector[T] alpha_treatment;

  vector[T] log_nu_treatment;

  vector<lower=0>[T] sigma_treatment;

  vector[C_one] beta_one;

}

transformed parameters {

  vector[N_one] one_pos_mean;

  vector[N_one] one_neg_mean;

  vector[N_one] one_sd;

  vector[N_two] two_pos_mean;

  vector[N_two] two_neg_mean;

  vector[N_two] two_sd;

  real term_one;

  vector<lower=0>[T] nu_treatment;

  nu_treatment <-  exp(log_nu_treatment);

  for(i in 1:N_one)  {

    term_one <- alpha_treatment[one_observation2treatment[i]] + X_one[one_observation2subject[i]]*beta_one;

    one_pos_mean[i] <-  nu_treatment[one_observation2treatment[i]] + term_one;

    one_neg_mean[i] <-  -nu_treatment[one_observation2treatment[i]] + term_one;

    one_sd[i] <-  sigma_treatment[one_observation2treatment[i]] + e_one_hat[i];

  }

  for(i in 1:N_two)  {

    two_pos_mean[i] <-  nu_treatment[two_observation2treatment[i]];

    two_neg_mean[i] <-  -nu_treatment[two_observation2treatment[i]];

    two_sd[i] <-  sigma_treatment[two_observation2treatment[i]] + e_two_hat[i];

  }

}

model {

  real log_prob_terms[4];

  sigma_alpha ~ inv_gamma(c, d);

  mu_alpha ~ normal(0, 10);

  sigma_nu ~ inv_gamma(e, f);

  log_mu_nu ~ lognormal(0, 10);

  alpha_treatment ~ normal(mu_alpha, sigma_alpha);

  log_nu_treatment ~ normal(log_mu_nu, sigma_nu);

  sigma_treatment ~ inv_gamma(a, b);

  /////vectorial implementation/////

  //  log_prob_terms[1] <- log(delta) + normal_log(y_one_hat, one_pos_mean, one_sd);

  //  log_prob_terms[2] <- log1m(delta) + normal_log(y_one_hat, one_neg_mean, one_sd);

  //  log_prob_terms[3] <- log(delta) + normal_log(y_two_hat, two_pos_mean, two_sd);

  //  log_prob_terms[4] <- log1m(delta) + normal_log(y_two_hat, two_neg_mean, two_sd);

  //  increment_log_prob(log_sum_exp(log_prob_terms));

  /////vectorial implementation/////

  /////sequential implementation/////

  for(i in 1:N_one)  {

      increment_log_prob(log_sum_exp(log(delta) + normal_log(y_one_hat[i], nu_treatment[one_observation2treatment[i]] + alpha_treatment[one_observation2treatment[i]] + X_one[one_observation2subject[i]]*beta_one, sigma_treatment[one_observation2treatment[i]] + e_one_hat[i]),log1m(delta) + normal_log(y_one_hat[i], -nu_treatment[one_observation2treatment[i]] + alpha_treatment[one_observation2treatment[i]] + X_one[one_observation2subject[i]]*beta_one, sigma_treatment[one_observation2treatment[i]] + e_one_hat[i])));

    }

  for(i in 1:N_two)  {

    increment_log_prob(log_sum_exp(log(delta) + normal_log(y_two_hat[i], nu_treatment[two_observation2treatment[i]], sigma_treatment[two_observation2treatment[i]] + e_two_hat[i]),log1m(delta) + normal_log(y_two_hat[i], -nu_treatment[two_observation2treatment[i]], sigma_treatment[two_observation2treatment[i]] + e_two_hat[i])));

  }

  /////sequential implementation/////

}

Attached are also an example data I ran my code with (4 chains, 1000 iterations, 100 iterations of warm up) and trace plots of running the vectorial implementation (vectorial.trace.pdf) and sequential implementation (sequential.trace.pdf), where the true simulated values are shown as horizontal black bars.

Thanks a lot,

rubi

[Bob Carpenter]

You have to do mixtures at the individual level --- there's
no good way to vectorize the computation.

Your vector-based code has a different meaning if you
write it out. Consider a simple case of a two-normal
mixture.

Individual level:

PROD_n [ lambda * normal(y[n] | mu1, sigma1)
+ (1 - lambda) * normal(y[n] | mu2, sigma2) ]

Model level:

[ lambda * PROD_n normal(y[n] | mu1, sigma1) ]
+ [ (1 - lambda) * PROD_n normal(y[n] | mu2, sigma2) ]

Two different densities. The latter mixes whole models whereas
the former mixes at the indivdual level

is trying to mix entire models,
not draw observations from a mixture.

- Bob

P.S. You don't need to CC Andrew --- he's on the Stan users list.

> <example.data.RData><sequential.trace.pdf><vectorial.trace.pdf><code.stan>

[rubi]

Thanks a lot.