SBC was primarily designed for continuous parameters, but can be used with models that have discrete parameters - whether the parameters are directly represented (e.g. in BUGS/JAGS) or marginalized out (as is usual in Stan).

Stan version and debugging

library(SBC); 
library(ggplot2)

use_cmdstanr <- getOption("SBC.vignettes_cmdstanr", TRUE) # Set to false to use rstan instead

if(use_cmdstanr) {
  library(cmdstanr)
} else {
  library(rstan)
  rstan_options(auto_write = TRUE)
}

# Multiprocessing support
library(future)
plan(multisession)

# The fits are very fast and we fit just a few, 
# so we force a minimum chunk size to reduce overhead of
# paralellization and decrease computation time.
options(SBC.min_chunk_size = 5)

# Setup caching of results
if(use_cmdstanr) {
  cache_dir <- "./_discrete_vars_SBC_cache"
} else {
  cache_dir <- "./_discrete_vars_rstan_SBC_cache"
}
cache_dir_jags <- "./_discrete_vars_SBC_cache"

if(!dir.exists(cache_dir)) {
  dir.create(cache_dir)
}
if(!dir.exists(cache_dir_jags)) {
  dir.create(cache_dir_jags)
}

We take the changepoint model from: https://mc-stan.org/docs/2_26/stan-users-guide/change-point-section.html

cat(readLines("stan/discrete_vars1.stan"), sep = "\n")
data {
  real<lower=0> r_e;
  real<lower=0> r_l;

  int<lower=1> T;
  array[T] int<lower=0> y;
}
transformed data {
  real log_unif;
  log_unif = -log(T);
}
parameters {
  real<lower=0> e;
  real<lower=0> l;
}
transformed parameters {
  vector[T] lp;
  lp = rep_vector(log_unif, T);
  for (s in 1:T)
    for (t in 1:T)
      lp[s] = lp[s] + poisson_lpmf(y[t] | t < s ? e : l);
}
model {
  e ~ exponential(r_e);
  l ~ exponential(r_l);
  target += log_sum_exp(lp);
}

