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")