def smooth_and_remove_step(x_lst, y_lst, x_min_flt,x_max_flt,rmv_step_bool):
'''
Takes entire data set, x and y
cuts down the spectra s.t x_min < x < x_max
THEN
Removes a step function from y_lst
'''
# Restrict the fit
x_fit = []
y_fit = []
top_lst = []
bottom_lst = []
for x,y in zip(x_lst, y_lst):
# Restrict the fitting region
if x_min_flt < x < x_max_flt:
x_fit.append(float(x))
y_fit.append(float(y))
# Find top and bottom of step
if x < x_min_flt + 7:
bottom_lst.append(float(y))
elif x > x_max_flt - 7:
top_lst.append(float(y))
x_fit = np.asarray(x_fit)
y_fit = np.asarray(y_fit)
top = np.mean(np.asarray(top_lst))
bottom = np.mean(np.asarray(bottom_lst))
delta = top-bottom
if (rmv_step_bool):
# Step Parameters
step_at = 100
step_width = 1
pp = Parameters()
pp.add_many(
('amplitude',delta),
('sigma',step_width),
('center',step_at)
)
step = StepModel(form = 'erf', prefix='', independent_vars=['x'])
y_fit = np.asarray([yy-bottom-step.eval(x=xx, params=pp) for xx,yy in zip(x_fit,y_fit)])
# rest is the same as smooth_the_data
# now we find the parameters using the - d^2/dx^2
ysmooth = interp.interp1d(x_fit, y_fit, kind='cubic')
# differentiate x 2
yp = np.gradient(ysmooth(x_fit))
ypp = np.gradient(yp)
# we want the peaks of -d2/dx2
ypp = np.asarray([-x for x in ypp])
return x_fit, y_fit, ysmooth, yp, ypp
def predictive_model(data: pd.DataFrame,
interesting_rows,
day_zero_n_patients: int = 20,
days_in_future: int = 30,
aggregated: bool = False):
data = data[interesting_rows].iloc[:, :]
from lmfit.models import StepModel, ExponentialModel
fig = plt.figure(figsize=(10, 5))
for c in range(len(data.index)):
if aggregated:
values = data.values[c, 4:][data.iloc[c, 4:] > day_zero_n_patients]
else:
values = np.concatenate(
([0],
np.diff(
data.values[c,
4:][data.iloc[c, 4:] > day_zero_n_patients])))
n = values.shape[0]
x = np.asarray(range(values.shape[0]), dtype='float64')
y = np.asarray(values, dtype='float64')
if len(x) == 0:
continue
label = "{}-{}".format(data.values[c, 0], data.values[c, 1])
plt.plot(x, y, label=label)
if data.values[c, 1] in ["China", "US"]:
continue
try:
model_step = StepModel()
model_exp = ExponentialModel()
params_step = model_step.guess(y, x=x)
params_exp = model_exp.guess(y, x=x)
result_step = model_step.fit(y, params_step, x=x)
result_exp = model_exp.fit(y, params_exp, x=x)
except Exception:
continue
x_pred = np.asarray(range(days_in_future))
plt.plot(x_pred,
model_step.eval(result_step.params, x=x_pred),
':',
label='fit-{}'.format(label))
plt.plot(x_pred,
model_exp.eval(result_exp.params, x=x_pred),
'.',
label='fit-{}'.format(label))
# print(result.fit_report())
# result.plot_fit()
plt.legend(prop={"size": 7})
plt.yscale('log')
plt.xticks(rotation=45)
plt.grid(which='both')
now = datetime.now()
dt_string = now.strftime("%d%m%Y-%H%M%S")
