📉 留存/生存分析

import numpy as np
from scipy import stats

def kaplan_meier(times, events):
    """
    Kaplan-Meier 生存分析
    times: 观察时间数组
    events: 事件发生标记（1=发生，0=删失）
    """
    # 排序
    order = np.argsort(times)
    times = np.array(times)[order]
    events = np.array(events)[order]

    unique_times = np.unique(times)
    survival = []
    at_risk = []
    n_events = []

    n = len(times)
    s = 1.0

    for t in unique_times:
        mask = times == t
        n_at_risk = np.sum(times >= t)
        n_event = np.sum(events[mask])

        if n_at_risk > 0:
            s *= (1 - n_event / n_at_risk)

        survival.append(s)
        at_risk.append(n_at_risk)
        n_events.append(n_event)

    # 计算中位生存时间
    median_idx = np.where(np.array(survival) <= 0.5)[0]
    median_time = unique_times[median_idx[0]] if len(median_idx) > 0 else None

    return {
        'times': unique_times,
        'survival': survival,
        'at_risk': at_risk,
        'events': n_events,
        'median': median_time,
        'final_survival': survival[-1]
    }