diff --git a/inst/stan/functions/observation_model.stan b/inst/stan/functions/observation_model.stan index e72dde771..1fb7519cf 100644 --- a/inst/stan/functions/observation_model.stan +++ b/inst/stan/functions/observation_model.stan @@ -195,6 +195,36 @@ vector report_log_lik(array[] int cases, vector reports, return(log_lik); } + +/** + * Custom safe version of the negative binomial sampler + * + * This function generates random samples of the negative binomial distribution + * whilst avoiding numerical overflows. In particular: + * - if the mu parameter is very small it always returns 0 + * - if the phi parameter is large it returns a sample from a Poisosn + * distribution + * - if the gamma rate of the gamma-Poisson mixture used for simulating from the + * distribution is very large, it returns 1e8 + * - in all other cases it returns a sample from the negative binomial + * distribution + * + * @param mu Real value for mean mu. + * @param phi Real value for phi. + * + * @return A random sample + */ +int neg_binomial_2_safe_rng(real mu, real phi) { + if (mu < 1e-8) { + return(0); + } else if (phi > 1e4) { + return(poisson_rng(mu > 1e8 ? 1e8 : mu)); + } else { + real gamma_rate = gamma_rng(phi, phi / mu); + return(poisson_rng(gamma_rate > 1e8 ? 1e8 : gamma_rate)); + } +} + /** * Generate random samples of reported cases * @@ -217,16 +247,7 @@ array[] int report_rng(vector reports, real dispersion, int model_type) { } for (s in 1:t) { - if (reports[s] < 1e-8) { - sampled_reports[s] = 0; - } else { - // defer to poisson if phi is large, to avoid overflow - if (phi > 1e4) { - sampled_reports[s] = poisson_rng(reports[s] > 1e8 ? 1e8 : reports[s]); - } else { - sampled_reports[s] = neg_binomial_2_rng(reports[s] > 1e8 ? 1e8 : reports[s], phi); - } - } + sampled_reports[s] = neg_binomial_2_safe_rng(reports[s], phi); } return(sampled_reports); }