stan_nb_dens.R

library( rstan )
options(mc.cores=6)

n <- 1000
s <- round( 10^rnorm( n, 3, .5 ) )
fracs <- 10^ifelse( runif(n)<.7, rnorm( n, -3.3, .4 ), rnorm( n, -1.7, .2 ) )
k <- rpois( n, fracs*s )

model <- stan_model( model_code="
data {
  int<lower=1> n;
  int<lower=1> ncomp;
  array[ncomp] real nbsize;
  array[n] int k;
  array[n] int s;
  array[ncomp] real lambda;
  
}

parameters {
  simplex[ncomp] theta;
}

model {
  vector[ncomp] log_theta = log(theta);
  for( i in 1:n ) {
     vector[ncomp] lps = log_theta;
     for( j in 1:ncomp ) {
        lps[j] += neg_binomial_lpmf( k[i] | nbsize[j], nbsize[j] / ( lambda[j] * s[i] ) );
     }
     target += log_sum_exp(lps);
  }
}
")



stepsize <- .3
lambda <- 10^seq( -5, 0, by=stepsize )
alpha <- .5*stepsize*lambda   # scale parameter
fit <- sampling( model, data=list( n=n, ncomp=length(lambda), nbsize=1/alpha, k=k, s=s, lambda=lambda ), iter=500 )

# Original fractions:
hist( log10(fracs), 100, freq=FALSE )

# Histogram of log-counts
hist( log10(k/s + 1e-4 ), 100, freq=FALSE )

# Original fractions again:
hist( log10(fracs), 100, freq=FALSE )
# Reconstruction of fractions density from counts
lines( log10(lambda), summary(fit)$summary[1:length(lambda),"mean"] / stepsize )

xg <- seq( -6, 0, length.out=1000 )      
lines( xg,
       sapply( 1:length(lambda), function(i)
          dgamma( 10^xg, shape=lambda[i]/alpha[i], scale=alpha[i] )*(10^xg)*log(10) ) %*% 
        summary(fit)$summary[1:length(lambda),"mean"], col="red" )