generated quantities {
  int<lower=1,upper=T> s;
  s = categorical_logit_rng(lp);
}
if(use_cmdstanr) {
  model_1 <- cmdstan_model("stan/discrete_vars1.stan")
  backend_1 <- SBC_backend_cmdstan_sample(model_1)
} else {
  model_1 <- stan_model("stan/discrete_vars1.stan")
  backend_1 <- SBC_backend_rstan_sample(model_1)
}
## In file included from stan/lib/stan_math/lib/boost_1.81.0/boost/multi_array/multi_array_ref.hpp:32,
##                  from stan/lib/stan_math/lib/boost_1.81.0/boost/multi_array.hpp:34,
##                  from stan/lib/stan_math/lib/boost_1.81.0/boost/numeric/odeint/algebra/multi_array_algebra.hpp:22,
##                  from stan/lib/stan_math/lib/boost_1.81.0/boost/numeric/odeint.hpp:63,
##                  from stan/lib/stan_math/stan/math/prim/functor/ode_rk45.hpp:9,
##                  from stan/lib/stan_math/stan/math/prim/functor/integrate_ode_rk45.hpp:6,
##                  from stan/lib/stan_math/stan/math/prim/functor.hpp:16,
##                  from stan/lib/stan_math/stan/math/rev/fun.hpp:200,
##                  from stan/lib/stan_math/stan/math/rev.hpp:12,
##                  from stan/lib/stan_math/stan/math.hpp:19,
##                  from stan/src/stan/model/model_header.hpp:4,
##                  from C:/Users/Martin/AppData/Local/Temp/RtmpA34pwG/model-2b747fd82a5f.hpp:2:
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:180:45: warning: 'template<class _Arg, class _Result> struct std::unary_function' is deprecated [-Wdeprecated-declarations]
##   180 |         : public boost::functional::detail::unary_function<typename unary_traits<Predicate>::argument_type,bool>
##       |                                             ^~~~~~~~~~~~~~
## In file included from C:/rtools43/ucrt64/include/c++/13.2.0/string:49,
##                  from C:/rtools43/ucrt64/include/c++/13.2.0/bits/locale_classes.h:40,
##                  from C:/rtools43/ucrt64/include/c++/13.2.0/bits/ios_base.h:41,
##                  from C:/rtools43/ucrt64/include/c++/13.2.0/ios:44,
##                  from C:/rtools43/ucrt64/include/c++/13.2.0/istream:40,
##                  from C:/rtools43/ucrt64/include/c++/13.2.0/sstream:40,
##                  from C:/rtools43/ucrt64/include/c++/13.2.0/complex:45,
##                  from stan/lib/stan_math/lib/eigen_3.4.0/Eigen/Core:50,
##                  from stan/lib/stan_math/lib/eigen_3.4.0/Eigen/Dense:1,
##                  from stan/lib/stan_math/stan/math/prim/fun/Eigen.hpp:22,
##                  from stan/lib/stan_math/stan/math/rev.hpp:4:
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:117:12: note: declared here
##   117 |     struct unary_function
##       |            ^~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:214:45: warning: 'template<class _Arg1, class _Arg2, class _Result> struct std::binary_function' is deprecated [-Wdeprecated-declarations]
##   214 |         : public boost::functional::detail::binary_function<
##       |                                             ^~~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:131:12: note: declared here
##   131 |     struct binary_function
##       |            ^~~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:252:45: warning: 'template<class _Arg, class _Result> struct std::unary_function' is deprecated [-Wdeprecated-declarations]
##   252 |         : public boost::functional::detail::unary_function<
##       |                                             ^~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:117:12: note: declared here
##   117 |     struct unary_function
##       |            ^~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:299:45: warning: 'template<class _Arg, class _Result> struct std::unary_function' is deprecated [-Wdeprecated-declarations]
##   299 |         : public boost::functional::detail::unary_function<
##       |                                             ^~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:117:12: note: declared here
##   117 |     struct unary_function
##       |            ^~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:345:57: warning: 'template<class _Arg, class _Result> struct std::unary_function' is deprecated [-Wdeprecated-declarations]
##   345 |     class mem_fun_t : public boost::functional::detail::unary_function<T*, S>
##       |                                                         ^~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:117:12: note: declared here
##   117 |     struct unary_function
##       |            ^~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:361:58: warnin
## g: 'template<class _Arg1, class _Arg2, class _Result> struct std::binary_function' is deprecated [-Wdeprecated-declarations]
##   361 |     class mem_fun1_t : public boost::functional::detail::binary_function<T*, A, S>
##       |                                                          ^~~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:131:12: note: declared here
##   131 |     struct binary_function
##       |            ^~~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:377:63: warning: 'template<class _Arg, class _Result> struct std::unary_function' is deprecated [-Wdeprecated-declarations]
##   377 |     class const_mem_fun_t : public boost::functional::detail::unary_function<const T*, S>
##       |                                                               ^~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:117:12: note: declared here
##   117 |     struct unary_function
##       |            ^~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:393:64: warning: 'template<class _Arg1, class _Arg2, class _Result> struct std::binary_function' is deprecated [-Wdeprecated-declarations]
##   393 |     class const_mem_fun1_t : public boost::functional::detail::binary_function<const T*, A, S>
##       |                                                                ^~~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:131:12: note: declared here
##   131 |     struct binary_function
##       |            ^~~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:438:61: warning: 'template<class _Arg, class _Result> struct std::unary_function' is deprecated [-Wdeprecated-declarations]
##   438 |     class mem_fun_ref_t : public boost::functional::detail::unary_function<T&, S>
##       |                                                             ^~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:117:12: note: declared here
##   117 |     struct unary_function
##       |            ^~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:454:62: warning: 'template<class _Arg1, class _Arg2, class _Result> struct std::binary_function' is deprecated [-Wdeprecated-declarations]
##   454 |     class mem_fun1_ref_t : public boost::functional::detail::binary_function<T&, A, S>
##       |                                                              ^~~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:131:12: note: declared here
##   131 |     struct binary_function
##       |            ^~~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:470:67: warning: 'template<class _Arg, class _Result> struct std::unary_function' is deprecated [-Wdeprecated-declarations]
##   470 |     class const_mem_fun_ref_t : public boost::functional::detail::unary_function<const T&, S>
##       |                                                                   ^~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:117:12: note: declared here
##   117 |     struct unary_function
##       |            ^~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:487:68: warning: 'template<class _Arg1, class _Arg2, class _Result> struct std::binary_function' is deprecated [-Wdeprecated-declarations]
##   487 |     class const_mem_fun1_ref_t : public boost::functional::detail::binary_function<const T&, A, S>
##       |                                                                    ^~~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bi
## ts/stl_function.h:131:12: note: declared here
##   131 |     struct binary_function
##       |            ^~~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:533:73: warning: 'template<class _Arg, class _Result> struct std::unary_function' is deprecated [-Wdeprecated-declarations]
##   533 |     class pointer_to_unary_function : public boost::functional::detail::unary_function<Arg,Result>
##       |                                                                         ^~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:117:12: note: declared here
##   117 |     struct unary_function
##       |            ^~~~~~~~~~~~~~
## stan/lib/stan_math/lib/boost_1.81.0/boost/functional.hpp:557:74: warning: 'template<class _Arg1, class _Arg2, class _Result> struct std::binary_function' is deprecated [-Wdeprecated-declarations]
##   557 |     class pointer_to_binary_function : public boost::functional::detail::binary_function<Arg1,Arg2,Result>
##       |                                                                          ^~~~~~~~~~~~~~~
## C:/rtools43/ucrt64/include/c++/13.2.0/bits/stl_function.h:131:12: note: declared here
##   131 |     struct binary_function
##       |            ^~~~~~~~~~~~~~~
## In file included from stan/lib/stan_math/stan/math/prim/prob/von_mises_lccdf.hpp:5,
##                  from stan/lib/stan_math/stan/math/prim/prob/von_mises_ccdf_log.hpp:4,
##                  from stan/lib/stan_math/stan/math/prim/prob.hpp:359,
##                  from stan/lib/stan_math/stan/math/prim.hpp:16,
##                  from stan/lib/stan_math/stan/math/rev.hpp:16:
## stan/lib/stan_math/stan/math/prim/prob/von_mises_cdf.hpp: In function 'stan::return_type_t<T_x, T_sigma, T_l> stan::math::von_mises_cdf(const T_x&, const T_mu&, const T_k&)':
## stan/lib/stan_math/stan/math/prim/prob/von_mises_cdf.hpp:194: note: '-Wmisleading-indentation' is disabled from this point onwards, since column-tracking was disabled due to the size of the code/headers
##   194 |       if (cdf_n < 0.0)
##       |
## stan/lib/stan_math/stan/math/prim/prob/von_mises_cdf.hpp:194: note: adding '-flarge-source-files' will allow for more column-tracking support, at the expense of compilation time and memory

