Format data for survival analysis using pandas

Question

I'm trying to figure out the quickest way to get survival analysis data into a format that will allow for time varying covariates. Basically this would be a python implementation of stsplit in Stata. To give a simple example, with the following set of information:

id start end  x1  x2  exit
 1   0    18  12  11   1

This tells us that an observation started at time 0, and ended at time 18. Exit tells us that this was a 'death' rather than right censoring. x1 and x2 are variables that are constant over time.

id   t   age
 1   0    30
 1   7    40
 1   17   50

I'd like to get:

id start end  x1  x2  exit  age
 1   0    7   12  11   0    30
 1   7    17  12  11   0    40
 1   17   18  12  11   1    50

Exit is only 1 at the end, signifying that t=18 is when the death occurred.

t is subsumed in the changes to start and end. Essentially start is now t. — Luke, Oct 22 '14 at 23:21

score 1 · Answer 1 · answered Oct 23 '14 at 01:29

Assuming:

>>> df1
id  start   end x1  x2  exit
0   1   0   18  12  11  1

and:

>>> df2
   id   t   age
0   1   0   30
1   1   7   40
2   1   17  50

You can do:

df = df2.copy()                                 # start with df2
df['x1'] = df1.ix[0, 'x1']                      # x1 column
df['x2'] = df1.ix[0, 'x2']                      # x2 column
df.rename(columns={'t': 'start'}, inplace=True) # start column
df['end'] = df['start'].shift(-1)               # end column
df.ix[len(df)-1, 'end'] = df1.ix[0, 'end']
df['exit'] = 0                                  # exit column
df.ix[len(df)-1, 'exit'] = 1                     
df = df[['id', 'start', 'end', 'x1', 'x2', 'exit', 'age']] # reorder columns

Output:

>>> df
    id  start   end x1  x2  exit    age
0   1   0       7   12  11  0       30
1   1   7       17  12  11  0       40
2   1   17      18  12  11  1       50

score 0 · Answer 2 · answered Apr 23 '19 at 19:04

This is possible using lifelines, specifically, the add_covariate_to_timeline function, example here. The function is quite flexible and can doing things like cumulative sums, etc.

For the example above:


"""
id start end  x1  x2  exit
 1   0    18  12  11   1
"""
long_df = pd.DataFrame([
    {'id': 1, 'start': 0, 'end': 18, 'x1': 12, 'x2': 11, 'exit': 1}
])


"""
id   t   age
 1   0    30
 1   7    40
 1   17   50
"""
tv_covariates = pd.DataFrame([
    {'id': 1, 't': 0, 'age': 30},
    {'id': 1, 't': 7, 'age': 40},
    {'id': 1, 't': 17, 'age': 50},
])

from lifelines.utils import add_covariate_to_timeline

add_covariate_to_timeline(long_df, tv_covariates, id_col='id', duration_col='t', event_col='exit', start_col='start', stop_col='end')

"""
   start  age    x1    x2   end  id   exit
0      0   30  12.0  11.0   7.0   1  False
1      7   40  12.0  11.0  17.0   1  False
2     17   50  12.0  11.0  18.0   1   True

"""

Format data for survival analysis using pandas

2 Answers2