python.py

"""Module 12 — Regression models: OLS, logistic regression, and Cox regression.

Runs standalone from the project root:
    python module/12-regression/code/python.py

Data: read from data/ (committed with the repo); if that folder is
missing, the same files are fetched from the published URL.
Packages: statsmodels, lifelines  (pip install statsmodels lifelines)
"""
from __future__ import annotations

import sys
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
import statsmodels.formula.api as smf  # noqa: E402
from lifelines import CoxPHFitter  # noqa: E402

from lib.ground_truth import (  # noqa: E402
    true_hazard_ratios,
    true_odds_ratios,
    true_odds_ratios_for,
)
from lib.helpers import SEED, load_cohort  # noqa: E402

pd.set_option("display.width", 100)
pd.set_option("display.float_format", "{:.4f}".format)


# ---------------------------------------------------------------------------
# Feature engineering
# ---------------------------------------------------------------------------

def prepare_features(df: pd.DataFrame) -> pd.DataFrame:
    """Derive binary indicators and centre age as in the data generator.

    Column names that stay German (clinical schema):
        raucherstatus, aufnahmegrund, alter, verweildauer_tage, verstorben_30d,
        sofa_score, crp_mg_l, diabetes, hypertonie
    New derived columns use English names to follow code conventions.
    """
    df = df.copy()
    df["active_smoker"] = (df["raucherstatus"] == "aktiv").astype(int)
    df["sepsis"]        = (df["aufnahmegrund"] == "Sepsis").astype(int)
    # Centre age on 64 — matches the synthetic data generator intercept
    df["age_centred"]   = df["alter"] - 64
    return df


# ---------------------------------------------------------------------------
# Section 1 — OLS: length of stay
# ---------------------------------------------------------------------------

def fit_ols(df: pd.DataFrame) -> None:
    """Ordinary Least Squares for a continuous outcome (length of stay).

    Note: verweildauer_tage is right-skewed; OLS is illustrative here.
    A Gamma-GLM would be more appropriate for production use.
    """
    print("=" * 65)
    print("SECTION 1 — LINEAR REGRESSION: length of stay (days)")
    print("=" * 65)
    print("Note: length of stay is right-skewed; OLS is illustrative.")
    print("A Gamma-GLM or log-transform would be more appropriate.\n")

    formula = "verweildauer_tage ~ sofa_score + age_centred + sepsis + diabetes"
    result = smf.ols(formula, data=df).fit()

    print(f"R²: {result.rsquared:.4f}   Adjusted R²: {result.rsquared_adj:.4f}\n")

    coef_table = pd.DataFrame({
        "β̂":            result.params,
        "SE":            result.bse,
        "95% CI lower":  result.conf_int()[0],
        "95% CI upper":  result.conf_int()[1],
        "p":             result.pvalues,
    })
    print("Coefficients:")
    print(coef_table.to_string())

    sofa_coef = result.params["sofa_score"]
    print(
        f"\nInterpretation: each additional SOFA point is associated with"
        f" {sofa_coef:.2f} extra days of stay, adjusting for the other predictors."
    )

    # The sepsis/diabetes rows above are DIRECT effects — they hold sofa_score
    # fixed, and sofa_score is a mediator. Reading "sepsis: p = 0.42" as "sepsis
    # does not prolong the stay" is exactly the overadjustment error §5 warns
    # about. Show the total effect next to it instead of leaving the trap open.
    total = smf.ols("verweildauer_tage ~ age_centred + sepsis + diabetes", data=df).fit()
    print("\nDirect vs. total effect (sofa_score is a MEDIATOR, see README §5):")
    print(f"{'':12s}{'direct (sofa in model)':>26}{'total (sofa omitted)':>24}")
    for term in ("sepsis", "diabetes"):
        print(f"{term:12s}"
              f"{f'{result.params[term]:+.2f} d (p={result.pvalues[term]:.3f})':>26}"
              f"{f'{total.params[term]:+.2f} d (p={total.pvalues[term]:.3f})':>24}")
    print("Sepsis lengthens the stay entirely THROUGH the SOFA score. Adjusting")
    print("for SOFA removes the effect you were asking about.")


# ---------------------------------------------------------------------------
# Section 2 — Logistic regression: 30-day mortality
# ---------------------------------------------------------------------------

