def show_fitting(args, N, data, country_name, params):
dates = list(data.index)
days = [int(i) for i in range(0, data.shape[0])]
dates_all = dates + cu.get_dates(dates[-1], 30)
# get the prediction
E = Epidemology(args.model_name, 'solve_ivp', data.shape[0] + 30)
E.set_init(N, data[0], 0)
sir = E.evolve(params)
title = country_name + r', $\beta_0$' +"=" + str('%.2f' %params[0])\
+ r', $\mu$= ' + str('%.2f' % params[1]) \
+ r', $\gamma$= '+ str('%.2f' % params[2])
t = cu.strip_year(dates)
t_all = cu.strip_year(dates_all)
fig = plt.figure(figsize=(15, 10))
ax = fig.add_subplot(111)
ax.plot(t_all, sir.y[1, :], c='r', label='Infected')
ax.plot(t_all, sir.y[2, :], c='g', label='Recovered')
plt.scatter(t, data.to_numpy(), c='b', label='Data')
ax.set_yscale('log')
ax.set_title(title)
ax.legend()
plt.setp(ax.get_xticklabels(), rotation=90, horizontalalignment='right')
plot_dir = args.output_dir + os.sep + "plots"
os.makedirs(plot_dir, exist_ok=True)
plt.savefig(plot_dir + os.sep + country_name + ".pdf")
print(count, row['date'], row['confirmed'], row['recovered'],
row['deaths'])
count += 1
# let us get all the dates
dates = df['date'].to_list()
# add some future dates also
dates_predict = dates + get_dates(dates[df.shape[0] - 1], args.num_predict)
# get the days
days = np.array([float(i) for i in range(0, len(dates))])
days_predict = np.array([float(i) for i in range(0, len(dates_predict))])
# let us get rid of the year for plotting
dates_predict = strip_year(dates_predict)
dates = strip_year(dates)
# now let us iterate over the three columns
cases = ["deaths", "recovered", "confirmed"]
fig = plt.figure(figsize=(18, 18))
ax = [plt.subplot(311)]
ax.append(plt.subplot(312))
ax.append(plt.subplot(313))
num_days = len(days)
for i in range(0, len(cases)):
y = df[cases[i]].to_numpy()
if end == 0:
end = df.shape[0]
x_fit = np.array(x[start:end] - x[start])
y_fit = np.log(y[start:end])
slope, intercept, r_value, p_value, std_err = stats.linregress(
x_fit, y_fit)
y_predict = intercept + slope * (x[start:] - x[start])
date_str = get_file_name(dates, start, end)
output_file = args.output_dir + os.sep + args.country + "_" + args.type + "_" + date_str + ".pdf"
dates = strip_year(df['date'].to_list())
fig = plt.figure()
ax = fig.add_subplot()
ax.set_title(args.type)
ax.plot(x, np.log(y), 'o', label='data')
ax.plot(x[start:], y_predict)
ax.plot(x[start:], y_predict,'k-',\
label='fit: a=%5.3f, b=%5.3f, fit window =[%d-%d]' % (intercept,slope,start,end))
ax.legend(loc='upper center')
ax.set_ylabel("log( " + args.type + " )")
plt.setp(ax.get_xticklabels(), rotation=90, horizontalalignment='right')
plt.savefig(output_file)
print("output file:", output_file)
dR = R.diff(periods=1).iloc[1:]
I1 = I.iloc[1:]
# equation (11) of arXiv:2003.00122v5
beta1 = (dI + dR).divide(I1, fill_value=0)
# For our reconstruction
N, L, T = lockdown_info (args.lockdown_file, args.country_name)
R = BetaSolver(df, N)
beta2 = R.solve(7.0, 7.0, 1.0).iloc[1:-2]
fig = plt.figure(figsize=(18,12))
ax = fig.add_subplot()
beta1.index = strip_year(beta1.index)
beta2.index = strip_year(beta2.index)
ax.plot(beta1,'o',c='b')
ax.plot(beta1,label='Direct',c='b')
ax.plot(beta2,label='Reconstructed',c='r')
ax.plot(beta2,'o',c='r')
ax.set_xlabel('t')
ax.set_ylabel(r'$\beta (t)$')
ax.set_title(args.country_name)
plt.legend()
ax.axhline(y=0,c='k',ls='--')
plt.setp(ax.get_xticklabels(), rotation=90, horizontalalignment='right',fontsize=10)
plt.show()