def estimate_r0(cases_df, place, sample_size, min_days, w_date): used_brazil = False incidence = (cases_df[place].query("totalCases > @MIN_CASES_TH").pipe( prepare_for_r0_estimation)[:w_date]) if len(incidence) < MIN_DAYS_r0_ESTIMATE: used_brazil = True incidence = (make_brazil_cases(cases_df).pipe( prepare_for_r0_estimation)[:w_date]) Rt = ReproductionNumber(incidence=incidence, prior_shape=5.12, prior_scale=0.64, si_pars={ 'mean': 4.89, 'sd': 1.48 }, window_width=MIN_DAYS_r0_ESTIMATE - 2) Rt.compute_posterior_parameters() return Rt.sample_from_posterior(sample_size=sample_size)
def run_Rt_estimation(incidence, prior_shape, prior_scale, mean_si, sd_si, t_max, window_width): si_pars = {'mean': mean_si, 'sd': sd_si} Rt = ReproductionNumber(incidence=incidence, si_pars=si_pars, prior_shape=prior_shape, prior_scale=prior_scale, window_width=window_width) Rt.compute_posterior_parameters() Rt_posterior_sample = Rt.sample_from_posterior(sample_size=N) Rt.compute_posterior_summaries(posterior_sample=Rt_posterior_sample, t_max=t_max) results = Rt return results