Skip to content
Open
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions inst/stan/functions/rt.stan
Original file line number Diff line number Diff line change
Expand Up @@ -98,15 +98,20 @@ void rt_lp(array[] real initial_infections_scale, vector bp_effects,
*
* @ingroup rt_estimation
*/
real R_to_r_newton_step(real R, real r, vector pmf) {
real R_to_r_newton_step(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(R, r, pmf, zero_series);
}

/**
* Estimate the growth rate r from reproduction number R
*
Expand All @@ -128,11 +133,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(R, r, gt_pmf, zero_series);
r -= step;
}

Expand Down
Loading