def fit_population(args):
tp_to_eval_seroprevalence = 30
from scripts.estimate_R0 import estimate_Re
from scripts.load_utils import cumulative_to_rolling_average
region, json_case_data, pop_params, age_distribution, return_fit_res = args
print(f"starting fit for {region}")
time_range_fit = 60
data = convert_to_vectors(json_case_data)
weekly_data = cumulative_to_rolling_average(data)
if weekly_data['cases'] is None or weekly_data['deaths'] is None:
return (region, None, None) if return_fit_res else (region, None)
weekly_data_for_fit = {
k: v[-time_range_fit + 7:]
for k, v in weekly_data.items()
}
# estimate Re and bound from above (usually a problem at the very beginning of a wave)
res = estimate_Re(weekly_data['cases'], imports=1, cutoff=5, chop=3)
Re = np.minimum(pop_params['r0'], res['Re'])
Refit = fit_REblocks(weekly_data['time'][:-8], Re)
piecewise_constant_Re = Refit.predict(weekly_data['time'].reshape(-1, 1))
change_points = np.concatenate(
[[-1],
np.where(np.diff(piecewise_constant_Re) != 0)[0], [len(Re) - 1]]) + 1
total_deaths_at_start = data['deaths'][-tp_to_eval_seroprevalence]
IFR = get_IFR(age_distribution)
reporting_fraction = get_reporting_fraction(data['cases'], data['deaths'],
IFR)
n_cases_at_start = total_deaths_at_start / IFR
seroprevalence = n_cases_at_start / pop_params['size']
if np.isnan(seroprevalence) or (type(seroprevalence)
== np.ma.core.MaskedConstant
and seroprevalence.mask == True):
seroprevalence = 0
mitigations = []
for ci, cp in enumerate(change_points[:-1]):
mitigations.append({
'tMin':
int(weekly_data['time'][cp]),
'tMax':
int(weekly_data['time'][change_points[ci + 1] - 1]),
'value':
max(
0,
min(
1,
float(1 - piecewise_constant_Re[cp] /
(1 - seroprevalence) / pop_params['r0'])))
})
tmin = data['time'][-time_range_fit]
average_Re = np.mean(Re[-30:])
fixed_params = {
'logR0': np.log(pop_params['r0']),
'efficacy': 1 - average_Re / pop_params['r0'],
'containment_start': tmin,
'seroprevalence': seroprevalence
}
if (weekly_data_for_fit['cases'][0]):
guess = {
'logInitial':
np.log(
(0.5 + weekly_data_for_fit['cases'][0]) / reporting_fraction),
'reported':
reporting_fraction
}
elif (weekly_data_for_fit['deaths'][0]) and reporting_fraction > 0:
guess = {
'logInitial':
np.log((0.5 + weekly_data_for_fit['deaths'][0]) /
reporting_fraction / IFR),
'reported':
reporting_fraction
}
elif (weekly_data_for_fit['deaths'][0]):
guess = {
'logInitial':
np.log((0.5 + weekly_data_for_fit['deaths'][0]) / 0.3 / IFR),
'reported':
reporting_fraction
}
else:
guess = {'logInitial': np.log(1), 'reported': reporting_fraction}
fit_result, success = fit_params(data['time'][-time_range_fit:],
weekly_data_for_fit,
guess,
age_distribution,
pop_params['size'],
fixed_params=fixed_params)
params = {}
params.update(fixed_params)
for p in guess:
params[p] = fit_result.__getattribute__(p)
for p in params:
params[p] = float(params[p])
params['mitigations'] = mitigations
params['tMin'] = datetime.fromordinal(tmin).date().strftime('%Y-%m-%d')
params['tMax'] = datetime.fromordinal(tmin +
90).date().strftime('%Y-%m-%d')
if return_fit_res:
return (region, params, fit_result)
else:
return (region, params)
import ipdb
ipdb.set_trace()
from matplotlib import pyplot as plt
case_counts = parse_tsv()
scenario_data = load_population_data()
age_distributions = load_distribution()
# region = 'JPN-Kagawa'
region = 'United States of America'
# region = 'Germany'
region = 'Switzerland'
region = 'USA-Texas'
age_dis = age_distributions[scenario_data[region]['ages']]
region, p, fit_params = fit_population(
(region, case_counts[region], scenario_data[region], age_dis, True))
model_data = generate_data(fit_params)
model_cases = model_data['cases'][7:] - model_data['cases'][:-7]
model_deaths = model_data['deaths'][7:] - model_data['deaths'][:-7]
model_time = fit_params.time[7:]
cases = cumulative_to_rolling_average(
convert_to_vectors(case_counts[region]))
print(fit_params.reported, np.exp(fit_params.logInitial))
print(get_IFR(age_dis))
plt.plot(model_time, model_cases)
plt.plot(model_time, model_deaths)
plt.plot(cases['time'], cases['cases'])
plt.plot(cases['time'], cases['deaths'])
plt.yscale('log')