risikogruppen_km.png

def fig_risikogruppen_km(risk_scores, events_test, times_test) -> None:
    """Kaplan-Meier curves stratified by model-predicted risk tertile."""
    # Split by tertiles of the risk score (not by outcome — that would be leakage).
    tertile_low = np.percentile(risk_scores, 33)
    tertile_high = np.percentile(risk_scores, 67)

    masks = {
        "Niedrig (< 33. Perz.)": risk_scores < tertile_low,
        "Mittel (33.–67. Perz.)": (risk_scores >= tertile_low) & (risk_scores < tertile_high),
        "Hoch (> 67. Perz.)": risk_scores >= tertile_high,
    }
    colors = [PRIMARY, SECONDARY, EVENT]

    fig, ax = plt.subplots(figsize=(7, 4))
    for (label, mask), color in zip(masks.items(), colors):
        if mask.sum() < 5:
            continue
        kmf = KaplanMeierFitter()
        kmf.fit(
            durations=times_test[mask],
            event_observed=events_test[mask],
            label=label,
        )
        kmf.plot_survival_function(ax=ax, color=color, ci_show=True)

    ax.set_xlabel("Zeit (Tage)")
    ax.set_ylabel("Überlebenswahrscheinlichkeit")
    ax.set_title("Kaplan-Meier nach Risikogruppe (Modell-Tertile)")
    ax.set_ylim(0, 1.05)
    ax.legend(loc="lower left")
    save(fig, ASSETS / "risikogruppen_km.png")