def fit_logistic(df: pd.DataFrame) -> None:
    """Logistic regression for a binary outcome (30-day mortality).

    Reports:
    - Estimated log-odds coefficients (beta_hat)
    - Odds ratios with 95 % confidence intervals
    - Truth comparison against the known generative betas
    - Model fit statistics (pseudo-R², AIC, EPV)
    """
    print("\n" + "=" * 65)
    print("SECTION 2 — LOGISTIC REGRESSION: 30-day mortality")
    print("=" * 65)

    formula = (
        "verstorben_30d ~ age_centred + sofa_score + crp_mg_l"
        " + diabetes + sepsis + active_smoker"
    )
    result = smf.logit(formula, data=df).fit(disp=False)

    # Odds ratios and 95 % CI on the exponentiated scale
    ci = result.conf_int()
    or_table = pd.DataFrame({
        "Log-odds (β̂)":   result.params,
        "OR":              np.exp(result.params),
        "OR 95% CI lower": np.exp(ci[0]),
        "OR 95% CI upper": np.exp(ci[1]),
        "p":               result.pvalues,
    })
    print("\nEstimated coefficients and odds ratios:")
    print(or_table.to_string())

    # ------------------------------------------------------------------
    # Truth comparison — against the true ODDS ratios.
    #
    # The generative model in data/generate_data.py is a Weibull proportional-
    # hazards process, so its coefficients are log-HAZARDS. Exponentiating them
    # gives true hazard ratios, not the odds ratios this logistic model
    # estimates. lib/ground_truth.py obtains the true ORs by replaying the same
    # data-generating process at N = 2 000 000 and fitting this exact model.
    # ------------------------------------------------------------------
    true_or = true_odds_ratios()
    true_hr = true_hazard_ratios()["direct"]

    print("\n--- Truth comparison: estimated OR vs. true OR ---")
    print(f"{'Predictor':<16}{'OR_hat':>10}{'95% CI':>20}{'true OR':>10}"
          f"{'true HR':>10}{'covered?':>10}")
    print("-" * 76)
    covered = 0
    for predictor, or_true in true_or.items():
        est = float(np.exp(result.params[predictor]))
        lo, hi = float(np.exp(ci.loc[predictor, 0])), float(np.exp(ci.loc[predictor, 1]))
        hit = lo <= or_true <= hi
        covered += hit
        print(f"{predictor:<16}{est:>10.3f}{f'[{lo:.3f}, {hi:.3f}]':>20}"
              f"{or_true:>10.3f}{true_hr[predictor]:>10.3f}{'yes' if hit else 'NO':>10}")

    print(
        f"\nThe 95 % CI covers the true OR for {covered}/{len(true_or)} predictors."
        "\nRemaining deviations are sampling variance (N=500, 78 events), not bias."
        "\n"
        "\nWhy two 'true' columns? The cohort is generated by a Weibull hazard"
        "\nmodel, so exp(beta) is a true HAZARD ratio (compare it in Module 17's"
        "\nCox model). A logistic model estimates an ODDS ratio, which at a 15.6 %"
        "\nevent rate sits systematically further from 1. Comparing an estimated OR"
        "\nagainst a true HR would make an unbiased estimate look biased."
    )

    # Model fit
    n_events = int(df["verstorben_30d"].sum())
    n_pred = len(true_or)                  # intercept is not a predictor
    epv = n_events / n_pred
    print(f"\nPseudo-R² (McFadden): {result.prsquared:.4f}")
    print(f"Log-likelihood:       {result.llf:.2f}")
    print(f"AIC:                  {result.aic:.2f}")
    print(f"\nEvents per variable (EPV): {n_events} events / {n_pred} predictors = {epv:.1f}")
    if epv < 10:
        print("  WARNING: EPV < 10 — model may be over-parametrised.")
    else:
        print("  EPV >= 10 — model complexity is defensible.")


# ---------------------------------------------------------------------------
# Section 3 — Cox regression: time-to-event analysis
# ---------------------------------------------------------------------------

