use analytic integration to speed up n-mixture estimation

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)]

vignettes/nmixtures.Rmd

@@ -368,7 +368,7 @@ trend_map %>%
   head(12)
 ```
 
-Now we are ready to fit a model using `mvgam()`. Here we will use penalized splines for each of the continuous covariate effects to detect possible nonlinear associations. We also showcase how `mvgam` can make use of the different approximation algorithms available in `Stan` by using the meanfield variational Bayes approximator (this reduces computation time to around 12 seconds for this example)
+Now we are ready to fit a model using `mvgam()`. Here we will use penalized splines for each of the continuous covariate effects to detect possible nonlinear associations. We also showcase how `mvgam` can make use of the different approximation algorithms available in `Stan` by using the meanfield variational Bayes approximator (this reduces computation time from around 90 seconds to around 12 seconds for this example)
 ```{r include = FALSE, results='hide'}
 mod <- mvgam(
   # effects of covariates on detection probability;