EDA CookBook - 生存分析とKaplanMeier推定量

1
2
3
4
5
6
7
import pandas as pd
from sksurv.datasets import load_veterans_lung_cancer

df_x, array_y  = load_veterans_lung_cancer()
df_y = pd.DataFrame(array_y)
df = pd.concat([df_y, df_x], axis=1)
df.head()

1
2
3
4
5
6
7
cat_cols = df.select_dtypes(exclude=np.float64).columns.tolist()
(pd.DataFrame(
            df[cat_cols].melt(var_name='column', 
                              value_name='value'
                             ).value_counts())
            .rename(columns={0: 'counts'})
            .sort_values(by=['column', 'counts']))

1
2
3
4
5
6
7
8
9
10
11
is_therapy = df['Prior_therapy'] == 'yes'
event = df['Status'].values
time = df['Survival_in_days'].values
x_total, y_total = kaplan_meier_estimator(event, time)
x_treated, y_treated = kaplan_meier_estimator(event[is_therapy], time[is_therapy])
x_control, y_control = kaplan_meier_estimator(event[~is_therapy], time[~is_therapy])

plt.step(x_treated, y_treated, where="post", label='treated')
plt.step(x_control, y_control, where="post", label='controled')
plt.ylim(0, 1)
plt.legend();

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
## x_total, y_totalの再現のみ
km_idx, km_freq = np.unique(time, return_counts=True)
km_freq = np.cumsum(km_freq[::-1])[::-1]
death_idx, death_freq = np.unique(time[event], return_counts=True)

df_left = pd.DataFrame(np.array([km_idx, km_freq],dtype=np.int64).T,columns=['time', 'obs'])
df_right = pd.DataFrame(np.array([death_idx, death_freq],dtype=np.int64).T,columns=['time', 'death'])
df_merge = (pd.merge(df_left, df_right, how='left', on='time')).fillna(0)
df_merge['ratio'] = np.cumprod(1 - df_merge['death']/df_merge['obs'])

## 95% confidence intervalの計算
df_merge['std_err'] = np.sqrt(np.cumsum(df_merge['death']/(df_merge['obs'] * (df_merge['obs'] - df_merge['death']))))
df_merge['CI_lower'] = np.clip(df_merge['ratio'] * np.exp(-1.96 * df_merge['std_err']), 0, 1)
df_merge['CI_upper'] = np.clip(df_merge['ratio'] * np.exp(1.96 * df_merge['std_err']), 0, 1)

## plotの再現
plt.step(df_merge['time'].values, df_merge['ratio'].values,where="post", label='survival')
plt.fill_between(df_merge['time'].values, df_merge['CI_lower'].values,df_merge['CI_upper'].values, color = 'b', alpha = 0.2, step="post")


plt.legend();

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
from scipy import stats
import numpy as np
import pandas as pd
from collections import OrderedDict

def survival_logrank_test(duration_array, event_array, group_array):
    """K-sample log-rank hypothesis test of identical survival functions.
    Compares the pooled hazard rate with each group-specific
    hazard rate. The alternative hypothesis is that the hazard
    rate of at least one group differs from the others at some time.

    Parameters
    ----------
    duration_array : time of event or time of censoring array, shape = (n_samples,)
    
    event_array : the binary event indicator array
                  1: event observed
                  0: censored
    
    group_array : array-like, shape = (n_samples,)
                  Group membership of each sample.
    
    Returns
    -------
    chisq : float
        Test statistic.
    pvalue : float
        Two-sided p-value with respect to the null hypothesis
        that the hazard rates across all groups are equal.
    stats : pandas.DataFrame
        Summary statistics for each group:  number of samples,
        observed number of events, expected number of events,
        and test statistic.
        Only provided if `return_stats` is True.
    covariance : array, shape=(n_groups, n_groups)
        Covariance matrix of the test statistic.
        Only provided if `return_stats` is True.
    
    References
    ----------
    .. [1] Fleming, T. R. and Harrington, D. P.
           A Class of Hypothesis Tests for One and Two Samples of Censored Survival Data.
           Communications In Statistics 10 (1981): 763-794.
    """

    n_samples = len(duration_array)
    groups, group_counts = np.unique(group_array, return_counts=True)
    n_groups = groups.shape[0]

    # sort descending
    o = np.argsort(-duration_array, kind="mergesort")
    x = group_array[o]
    event = event_array[o]
    time = duration_array[o]

    at_risk = np.zeros(n_groups, dtype=int)
    observed = np.zeros(n_groups, dtype=int)
    expected = np.zeros(n_groups, dtype=float)
    covar = np.zeros((n_groups, n_groups), dtype=float)
    covar_indices = np.diag_indices(n_groups)

    k = 0
    while k < n_samples:
        ti = time[k]
        total_events = 0
        while k < n_samples and ti == time[k]:
            idx = np.searchsorted(groups, x[k])
            if event[k]:
                observed[idx] += 1
                total_events += 1
            at_risk[idx] += 1
            k += 1
        if total_events != 0:
            total_at_risk = k
            expected += at_risk * (total_events / total_at_risk)
            if total_at_risk > 1:
                multiplier = total_events * (total_at_risk - total_events) / (total_at_risk * (total_at_risk - 1))
                temp = at_risk * multiplier
                covar[covar_indices] += temp
                covar -= np.outer(temp, at_risk) / total_at_risk
    
    chi_df = n_groups - 1
    numerator = observed[:chi_df] - expected[:chi_df]
    
    chisq = np.linalg.solve(covar[:chi_df, :chi_df], numerator) @ numerator
    pval = stats.chi2.sf(chisq, chi_df)
    
    table = OrderedDict()
    table["counts"] = group_counts
    table["observed"] = observed
    table["expected"] = expected
    table["statistic"] = observed - expected
    table = pd.DataFrame.from_dict(table)
    table.index = pd.Index(groups, name="group", dtype=groups.dtype)
    
    return chisq, pval, table, covar, chi_df


# numpy example
G = np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2])
T = np.array([5, 3, 9, 8, 7, 4, 4, 3, 2, 5, 6, 7])
E = np.array([1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 0])
survival_logrank_test(duration_array=T, event_array=E, group_array=G)