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

# <style>
#     @font-face {
#         font-family: "Computer Modern";
#         src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
#     }
#     div.cell{
#         width:800px;
#         margin-left:16% !important;
#         margin-right:auto;
#     }
#     h1 {
#         font-family: Helvetica, serif;
#     }
#     h4{
#         margin-top:12px;
#         margin-bottom: 3px;
#        }
#     div.text_cell_render{
#         font-family: Computer Modern, "Helvetica Neue", Arial, Helvetica, Geneva, sans-serif;
#         line-height: 145%;
#         font-size: 130%;
#         width:800px;
#         margin-left:auto;
#         margin-right:auto;
#     }
#     .CodeMirror{
#             font-family: "Source Code Pro", source-code-pro,Consolas, monospace;
#     }
#     .prompt{
#         display: None;
#     }
#     .text_cell_render h5 {
#         font-weight: 300;
#         font-size: 22pt;
#         color: #4057A1;
#         font-style: italic;
#         margin-bottom: .5em;
#         margin-top: 0.5em;
#         display: block;
#     }
#     
#     .warning{
#         color: rgb( 240, 20, 20 )
#         }  
# </style>
# <script>
#     MathJax.Hub.Config({
#                         TeX: {
#                            extensions: ["AMSmath.js"]
#                            },
#                 tex2jax: {
#                     inlineMath: [ ['$','$'], ["\\(","\\)"] ],
#                     displayMath: [ ['$$','$$'], ["\\[","\\]"] ]
#                 },
#                 displayAlign: 'center', // Change this to 'center' to center equations.
#                 "HTML-CSS": {
#                     styles: {'.MathJax_Display': {"margin": 4}}
#                 }
#         });
# </script>

# # Divorce Rates

# In[1]:


from lifelines.estimation import AalenAdditiveFitter
import pandas as pd
import numpy as np
import patsy
get_ipython().run_line_magic('pylab', 'inline')


# In[3]:


data = pd.read_csv('../lifelines/datasets/divorce.dat', sep="\s{2,10}")


# In[4]:


data.head(10)


# In[5]:


df = patsy.dmatrix('heduc + heblack + heblack:mixed + years + mixed + div -1 ', data, return_type='dataframe')


# In[6]:


df.head()


# In[130]:


aaf = AalenAdditiveFitter(fit_intercept=True, coef_penalizer=1.)


# In[131]:


df.columns


# In[132]:


timeline = np.linspace(0, 35, 1000)


# In[133]:


aaf.fit(df, 'years', event_col='div[T.Yes]', timeline=timeline)


# In[134]:


figsize(12.5, 8.5)
#aaf.cumulative_hazards_.ix[:][["baseline"]].plot()
#aaf.cumulative_hazards_.ix[:][["bb","bw","wb","ww"]].plot()
aaf.cumulative_hazards_.plot()
plt.legend(loc='upper left')


# In[ ]:





# In[ ]:





# In[ ]: