#!/usr/bin/env python
# coding: utf-8

# In[1]:


from lifelines.datasets import load_rossi
rossi = load_rossi()


# In[2]:


rossi.head()


# In[3]:


# let's b-spline age
cph = CoxPHFitter().fit(rossi, "week", "arrest", formula="fin + bs(age, df=4) + wexp + mar + paro + prio")


# In[4]:


# now we need to "extend" our data to plot it
# we'll plot age over it's observed range
age_range = np.linspace(rossi['age'].min(), rossi['age'].max(), 50)

# need to create a matrix of variables at their means, _except_ for age. 
x_bar = cph._central_values
df_varying_age = pd.concat([x_bar] * 50).reset_index(drop=True)
df_varying_age['age'] = age_range

df_varying_age.head()


# In[5]:


cph.predict_log_partial_hazard(df_varying_age).plot()


# In[6]:


# compare to _not_ bspline-ing:
cph = CoxPHFitter().fit(rossi, "week", "arrest", formula="fin + age + wexp + mar + paro + prio")

age_range = np.linspace(rossi['age'].min(), rossi['age'].max(), 50)

# need to create a matrix of variables at their means, _except_ for age. 
x_bar = cph._central_values
df_varying_age = pd.concat([x_bar] * 50).reset_index(drop=True)
df_varying_age['age'] = age_range

cph.predict_log_partial_hazard(df_varying_age).plot()


# In[ ]:





# In[ ]: