diff --git a/R/add_nmixture.R b/R/add_nmixture.R
index dd991e26..3b905cef 100644
--- a/R/add_nmixture.R
+++ b/R/add_nmixture.R
@@ -351,25 +351,32 @@ add_nmix_model = function(model_file,
model_file[grep('flat_ys ~ poisson_log_glm(append_col(flat_xs, flat_trends),',
model_file, fixed = TRUE)] <-
paste0('// loop over replicate sampling window (each site*time*species combination)\n',
- 'for ( k in 1 : K_groups ) {\n',
+ 'for (k in 1 : K_groups) {\n',
'// all log_lambdas are identical because they represent site*time\n',
'// covariates; so just use the first measurement\n',
'real log_lambda = flat_trends[K_inds[k, 1]];\n',
- 'vector[N_max[k] - Y_max[k] + 1] terms;\n',
- 'int l = 0;\n',
- '// marginalize over latent abundance\n',
- 'for ( Ni in Y_max[k] : N_max[k] ) {\n',
- 'l = l + 1;\n',
- '// factor for poisson prob of latent Ni; compute\n',
- '// only once per sampling window\n',
- 'terms[l] = poisson_log_lpmf( Ni | log_lambda ) +\n',
- '// for each replicate observation, binomial prob observed is\n',
- '// computed in a vectorized statement\n',
- 'binomial_logit_lpmf( flat_ys[K_inds[k, K_starts[k] : K_stops[k]]] |\n',
- 'Ni,\n',
- 'flat_ps[K_inds[k, K_starts[k] : K_stops[k]]] );\n',
+ '// logit-scale detection probilities for the replicate observations\n',
+ 'vector[size(K_inds[k, K_starts[k] : K_stops[k]])] logit_p = to_vector(flat_ps[K_inds[k, K_starts[k] : K_stops[k]]]);\n',
+ '// K values and observed counts for these replicates\n',
+ 'int K_max = N_max[k];\n',
+ 'int K_min = Y_max[k];\n',
+ 'array[size(K_inds[k, K_starts[k] : K_stops[k]])] int N_obs = flat_ys[K_inds[k, K_starts[k] : K_stops[k]]];\n',
+ 'int possible_N = K_max - K_min;\n',
+ '// marginalize over possible latent counts analytically\n',
+ 'real ff = exp(log_lambda) * prod(1 - inv_logit(logit_p));\n',
+ 'real prob_n = 1;\n',
+ 'for (i in 1 : possible_N){\n',
+ 'real N = K_max - i + 1;\n',
+ 'real k_obs = 1;\n',
+ 'for (j in 1 : size(N_obs)){\n',
+ 'k_obs *= N / (N - N_obs[j]);\n',
'}\n',
- 'target += log_sum_exp( terms );\n',
+ 'prob_n = 1 + prob_n * ff * k_obs / N;\n',
+ '}\n',
+ '// add log(pr_n) to prob(K_min)\n',
+ 'target += poisson_log_lpmf(K_min | log_lambda) +\n',
+ 'binomial_logit_lpmf(N_obs | K_min, logit_p) +\n',
+ 'log(prob_n);\n',
'}')
model_file <- model_file[-grep('0.0,append_row(b, 1.0));',
model_file, fixed = TRUE)]
diff --git a/docs/articles/nmixtures.html b/docs/articles/nmixtures.html
index e61aafaf..7abd6da0 100644
--- a/docs/articles/nmixtures.html
+++ b/docs/articles/nmixtures.html
@@ -13,6 +13,7 @@
+