python.py

"""Module 20 - competing risks and time-varying Cox models.

Reads the coherent competing-risks dataset data/konkurrenz_risiken.csv
(generated by data/generate_data.py with its OWN RNG streams, so it changes no
other course number):

    patient_id
    event_time             day of the exit event (death day, discharge day, or 30)
    event_state            0 = censored (still in hospital at day 30, or lost
                               to follow-up earlier)
                           1 = in-hospital death (event of interest)
                           2 = discharge alive (competing event)
    vasopressor_start_tag  day a time-varying vasopressor exposure begins, or
                           empty if the patient is never exposed in their window

Discharge is a GENUINE competing event on the same day axis as death, drawn
with a hazard that FALLS with severity (higher sofa_score -> later, or no,
discharge), and administrative censoring at day 30 is preserved (the day-30
risk set stays non-degenerate at ~90 patients). In-hospital deaths are kept
exactly as observed in the cohort (78 of 500 = 15.6 %); the Aalen-Johansen CIF
is slightly higher (16.0 %) because 18 patients are lost to follow-up before
day 30. The cause-specific Cox hazards land close to Module 17.

Baseline covariates (alter, sofa_score, diabetes) come from kohorte.csv. No
files under data/ are edited by this script.

Run from project root:
    python module/20-competing-risks-timevariant/code/python.py
"""
from __future__ import annotations

import sys
import warnings
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT))

import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
from lifelines import CoxPHFitter, CoxTimeVaryingFitter, KaplanMeierFitter  # noqa: E402

from lib.helpers import _load, load_cohort  # noqa: E402

ADMIN_HORIZON = 30  # administrative censoring day


def build_competing_risk_data() -> pd.DataFrame:
    """Load the 3-state competing-risks outcome and attach baseline covariates.

    The three states already live in konkurrenz_risiken.csv (a properly
    generated file); here we only join the baseline covariates the models need.
    """
    cr = _load("konkurrenz_risiken.csv")
    cov = load_cohort()[["patient_id", "alter", "sofa_score", "diabetes"]]
    df = cr.merge(cov, on="patient_id")
    df["event_time"] = df["event_time"].astype(float)
    return df


def aalen_johansen_cif(time: pd.Series, state: pd.Series, cause: int = 1) -> pd.Series:
    """Aalen-Johansen cumulative incidence of `cause` at each distinct time.

    state: 0 = censored, any other value = an event of that type. Tied days are
    handled exactly. lifelines' AalenJohansenFitter instead jitters tied times
    at random; with day-resolution data (many ties) that shifts the day-30
    estimate (15.6 % instead of the exact 16.0 % here, seed-dependent).
    """
    cif, surv, out = 0.0, 1.0, {}
    for t in np.sort(time.unique()):
        at_risk = (time >= t).sum()
        d_cause = ((time == t) & (state == cause)).sum()
        d_any = ((time == t) & (state != 0)).sum()
        cif += surv * d_cause / at_risk   # S(t-) * cause-specific hazard at t
        surv *= 1 - d_any / at_risk       # all-cause survival (any event counts)
        out[t] = cif
    return pd.Series(out, name=f"cif_{cause}")


def build_time_varying_data(df: pd.DataFrame) -> pd.DataFrame:
    """Expand the shared vasopressor exposure into the start-stop (long) format.

    Uses `vasopressor_start_tag` from konkurrenz_risiken.csv, so Python and R
    build byte-identical rows (they previously drew from different RNG streams
    and disagreed). A patient exposed at day `s < event_time` gets two
    contiguous intervals [0, s) unexposed and [s, event_time) exposed, with the
    death indicator only in the LAST interval. Everyone else gets one interval.
    """
    rows = []
    for _, r in df.iterrows():
        pid = int(r["patient_id"])
        t_end = float(r["event_time"])
        death = int(r["event_state"] == 1)
        s = float(r["sofa_score"])
        start = r["vasopressor_start_tag"]
        if pd.notna(start) and start < t_end:
            rows.append(dict(patient_id=pid, start=0.0, stop=float(start),
                             vasopressor=0, death=0, sofa=s))
            rows.append(dict(patient_id=pid, start=float(start), stop=t_end,
                             vasopressor=1, death=death, sofa=s))
        else:
            rows.append(dict(patient_id=pid, start=0.0, stop=t_end, vasopressor=0,
                             death=death, sofa=s))
    tv = pd.DataFrame(rows)
    return tv[tv["stop"] > tv["start"]]


