diff --git a/NEWS.md b/NEWS.md index 6152fe738..184ddc14a 100644 --- a/NEWS.md +++ b/NEWS.md @@ -49,6 +49,7 @@ - Delay distribution discretisation now properly accounts for primary event censoring during model fitting, matching the correction already applied on the R side since v1.8.0. This improves accuracy for short delays where the observation window is large relative to the delay. - Left truncation of delay distributions (e.g. excluding generation times of zero) is now handled analytically rather than by zeroing and renormalising, giving more accurate PMFs near the truncation point. +- Reduced repeated Stan work in the Rt-to-growth Newton solver by reusing the generation-time support sequence across iterations. ## Package changes diff --git a/inst/stan/functions/rt.stan b/inst/stan/functions/rt.stan index 0c1a85ce2..fb7c55975 100644 --- a/inst/stan/functions/rt.stan +++ b/inst/stan/functions/rt.stan @@ -94,19 +94,28 @@ void rt_lp(array[] real initial_infections_scale, vector bp_effects, * @param R Reproduction number * @param r Current estimate of the growth rate * @param pmf Generation time probability mass function (first index: 0) + * @param zero_series Precomputed vector of indices `(0, 1, ..., len - 1)` + * used to align `pmf` with lagged terms in the calculation; expected to have + * the same length as `pmf`. * @return The Newton step for updating r * * @ingroup rt_estimation */ -real R_to_r_newton_step(real R, real r, vector pmf) { +real R_to_r_newton_step_with_support(real R, real r, vector pmf, + vector zero_series) { int len = num_elements(pmf); - vector[len] zero_series = linspaced_vector(len, 0, len - 1); vector[len] exp_r = exp(-r * zero_series); real ret = (R * dot_product(pmf, exp_r) - 1) / (- R * dot_product(pmf .* zero_series, exp_r)); return(ret); } +real R_to_r_newton_step(real R, real r, vector pmf) { + int len = num_elements(pmf); + vector[len] zero_series = linspaced_vector(len, 0, len - 1); + return R_to_r_newton_step_with_support(R, r, pmf, zero_series); +} + /** * Estimate the growth rate r from reproduction number R * @@ -128,11 +137,12 @@ real R_to_r_newton_step(real R, real r, vector pmf) { real R_to_r(real R, vector gt_rev_pmf, real abs_tol) { int gt_len = num_elements(gt_rev_pmf); vector[gt_len] gt_pmf = reverse(gt_rev_pmf); - real mean_gt = dot_product(gt_pmf, linspaced_vector(gt_len, 0, gt_len - 1)); + vector[gt_len] zero_series = linspaced_vector(gt_len, 0, gt_len - 1); + real mean_gt = dot_product(gt_pmf, zero_series); real r = fmax((R - 1) / (R * mean_gt), -1); real step = abs_tol + 1; while (abs(step) > abs_tol) { - step = R_to_r_newton_step(R, r, gt_pmf); + step = R_to_r_newton_step_with_support(R, r, gt_pmf, zero_series); r -= step; }