class BasicModel(ModelPipeline):
def __init__(self, fit_dict, basic_model_dict, **pipeline_kwargs):
"""
Generic class for a function to produce predictions from a model
with the following attributes.
Args:
**pipeline_kwargs: keyword arguments for the base class of ModelPipeline
predict_group: (str) which group to make predictions for
fit_dict: keyword arguments to CurveModel.fit_params()
basic_model_dict: additional keyword arguments to the CurveModel class
col_obs_se: (str) of observation standard error
col_covs: List[str] list of names of covariates to put on the parameters
param_names (list{str}):
Names of the parameters in the specific functional form.
link_fun (list{function}):
List of link functions for each parameter.
var_link_fun (list{function}):
List of link functions for the variables including fixed effects
and random effects.
"""
super().__init__(**pipeline_kwargs)
self.fit_dict = fit_dict
self.basic_model_dict = basic_model_dict
self.basic_model_dict.update({'col_obs_se': self.col_obs_se})
generator_kwargs = pipeline_kwargs
for arg in self.pop_cols:
generator_kwargs.pop(arg)
self.basic_model_dict.update(**generator_kwargs)
self.mod = None
self.setup_pipeline()
def refresh(self):
self.mod = None
def fit(self, df, group=None):
self.mod = CurveModel(df=df, **self.basic_model_dict)
self.mod.fit_params(**self.fit_dict)
def predict(self, times, predict_space, predict_group):
predictions = self.mod.predict(
t=times,
group_name=predict_group,
prediction_functional_form=predict_space)
return predictions
# Set up the CurveModel
model = CurveModel(
df=df,
col_t='time',
col_obs='ln_death_rate',
col_group='group',
col_covs=[['intercept'], ['intercept'], ['intercept']],
param_names=['alpha', 'beta', 'p'],
link_fun=[lambda x: x, lambda x: x, lambda x: x],
var_link_fun=[lambda x: x, lambda x: x, lambda x: x],
fun=ln_gaussian_cdf
)
# Fit the model to estimate parameters
model.fit_params(fe_init=[0, 0, 1.],
fe_gprior=[[0, np.inf], [0, np.inf], [1., np.inf]],
options={'disp':True})
# Get predictions
y_pred = model.predict(
t=np.linspace(0,100,num=100),
group_name="Kerala"
)
ground_truth = df.ln_death_rate[df['group']=="Hubei"].reset_index(drop=True)
print(ground_truth)
print(np.exp(y_pred))
# Plot results
plt.plot(np.linspace(0,100,num=100), y_pred, '-')
plt.plot(ground_truth, '.')
plt.show()
class GaussianCDF:
"""Fit a single Gaussian Atom to cumulative daily deaths"""
def __init__(self):
self.model = None
def fit(self, daily_deaths, social_distance=None):
daily_deaths = np.array(daily_deaths)
n_data = daily_deaths.shape[0]
# Prepare the data frame
df = pd.DataFrame()
df['death_rate'] = np.cumsum(daily_deaths)
df['time'] = np.arange(df['death_rate'].shape[0])
df['ln_death_rate'] = np.log(
df['death_rate'] + 1) # Add 1 to pad in case the #deaths are zero
df['group'] = ['all'] * n_data
df['cov_one'] = [1.0] * n_data
if social_distance is not None:
df['social_distance'] = social_distance
col_covs = [['cov_one'], ['cov_one', 'social_distance'],
['cov_one']]
num_fe = 4
fe_init = [-3, 100, 0, 10]
# col_covs = [['cov_one', 'social_distance'], ['cov_one', 'social_distance'], ['cov_one', 'social_distance']]
# num_fe = 6
# fe_init = [-3, 0, 100, 0, 10, 0]
else:
col_covs = [['cov_one'], ['cov_one'], ['cov_one']]
num_fe = 3
fe_init = [-3, 100, 1]
# Set up the CurveModel
self.model = CurveModel(
df=df,
col_t='time',
col_obs='ln_death_rate',
col_group='group',
col_covs=col_covs,
param_names=['alpha', 'beta', 'p'],
link_fun=[lambda x: np.exp(x), lambda x: x, lambda x: np.exp(x)],
var_link_fun=[lambda x: x] * num_fe,
fun=ln_gaussian_cdf)
# Fit the model to estimate parameters
self.model.fit_params(
fe_init=fe_init,
options={
'ftol': 1e-14,
'maxiter': 500
},
re_bounds=[[0, 0]] * num_fe # No random effects
)
def predict(self, t):
"""Get predictions for values in t"""
return self.model.predict(t=t, group_name='all')
def get_params(self):
return np.squeeze(self.model.params)