def main() -> None:
    df = build_competing_risk_data()

    print("=== 3-state competing-risks outcome (data/konkurrenz_risiken.csv) ===")
    print(df["event_state"].value_counts().rename({0: "censored", 1: "death", 2: "discharged"}))
    print(df[["patient_id", "event_time", "event_state", "vasopressor_start_tag"]].head())

    # --- 1) Cumulative Incidence Function (CIF) vs. naive 1-KM ---------------
    print("\n=== 1) CIF (correct) vs. naive 1-KM (treats discharge as censoring) ===")
    cif = aalen_johansen_cif(df["event_time"], df["event_state"], cause=1)
    cif_30 = float(cif.loc[:ADMIN_HORIZON].iloc[-1])

    kmf = KaplanMeierFitter()
    kmf.fit(df["event_time"], event_observed=(df["event_state"] == 1))
    naive_30 = float(1 - kmf.survival_function_.iloc[-1, 0])

    print(f"CIF for death at day {ADMIN_HORIZON} (competing risks respected): {cif_30:.1%}")
    print(f"Naive 1-KM 'risk of death' at the same horizon (discharge as censoring): {naive_30:.1%}")
    print("The naive number OVERSHOOTS the truth - visibly, but bounded - because discharged")
    print("patients drop out of the risk set, leaving a smaller denominator whose deaths get")
    print("extrapolated to the whole population. This is the realistic textbook bias, not an")
    print("artefact: the day-30 risk set is still ~90 patients, so 1-KM does not blow up to 100 %.")

    # --- 2) Cause-specific Cox model (competing event = ordinary censoring) --
    print("\n=== 2) Cause-specific Cox model (death only; discharge treated as censoring) ===")
    cs_df = df[["event_time", "event_state", "alter", "sofa_score", "diabetes"]].copy()
    cs_df["death"] = (cs_df["event_state"] == 1).astype(int)
    cs_cox = CoxPHFitter()
    cs_cox.fit(cs_df[["event_time", "death", "alter", "sofa_score", "diabetes"]],
              duration_col="event_time", event_col="death")
    print(cs_cox.summary[["coef", "exp(coef)", "se(coef)", "p"]].round(4))
    print("Interpretation: the instantaneous death rate among patients who have not yet died OR")
    print("been discharged - good for etiology, not for population-level absolute risk.")
    print("\nNote on Fine-Gray (subdistribution hazard): no maintained Python package implements it")
    print("(lifelines/scikit-survival don't). See code/r.R, which uses R's built-in")
    print("survival::finegray() + a weighted coxph() - the standard, well-tested route. There the")
    print("SOFA subdistribution HR is LARGER than the cause-specific HR, because higher-SOFA")
    print("patients are also less likely to be discharged (the competing event), a channel that")
    print("cause-specific hazards ignore by design.")
    print("\nPower caveat: with only", int(cs_df['death'].sum()), "deaths total, splitting events")
    print("across two causes leaves few events per cause - single-covariate p-values above are")
    print("interpretable, but a richly adjusted competing-risks model would be underpowered.")

    # --- 3) Immortal time bias: naive time-fixed vs. correct time-varying ----
    print("\n=== 3) Time-varying Cox: vasopressor exposure (start-stop format) ===")
    df["death"] = (df["event_state"] == 1).astype(int)
    df["ever_vasopressor"] = df["vasopressor_start_tag"].notna().astype(int)

    # NAIVE: classify the exposure as a fixed baseline covariate (WRONG - this
    # credits treated patients with the "immortal" time before their exposure).
    naive_cox = CoxPHFitter()
    naive_cox.fit(df[["event_time", "death", "ever_vasopressor", "sofa_score"]],
                  duration_col="event_time", event_col="death")
    hr_naive = float(np.exp(naive_cox.params_["ever_vasopressor"]))
    print(f"NAIVE time-fixed HR for 'ever vasopressor' (immortal time bias): {hr_naive:.2f} "
          f"(p={naive_cox.summary.loc['ever_vasopressor', 'p']:.3f})")
    print("-> a spurious PROTECTIVE effect: treated patients had to survive long enough to be")
    print("   treated, so their early deaths are misattributed to the untreated group.")

    # CORRECT: time-varying exposure, switched on only from the start day.
    tv = build_time_varying_data(df)
    print(f"\nStart-stop rows: {len(tv)} | patients: {tv['patient_id'].nunique()} | "
          f"exposed: {int(df['ever_vasopressor'].sum())}")
    print(pd.crosstab(tv["vasopressor"], tv["death"], rownames=["vasopressor"], colnames=["death"]))

    ctv = CoxTimeVaryingFitter()
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        ctv.fit(tv, id_col="patient_id", start_col="start", stop_col="stop",
               event_col="death", show_progress=False)
    if caught:
        print("Warnings:", [str(w.message) for w in caught])
    else:
        print("Converged without warnings (contrast with a 4-patient toy example, which wouldn't).")
    print(ctv.summary[["coef", "exp(coef)", "se(coef)", "p"]].round(4))
    hr_tv = float(np.exp(ctv.params_["vasopressor"]))
    print(f"\nCORRECT time-varying HR for vasopressor: {hr_tv:.2f} - the spurious protective effect")
    print("collapses toward 1 once the immortal pre-exposure time is credited to the UNEXPOSED")
    print("state. The generating process gives vasopressor no causal effect on death (it only")
    print("marks severity), so ~1.0 is the honest answer; the naive model manufactured a benefit.")


if __name__ == "__main__":
    main()