def confounder_vs_mediator(df: pd.DataFrame) -> None:
    """Crude -> +confounder -> +mediator, for sepsis and for diabetes.

    README §4 and the exercises quote these numbers. Printing them here means no
    number in the lesson exists only on paper: the reader can reproduce every
    cell, and a future edit that breaks one is caught by tools/check_numbers.py.
    """
    print("\n" + "=" * 65)
    print("SECTION 2b — CONFOUNDER vs. MEDIATOR (README §4, Aufgabe 3)")
    print("=" * 65)
    print("alter is a confounder (adjust). sofa_score is a MEDIATOR (do not adjust,")
    print("if the total effect is the question).")
    print("True ORs replayed at N = 400 000 (Monte-Carlo error ~1 %, far below the")
    print("width of the N=500 confidence intervals they sit next to).\n")

    # Never hardcode a truth: recompute it from the generating process, at a
    # cheaper N because these auxiliary values are compared against N=500 CIs
    # that are orders of magnitude wider than the Monte-Carlo error.
    spezifikationen = [
        ("crude", ()),
        ("+ alter", ("age_centred",)),
        ("+ alter + sofa", ("age_centred", "sofa_score")),
    ]

    for exposure in ("sepsis", "diabetes"):
        print(f"  {exposure}")
        print(f"    {'model':18s}{'OR_hat':>9}{'95% CI':>20}{'true OR':>10}")
        for label, extra in spezifikationen:
            terms = " + ".join((exposure,) + extra)
            fit = smf.logit(f"verstorben_30d ~ {terms}", data=df).fit(disp=False)
            est = float(np.exp(fit.params[exposure]))
            lo, hi = np.exp(fit.conf_int().loc[exposure])
            wahr = true_odds_ratios_for((exposure,) + extra, n=400_000)[exposure]
            print(f"    {label:18s}{est:>9.2f}{f'[{lo:.2f}, {hi:.2f}]':>20}{wahr:>10.3f}")
        print()

    print("Sepsis: 5.75 -> 1.92 when the mediator enters, and the CIs barely overlap —")
    print("the overadjustment bias is demonstrable. Diabetes shows the same pattern,")
    print("but its CI [0.52, 1.87] contains both the true total and direct effects:")
    print("at 78 events the bias can be derived, not shown.")


def fit_cox(df: pd.DataFrame) -> None:
    """Cox proportional-hazards regression for time-to-event data.

    Uses fu_zeit_tage (follow-up time until event/censoring) as the time
    variable and status (1 = died, 0 = censored) as the event indicator.
    verweildauer_tage (length of stay) is NOT a time-to-event variable, it
    is generated independently of the survival process — see data/README.md.
    """
    print("\n" + "=" * 65)
    print("SECTION 3 — COX REGRESSION: time-to-event analysis")
    print("=" * 65)
    print(
        "Duration: fu_zeit_tage | Event: status (1 = died, 0 = censored)"
        "\nPatients without the event are right-censored (not a data loss).\n"
    )

    cox_cols = [
        "fu_zeit_tage", "status",
        "sofa_score", "age_centred", "sepsis", "active_smoker",
    ]
    cox_df = df[cox_cols].copy()

    cph = CoxPHFitter()
    cph.fit(
        cox_df,
        duration_col="fu_zeit_tage",
        event_col="status",
    )

    summary_cols = ["exp(coef)", "exp(coef) lower 95%", "exp(coef) upper 95%", "p"]
    hr_table = cph.summary[summary_cols].copy()
    hr_table.columns = ["HR", "HR 95% CI lower", "HR 95% CI upper", "p"]
    print("Hazard ratios (HR) with 95 % CI:")
    print(hr_table.to_string())

    print(
        "\nHR interpretation:"
        "\n  HR > 1 → elevated hazard (event occurs sooner / more often)"
        "\n  HR < 1 → reduced hazard (protective)"
        "\n  HR = 1 → no difference from reference"
        "\n\nNote: the proportional-hazards assumption should be checked with"
        "\nSchoenfeld residuals — use cph.check_assumptions(cox_df) in lifelines."
    )


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    df_raw = load_cohort()
    df = prepare_features(df_raw)

    mortality_rate = df["verstorben_30d"].mean()
    print(f"Cohort: {len(df)} patients | 30-day mortality: {mortality_rate:.1%}")
    print(
        "This dataset has a KNOWN ground truth: a Weibull proportional-hazards"
        "\nprocess (data/generate_data.py). Its coefficients are log-HAZARDS, so"
        "\nwe compare the Cox model against them directly, and the logistic model"
        "\nagainst the true ODDS ratios from lib/ground_truth.py."
    )

    fit_ols(df)
    fit_logistic(df)
    confounder_vs_mediator(df)
    fit_cox(df)

    print("\nDone.")


if __name__ == "__main__":
    main()