Now, let’s generate data from the model.

generate_single_sim_1 <- function(T, r_e, r_l) {
  e <- rexp(1, r_e)
  l <- rexp(1, r_l)
  s <- sample.int(T, size = 1)
  
  y <- array(NA_real_, T)
  for(t in 1:T) {
    if(t <= s) {
      rate <- e
    } else {
      rate <- l
    }
    y[t] <- rpois(1, rate) 
  }
  
  list(
    variables = list(
      e = e, l = l, s = s
    ), generated = list(
      T = T,
      r_e = r_e,
      r_l = r_l,
      y = y
    )
  )
}

generator_1 <- SBC_generator_function(generate_single_sim_1, T = 5, r_e = 0.5, r_l = 0.1)
set.seed(85394672)
datasets_1 <- generate_datasets(generator_1, 30)

Additionally, we’ll add a derived quantity expressing the total log-likelihood of data given the fitted parameters. The expression within the derived_quantities() call is evaluated for both prior and posterior draws and included as another variable in SBC checks. It turns out this type of derived quantities can increase the sensitivity of the SBC against some issues in the model. See vignette("limits_of_SBC") for a more detailed discussion of this.

log_lik_dq <- derived_quantities(log_lik = sum(dpois(y, ifelse(1:T < s, e, l), log = TRUE)))

So finally, lets actually compute SBC:

