Getting survival function estimates group by attribute level in Lifelines

Question

I have a challenge with using Lifelines for KM estimates. I have a variable column called worker type (Full Time, Part Time, etc) that I would like to group the KM estimates for, then output to a CSV file. Here's a snippet:

worker_types = df['Emp_Status'].unique() 

for i, worker_type in enumerate(worker_types): 
    ix = df['Emp_Status'] == worker_type 
    kmf.fit(T[ix], C[ix]) 
    kmf.survival_function_['worker'] = worker_type 
    #print kmf.survival_function_ 
    kmf.surviva

l_function_.to_csv('C:\Users\Downloads\test.csv')

When I use the print function, I get each iteration of the KM estimate per worker_type; however, when trying to export to a csv file, I only get the last estimate of worker type.

I've read the lifelines docs, and seen the examples for the plotting of different levels, but not sure how to bridge that to exporting to csv.

You should provide links to docs and DYA (define your acronyms). — tmthydvnprt, Jul 26 '16 at 14:14

mechanical_meat · Accepted Answer · 2016-07-26T14:14:03.453

0

You can open the file in append mode at the top of the loop and then append each row, e.g.:

worker_types = df['Emp_Status'].unique() 
with open('C:/Users/Downloads/test.csv', 'a') as fou:
    for i, worker_type in enumerate(worker_types): 
        ix = df['Emp_Status'] == worker_type 
        kmf.fit(T[ix], C[ix]) 
        kmf.survival_function_['worker'] = worker_type 
        if i == 0:
            kmf.survival_function_.to_csv(fou) # write header on first iteration
        else:
            kmf.survival_function_.to_csv(fou, header=False)

Side note: Please do not use backwards slashes for Windows paths within Python. Instead use forward slashes.

edited Jul 26 '16 at 14:14

answered Jul 26 '16 at 14:08

mechanical_meat

163,903
24
228
223

Thank you for the assistance Bernie, that did the trick. My apologies for the sloppy post; it was my first one. – TBoneATL Jul 26 '16 at 19:14
You're very welcome, sir. No worries about the post it was fine. – mechanical_meat Jul 26 '16 at 19:16
Hi Bernie, can I ask a follow up question? I've also tried to append each iteration of the survival_function_ df to another dataframe, but the result dataframe only has the very last iteration of the Emp Status. Can you help point out what I'm doing wrong? `data = pd.DataFrame() worker_types = df['Emp_Status'].unique() for i, worker_type in enumerate(worker_types): ix = df['Emp_Status'] == worker_type kmf.fit(T[ix], C[ix]) kmf.survival_function_['worker'] = worker_type data.append(kmf.survival_function_)` – TBoneATL Jul 26 '16 at 21:13
@Travis1585: you need to ensure that the `.append()` operation is occurring in the `for` loop. – mechanical_meat Jul 26 '16 at 21:41

Getting survival function estimates group by attribute level in Lifelines

1 Answers1