Add burden trajectory plotting script

plot_burden_index_trajectories.py

@@ -0,0 +1,258 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Iterable
+
+try:
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+except ModuleNotFoundError as exc:
+    raise ModuleNotFoundError(
+        "plot_burden_index_trajectories.py requires matplotlib. "
+        "Install it in the server environment before running this script."
+    ) from exc
+import numpy as np
+import pandas as pd
+
+
+REQUIRED_COLUMNS = {
+    "patient_id",
+    "sex",
+    "landmark_age",
+    "burden_type",
+    "burden_dimension",
+    "bi_historical",
+    "bi_future",
+    "bi_total",
+}
+
+
+def _sex_label(value: object) -> str:
+    text = str(value).strip().lower()
+    if text in {"0", "0.0", "female", "f", "woman"}:
+        return "female"
+    if text in {"1", "1.0", "male", "m", "man"}:
+        return "male"
+    return text or "unknown"
+
+
+def _load_bi_csv(path: Path) -> pd.DataFrame:
+    header = pd.read_csv(path, nrows=0)
+    missing = sorted(REQUIRED_COLUMNS - set(header.columns))
+    if missing:
+        raise ValueError(f"{path} is missing required columns: {missing}")
+
+    df = pd.read_csv(
+        path,
+        usecols=[
+            "patient_id",
+            "sex",
+            "landmark_age",
+            "burden_type",
+            "burden_dimension",
+            "bi_historical",
+            "bi_future",
+            "bi_total",
+        ],
+    )
+    df["sex_label"] = df["sex"].map(_sex_label)
+    df["landmark_age"] = pd.to_numeric(df["landmark_age"], errors="raise")
+    for col in ["bi_historical", "bi_future", "bi_total"]:
+        df[col] = pd.to_numeric(df[col], errors="raise")
+    return df
+
+
+def _sample_patients(
+    df: pd.DataFrame,
+    *,
+    n_per_sex: int,
+    seed: int,
+) -> dict[str, np.ndarray]:
+    rng = np.random.default_rng(int(seed))
+    samples: dict[str, np.ndarray] = {}
+    for sex_label in ["female", "male"]:
+        ids = np.asarray(sorted(df.loc[df["sex_label"] == sex_label, "patient_id"].unique()))
+        if ids.size == 0:
+            samples[sex_label] = np.asarray([], dtype=np.int64)
+            continue
+        take = min(int(n_per_sex), int(ids.size))
+        samples[sex_label] = np.asarray(rng.choice(ids, size=take, replace=False))
+    return samples
+
+
+def _aggregate_total(df: pd.DataFrame) -> pd.DataFrame:
+    return (
+        df.groupby(["patient_id", "sex_label", "burden_type", "landmark_age"], as_index=False)
+        [["bi_historical", "bi_future", "bi_total"]]
+        .sum()
+    )
+
+
+def _plot_selected_trajectories(
+    total_df: pd.DataFrame,
+    *,
+    burden_type: str,
+    selected: dict[str, np.ndarray],
+    output_dir: Path,
+) -> None:
+    sub = total_df[total_df["burden_type"] == burden_type].copy()
+    mean_df = (
+        sub.groupby(["sex_label", "landmark_age"], as_index=False)["bi_total"]
+        .mean()
+        .sort_values("landmark_age")
+    )
+
+    fig, ax = plt.subplots(figsize=(9.5, 5.5), dpi=160)
+    colors = {"female": "#b83280", "male": "#2563eb"}
+    for sex_label in ["female", "male"]:
+        patient_ids = selected.get(sex_label, np.asarray([], dtype=np.int64))
+        for pid in patient_ids:
+            one = sub[sub["patient_id"] == pid].sort_values("landmark_age")
+            if one.empty:
+                continue
+            ax.plot(
+                one["landmark_age"],
+                one["bi_total"],
+                color=colors.get(sex_label, "0.4"),
+                alpha=0.22,
+                linewidth=1.0,
+            )
+        mean_one = mean_df[mean_df["sex_label"] == sex_label]
+        if not mean_one.empty:
+            ax.plot(
+                mean_one["landmark_age"],
+                mean_one["bi_total"],
+                color=colors.get(sex_label, "0.4"),
+                linewidth=2.8,
+                label=f"{sex_label} mean",
+            )
+
+    ax.set_title(f"{burden_type}: sampled individual trajectories and sex-specific means")
+    ax.set_xlabel("Landmark age")
+    ax.set_ylabel("Total burden index")
+    ax.grid(True, alpha=0.25)
+    ax.legend(frameon=False)
+    fig.tight_layout()
+    fig.savefig(output_dir / f"{burden_type}_sampled_trajectories_by_sex.png")
+    plt.close(fig)
+
+
+def _top_dimensions(df: pd.DataFrame, *, burden_type: str, top_n: int) -> list[str]:
+    sub = df[df["burden_type"] == burden_type]
+    if sub.empty:
+        return []
+    ranked = (
+        sub.groupby("burden_dimension")["bi_total"]
+        .mean()
+        .sort_values(ascending=False)
+        .head(int(top_n))
+    )
+    return [str(x) for x in ranked.index.tolist()]
+
+
+def _plot_dimension_mean_panels(
+    df: pd.DataFrame,
+    *,
+    burden_type: str,
+    dimensions: Iterable[str],
+    output_dir: Path,
+) -> None:
+    dims = list(dimensions)
+    if not dims:
+        return
+
+    mean_df = (
+        df[(df["burden_type"] == burden_type) & (df["burden_dimension"].isin(dims))]
+        .groupby(["sex_label", "burden_dimension", "landmark_age"], as_index=False)["bi_total"]
+        .mean()
+        .sort_values("landmark_age")
+    )
+    n_cols = min(3, len(dims))
+    n_rows = int(np.ceil(len(dims) / n_cols))
+    fig, axes = plt.subplots(
+        n_rows,
+        n_cols,
+        figsize=(4.2 * n_cols, 3.2 * n_rows),
+        dpi=160,
+        squeeze=False,
+    )
+    colors = {"female": "#b83280", "male": "#2563eb"}
+    for ax, dim in zip(axes.ravel(), dims):
+        panel = mean_df[mean_df["burden_dimension"] == dim]
+        for sex_label in ["female", "male"]:
+            one = panel[panel["sex_label"] == sex_label]
+            if one.empty:
+                continue
+            ax.plot(
+                one["landmark_age"],
+                one["bi_total"],
+                color=colors.get(sex_label, "0.4"),
+                linewidth=2.0,
+                label=sex_label,
+            )
+        ax.set_title(str(dim), fontsize=9)
+        ax.set_xlabel("Age")
+        ax.set_ylabel("Mean BI")
+        ax.grid(True, alpha=0.25)
+
+    for ax in axes.ravel()[len(dims):]:
+        ax.axis("off")
+    handles, labels = axes.ravel()[0].get_legend_handles_labels()
+    if handles:
+        fig.legend(handles, labels, loc="upper right", frameon=False)
+    fig.suptitle(f"{burden_type}: sex-specific mean trajectories for top dimensions")
+    fig.tight_layout(rect=(0, 0, 0.98, 0.95))
+    fig.savefig(output_dir / f"{burden_type}_top_dimension_means_by_sex.png")
+    plt.close(fig)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Plot DeepHealth burden-index trajectories by sex."
+    )
+    parser.add_argument("--input_csv", type=str, required=True)
+    parser.add_argument("--output_dir", type=str, default=None)
+    parser.add_argument("--n_per_sex", type=int, default=10)
+    parser.add_argument("--seed", type=int, default=2026)
+    parser.add_argument("--top_dimensions", type=int, default=12)
+    args = parser.parse_args()
+
+    input_csv = Path(args.input_csv)
+    output_dir = Path(args.output_dir) if args.output_dir else input_csv.with_suffix("")
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    df = _load_bi_csv(input_csv)
+    total_df = _aggregate_total(df)
+    selected = _sample_patients(total_df, n_per_sex=int(args.n_per_sex), seed=int(args.seed))
+
+    for burden_type in sorted(df["burden_type"].dropna().unique().tolist()):
+        burden_type = str(burden_type)
+        _plot_selected_trajectories(
+            total_df,
+            burden_type=burden_type,
+            selected=selected,
+            output_dir=output_dir,
+        )
+        dims = _top_dimensions(df, burden_type=burden_type, top_n=int(args.top_dimensions))
+        _plot_dimension_mean_panels(
+            df,
+            burden_type=burden_type,
+            dimensions=dims,
+            output_dir=output_dir,
+        )
+
+    selected_rows = []
+    for sex_label, patient_ids in selected.items():
+        for pid in patient_ids.tolist():
+            selected_rows.append({"sex": sex_label, "patient_id": int(pid)})
+    pd.DataFrame(selected_rows).to_csv(output_dir / "sampled_patients_by_sex.csv", index=False)
+    print(f"Input: {input_csv}")
+    print(f"Output directory: {output_dir}")
+    print(f"Sampled patients per sex: {int(args.n_per_sex)}")
+
+
+if __name__ == "__main__":
+    main()