results_1 <- compute_SBC(datasets_1, backend_1, 
                    cache_mode = "results", 
                    cache_location = file.path(cache_dir, "model1"),
                    dquants = log_lik_dq)
## Results loaded from cache file 'model1'
##  - 3 (10%) fits had at least one Rhat > 1.01. Largest Rhat was NA.
##  - 20 (67%) fits had tail ESS undefined or less than half of the maximum rank, potentially skewing 
## the rank statistics. The lowest tail ESS was NA.
##  If the fits look good otherwise, increasing `thin_ranks` (via recompute_SBC_statistics) 
## or number of posterior draws (by refitting) might help.
##  - 2 (7%) fits had divergent transitions. Maximum number of divergences was 2.
## Not all diagnostics are OK.
## You can learn more by inspecting $default_diagnostics, $backend_diagnostics 
## and/or investigating $outputs/$messages/$warnings for detailed output from the backend.

Here we also use the caching feature to avoid recomputing the fits when recompiling this vignette. In practice, caching is not necessary but is often useful.

TODO the diagnostic failures are false positives, because Rhat and ESS don’t work very well for discrete parameters. We need to figure out how to handle this better.

We can quickly note that the statistics for the s parameter are extreme - many ranks of 0 and extreme z-scores, including -Infinity. Seing just one or two such fits should be enough to convince us that there is something fundamentally wrong.

dplyr::filter(results_1$stats, variable == "s") 
## # A tibble: 30 × 15
##    sim_id variable simulated_value  rank   z_score  mean median    sd   mad
##     <int> <chr>              <dbl> <dbl>     <dbl> <dbl>  <dbl> <dbl> <dbl>
##  1      1 s                      3   179    0.0523  2.92      3 1.61   2.97
##  2      2 s                      1    14   -1.91    2.02      2 0.534  0   
##  3      3 s                      4   128   -1.36    4.66      5 0.485  0   
##  4      4 s                      1     4   -3.51    2.91      3 0.545  0   
##  5      5 s                      5   398    2.81    2.86      3 0.762  0   
##  6      6 s                      2   275    0.0169  1.98      1 1.44   0   
##  7      7 s                      3     0 -Inf       4         4 0      0   
##  8      8 s                      2   128   -0.580   2.84      3 1.45   1.48
##  9      9 s                      2     0   -7.97    2.99      3 0.124  0   
## 10     10 s                      2     2   -8.42    3.00      3 0.119  0   
## # ℹ 20 more rows
## # ℹ 6 more variables: q5 <dbl>, q95 <dbl>, rhat <dbl>, ess_bulk <dbl>,
## #   ess_tail <dbl>, max_rank <int>

Inspecting the statistics shows that quite often, the model is quite sure of the value of s while the simulated value is just one less.

Looking at the ecdf_diff plot we see that this seems to compromise heavily the inference for s, but the other parameters do not show such bad behaviour. Note that the log_lik derived quantity shows even starker failure than s, so it indeed poses a stricter check in this scenario.

plot_ecdf_diff(results_1)

plot_rank_hist(results_1)

An important note: you may wonder, how we got such a wiggly line for the s parameter - doesn’t it have just 5 possible values? Shouldn’t therefore the ECDF be one big staircase? In fact the package does a little trick to make discrete parameters comparable to continuous - the rank of a discrete parameter is chosen uniformly randomly across all possible ranks (i.e. posterior draws that have exactly equal value). This means that if the model is well behaved, ranks for the discrete parameter will be uniformly distributed across the whole range of possible ranks and we can use exactly the same diagnostics for a discrete parameter as we do for the continuous ones.

But back to the model - what happened? What is wrong with it? After some inspection, you may notice that the simulator does not match the model - the model takes the early rate (e) for points t < s while the simulator takes e for points t <= s, so there is effectively a shift by one time point between the simulator and the model. So let’s assume that we beleive that the Stan model is in fact right. We therefore updated the simulator to match the model:

