<jupyter_start><jupyter_text>Section 4. from "Piecewise exponential models for survival data with covariates" by <NAME><jupyter_code>from lifelines.datasets import load_lupus

df = load_lupus()

# preprocessing
T_col = 'time_between_diagnosis_and_last_observation_(years)'
E_col = 'dead'
df['time_elapsed_between_estimated_onset_and_diagnosis_binary'] = df['time_elapsed_between_estimated_onset_and_diagnosis_(months)'] <= 2*12
df['recent'] = df['year_month_of_diagnosis'] < '1951-07'

columns = ['is_male', 'is_white', 
           'age_at_diagnosis', 
           'time_elapsed_between_estimated_onset_and_diagnosis_binary',
           'recent'] + [T_col, E_col]

df = df[columns]
df = df.dropna() # drop the individual with NaN

# these models can naturally handle 0 durations, so we fudge a bit.
df.loc[df[T_col] == 0, T_col] = 0.000001
# add a constant column (only needed for non-Cox models.)
df['constant'] = 1.

from lifelines import PiecewiseExponentialRegressionFitter, CoxPHFitter


breakpoints = [
    [0.5, 0.8, 1.1, 1.7, 2.5, 3.1], 
    [0.3, 0.8, 1.0, 2.0, 3.0],
    [0.4, 0.9, 1.5, 2.5],
    [0.3, 1.0, 2.0, 3.0],
    [0.4],
    [0.3]
]

results = dict()
for i, bp in enumerate(breakpoints, start=1):
    # by forcing the penalizer to be 1000, the coefs between periods are constrainted to be identical. 
    pcf = PiecewiseExponentialRegressionFitter(penalizer=1000., breakpoints=bp)
    pcf.fit(df, T_col, E_col)
    # Note the negative sign. We use a different parameterization than the paper.
    results["model %d" % i] = -pcf.params_['lambda_0_'].drop('constant')
    
cph = CoxPHFitter().fit(df.drop('constant', axis=1), T_col, E_col)
results['Cox'] = cph.params_
results = pd.DataFrame(results).T
results
pcf = PiecewiseExponentialRegressionFitter(penalizer=1000., breakpoints=breakpoints[0])
pcf.fit(df, T_col, E_col)
pcf.plot_covariate_groups(['is_male', 'is_white'], [[1, 1], [1,0], [0, 1], [0, 0]], figsize=(10, 6))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14358193810248604, 'fraction_numerical': 0.04515474378488077, 'mean_word_length': 4.699421965317919, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '167117', 'license': 'MIT', 'n_tokens_mistral': 819, 'n_tokens_neox': 741, 'n_words': 173, 'path': 'replications/Friedman_1982.ipynb', 'repo_name': 'CamDavidsonPilon/lifelines-replications', 'size': '1981'}