Testing Reduce Sum in CmdStan 2.24 using N-mixture model

test_code.stan

functions {
  real partial_sum(int[] dummy_k,
                   int start, int end,
                   int[] count,
                   real lambda,
                   real p) {
    vector[end - start + 1] lp;

    for (k in start:end) {
      if (k - 1 < max(count))
        lp[k - start + 1] = negative_infinity();
      else
        lp[k - start + 1] = poisson_lpmf(k - 1 | lambda)
                            + binomial_lpmf(count | k - 1, p);
    }
    return log_sum_exp(lp);
  }
}

data {
  int<lower = 0> M;       // number of sites
  int<lower = 0> J;       // number of replications
  int<lower = 0> K;       // upper limit of the abundance
  int<lower = 0> N[M];    // abundance for each site
  int<lower = 0> C[M, J]; // count for each site and replication
}

transformed data {
  int dummy_k[K + 1] = rep_array(0, K + 1);
}

parameters {
  real<lower = 0> lambda;       // mean abundance
  real<lower = 0, upper = 1> p; // detection probability
}

model {
  int grainsize = 50;

  lambda ~ normal(0, 20);
  for (m in 1:M)
    target += reduce_sum_static(partial_sum, dummy_k, grainsize, C[m], lambda, p);
}

test_reduce_sum.R

# The simplest possible N-mixture model from Section 6.3 of
# Applied Hierarchical Modeling in Ecology.

library(cmdstanr)
set_cmdstan_path("/usr/local/cmdstan")

# generate simulated data

lambda <- 2.5 # mean abundance
p <- 0.4      # detection probability
M <- 150      # Number of sites
J <- 2        # Number of replications

set.seed(123)
N <- rpois(M, lambda) # abundance for each site
C <- matrix(0, nrow = M, ncol = J) # count for each site
for (j in 1:J)
  C[, j] <-rbinom(M, N, p)

# not multithreaded 

stan_data <- list(M = M, J = J, K = 100, N = N, C = C)
model1 <- cmdstan_model("test_code.stan")
fit1 <- model1$sample(data = stan_data,
                      chains = 4,
                      parallel_chains = 4,
                      iter_warmup = 1000, iter_sampling = 1000,
                      refresh = 500)
fit1$summary()

# multithreaded

stan_data <- list(M = M, J = J, K = 100, N = N, C = C)
model2 <- cmdstan_model("test_code.stan",
                        cpp_options = list(stan_threads = TRUE))
fit2 <- model2$sample(data = stan_data,
                      chains = 4,
                      parallel_chains = 4,
                      threads_per_chain = 2,
                      iter_warmup = 1000, iter_sampling = 1000,
                      refresh = 500)
fit2$summary()