generate_single_sim_2 <- function(T, r_e, r_l) {
  e <- rexp(1, r_e)
  l <- rexp(1, r_l)
  s <- sample.int(T, size = 1)
  
  y <- array(NA_real_, T)
  for(t in 1:T) {
    if(t < s) { ### <--- Only change here
      rate <- e
    } else {
      rate <- l
    }
    y[t] <- rpois(1, rate) 
  }
  
  list(
    variables = list(
      e = e, l = l, s = s
    ), generated = list(
      T = T,
      r_e = r_e,
      r_l = r_l,
      y = y
    )
  )
}

generator_2 <- SBC_generator_function(generate_single_sim_2, T = 5, r_e = 0.5, r_l = 0.1)

And we can recompute:

set.seed(5846502)
datasets_2 <- generate_datasets(generator_2, 30)
results_2 <- compute_SBC(datasets_2, backend_1,
                    dquants = log_lik_dq, 
                    cache_mode = "results", 
                    cache_location = file.path(cache_dir, "model2"))
## Results loaded from cache file 'model2'
##  - 8 (27%) fits had at least one Rhat > 1.01. Largest Rhat was NA.
##  - 24 (80%) fits had tail ESS undefined or less than half of the maximum rank, potentially skewing 
## the rank statistics. The lowest tail ESS was NA.
##  If the fits look good otherwise, increasing `thin_ranks` (via recompute_SBC_statistics) 
## or number of posterior draws (by refitting) might help.
##  - 2 (7%) fits had divergent transitions. Maximum number of divergences was 4.
## Not all diagnostics are OK.
## You can learn more by inspecting $default_diagnostics, $backend_diagnostics 
## and/or investigating $outputs/$messages/$warnings for detailed output from the backend.
plot_rank_hist(results_2)

plot_ecdf_diff(results_2)

Looks good, so let us add some more simulations to make sure the model behaves well.

set.seed(54321488)
datasets_2_more <- generate_datasets(generator_2, 100)
results_2_more <- compute_SBC(datasets_2_more, backend_1,
                    dquants = log_lik_dq, 
                    cache_mode = "results", 
                    cache_location = file.path(cache_dir, "model3"))
## Results loaded from cache file 'model3'
##  - 16 (16%) fits had at least one Rhat > 1.01. Largest Rhat was NA.
##  - 71 (71%) fits had tail ESS undefined or less than half of the maximum rank, potentially skewing 
## the rank statistics. The lowest tail ESS was NA.
##  If the fits look good otherwise, increasing `thin_ranks` (via recompute_SBC_statistics) 
## or number of posterior draws (by refitting) might help.
##  - 7 (7%) fits had divergent transitions. Maximum number of divergences was 19.
## Not all diagnostics are OK.
## You can learn more by inspecting $default_diagnostics, $backend_diagnostics 
## and/or investigating $outputs/$messages/$warnings for detailed output from the backend.
results_2_all <- bind_results(results_2, results_2_more)
plot_rank_hist(results_2_all)

plot_ecdf_diff(results_2_all)

Now - as far as this amount of SBC steps can see, the model is good and we get good behaviour for both the continuous and the discrete parameters and the log_lik derived quantity. Hooray!

JAGS version

We can now write the same model in JAGS. This becomes a bit easier as JAGS lets us represent discrete parameters directly:

cat(readLines("other_models/changepoint.jags"), sep = "\n")
data {
  for(i in 1:T) {
    prior_s[i] = 1.0/T
  }
}

model {
  e ~ dexp(r_e);
  l ~ dexp(r_l);
  s ~ dcat(prior_s)
  for(i in 1:T) {
      y[i] ~ dpois(ifelse(i < s, e, l))
  }
}

We will use the rjags package, let us verify it is installed correctly.

## Loading required package: coda
## Linked to JAGS 4.3.1
## Loaded modules: basemod,bugs

We will also default to relatively large number of samples as we can expect some autocorrelation in the Gibbs sampler.

backend_jags <- SBC_backend_rjags("other_models/changepoint.jags",
                                  variable.names = c("e","l","s"),
                                  n.iter = 10000,
                                  n.burnin = 1000,
                                  n.chains = 4,
                                  thin = 10)

Running SBC with all the corrected datasets from before (rJAGS accepts input data in exactly the same format as Stan, so we can reuse the datasets without any change):

datasets_2_all <- bind_datasets(datasets_2, datasets_2_more)
results_jags <- compute_SBC(datasets_2_all, backend_jags,
                            dquants = log_lik_dq,
                        cache_mode = "results",
                        cache_location = file.path(cache_dir_jags, "rjags"))
## Results loaded from cache file 'rjags'
##  - 19 (15%) fits had at least one Rhat > 1.01. Largest Rhat was NA.
##  - 95 (73%) fits had tail ESS undefined or less than half of the maximum rank, potentially skewing 
## the rank statistics. The lowest tail ESS was NA.
##  If the fits look good otherwise, increasing `thin_ranks` (via recompute_SBC_statistics) 
## or number of posterior draws (by refitting) might help.
## Not all diagnostics are OK.
## You can learn more by inspecting $default_diagnostics, $backend_diagnostics 
## and/or investigating $outputs/$messages/$warnings for detailed output from the backend.

Similarly to the case above, the Rhat and ESS warnings are false positives due to the s parameter, which we need to handle better.

The rank plots show no problems.

plot_rank_hist(results_jags)

plot_ecdf_diff(results_jags)

As an exercise, we can also write the marginalized version of the model in JAGS. In some cases, marginalization improves performance even for JAGS models, however, for this model it is actually not an improvement, presumably because the model is very simple.

cat(readLines("other_models/changepoint_marginalized.jags"), sep = "\n")
data {
  for(i in 1:T) {
    prior_unif[i] = -log(T)
  }

  # Using the zeroes crossing trick to compute the likelihood
  # See e.g. https://667-per-cm.net/2014/02/17/the-zero-crossings-trick-for-jags-finding-roots-stochastically/
  z = 0
}

model {
  e ~ dexp(r_e);
  l ~ dexp(r_l);

  # Prepare the zero trick
  z ~ dpois(z_mean)

  # Compute the likelihood
  # The lp is a matrix to avoid having to redefine nodes
  lp[1, 1:T] = prior_unif
  for (s in 1:T) {
    for (t in 1:T) {
      lp[1 + t, s] = lp[t, s] + log(ifelse(t < s, e, l)) * y[t] - ifelse(t < s, e, l)
    }
    p[s] = exp(lp[T + 1, s])
  }

  # log-sum-exp to compute the log likelihood in a numerically stable way
  m = max(lp[T + 1, ])
  sum_exp_rest[1] = 0
  for(t in 1:T) {
    sum_exp_rest[1 + t] = sum_exp_rest[t] + exp(lp[T + 1, s] - m)
  }
  lp_total = m + log(sum_exp_rest[T + 1])

  # We have the likelihood now add it to z_mean for the zeros trick
  z_mean = -lp_total + 10000

  s ~ dcat(p)
}

The code got quite a bit more complex, se let’s check if we didn’t mess up the rewrite - first we build a backend with this new representation:

backend_jags_marginalized <- SBC_backend_rjags("other_models/changepoint_marginalized.jags",
                                  variable.names = c("e","l","s"),
                                  n.iter = 10000,
                                  n.burnin = 1000,
                                  n.chains = 4,
                                  thin = 10)

Then we run the actual SBC:

results_jags_marginalized <- compute_SBC(datasets_2_all, backend_jags_marginalized,
                                         dquants = log_lik_dq,
                        cache_mode = "results",
                        cache_location = file.path(cache_dir_jags, "rjags_marginalized"))
## Results loaded from cache file 'rjags_marginalized'
##  - 24 (18%) fits had at least one Rhat > 1.01. Largest Rhat was NA.
##  - 95 (73%) fits had tail ESS undefined or less than half of the maximum rank, potentially skewing 
## the rank statistics. The lowest tail ESS was NA.
##  If the fits look good otherwise, increasing `thin_ranks` (via recompute_SBC_statistics) 
## or number of posterior draws (by refitting) might help.
## Not all diagnostics are OK.
## You can learn more by inspecting $default_diagnostics, $backend_diagnostics 
## and/or investigating $outputs/$messages/$warnings for detailed output from the backend.

And the ranks plots look good, so we indeed probably did succeed in correctly marginalizing the s parameter!

plot_rank_hist(results_jags_marginalized)

plot_ecdf_diff(results_jags_marginalized)