DeepHealth/plot_deephealth_index_trajectories.py

from __future__ import annotations

import argparse
from pathlib import Path

try:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
except ModuleNotFoundError as exc:
    raise ModuleNotFoundError(
        "plot_deephealth_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",
    "index_type",
    "index_id",
    "index_label",
    "index_value",
}


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_index_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=sorted(REQUIRED_COLUMNS))
    df["sex_label"] = df["sex"].map(_sex_label)
    df["landmark_age"] = pd.to_numeric(df["landmark_age"], errors="raise")
    df["index_value"] = pd.to_numeric(df["index_value"], 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
        samples[sex_label] = np.asarray(
            rng.choice(ids, size=min(int(n_per_sex), int(ids.size)), replace=False)
        )
    return samples


def _plot_sampled_trajectories(
    df: pd.DataFrame,
    *,
    index_type: str,
    selected: dict[str, np.ndarray],
    output_dir: Path,
) -> None:
    sub = df[df["index_type"] == index_type].copy()
    if sub.empty:
        return
    if index_type == "organ_involvement":
        total = (
            sub.groupby(["patient_id", "sex_label", "landmark_age"], as_index=False)[
                "index_value"
            ]
            .mean()
            .rename(columns={"index_value": "trajectory_value"})
        )
        title = "mean organ involvement"
    else:
        total = sub.rename(columns={"index_value": "trajectory_value"})
        title = "frailty risk"

    mean_df = (
        total.groupby(["sex_label", "landmark_age"], as_index=False)["trajectory_value"]
        .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"]:
        for pid in selected.get(sex_label, np.asarray([], dtype=np.int64)):
            one = total[total["patient_id"] == pid].sort_values("landmark_age")
            if one.empty:
                continue
            ax.plot(
                one["landmark_age"],
                one["trajectory_value"],
                color=colors.get(sex_label, "0.4"),
                alpha=0.22,
                linewidth=1.2,
            )
        mean_one = mean_df[mean_df["sex_label"] == sex_label]
        if not mean_one.empty:
            ax.plot(
                mean_one["landmark_age"],
                mean_one["trajectory_value"],
                color=colors.get(sex_label, "0.4"),
                linewidth=2.6,
                label=f"{sex_label} mean",
            )

    ax.set_title(f"{index_type}: sampled trajectories and sex-specific means")
    ax.set_xlabel("Landmark age")
    ax.set_ylabel(title)
    ax.grid(True, alpha=0.25)
    ax.legend(frameon=False)
    fig.tight_layout()
    fig.savefig(output_dir / f"{index_type}_sampled_trajectories_by_sex.png")
    plt.close(fig)


def _plot_top_dimensions(
    df: pd.DataFrame,
    *,
    output_dir: Path,
    top_n: int,
) -> None:
    sub = df[df["index_type"] == "organ_involvement"].copy()
    if sub.empty:
        return
    order = (
        sub.groupby("index_id")["index_value"]
        .mean()
        .sort_values(ascending=False)
        .head(int(top_n))
        .index.tolist()
    )
    n = len(order)
    if n == 0:
        return
    ncols = min(3, n)
    nrows = int(np.ceil(n / ncols))
    fig, axes = plt.subplots(nrows, ncols, figsize=(4.0 * ncols, 3.0 * nrows), dpi=160)
    axes_arr = np.asarray(axes).reshape(-1)
    colors = {"female": "#b83280", "male": "#2563eb"}
    for ax, index_id in zip(axes_arr, order):
        one = sub[sub["index_id"] == index_id]
        mean_df = (
            one.groupby(["sex_label", "landmark_age"], as_index=False)["index_value"]
            .mean()
            .sort_values("landmark_age")
        )
        for sex_label in ["female", "male"]:
            m = mean_df[mean_df["sex_label"] == sex_label]
            if not m.empty:
                ax.plot(
                    m["landmark_age"],
                    m["index_value"],
                    color=colors.get(sex_label, "0.4"),
                    linewidth=1.8,
                    label=sex_label,
                )
        ax.set_title(str(index_id), fontsize=9)
        ax.set_xlabel("Age")
        ax.set_ylabel("Organ involvement")
        ax.grid(True, alpha=0.22)
    for ax in axes_arr[n:]:
        ax.axis("off")
    handles, labels = axes_arr[0].get_legend_handles_labels()
    if handles:
        fig.legend(handles, labels, loc="upper right", frameon=False)
    fig.suptitle("organ_involvement: sex-specific mean trajectories for top dimensions")
    fig.tight_layout(rect=(0, 0, 0.98, 0.96))
    fig.savefig(output_dir / "organ_involvement_top_dimensions_by_sex.png")
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Plot DeepHealth organ involvement and frailty risk trajectories."
    )
    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_n", 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.parent
    output_dir.mkdir(parents=True, exist_ok=True)
    df = _load_index_csv(input_csv)
    selected = _sample_patients(df, n_per_sex=int(args.n_per_sex), seed=int(args.seed))
    for index_type in ["organ_involvement", "frailty_risk"]:
        _plot_sampled_trajectories(
            df,
            index_type=index_type,
            selected=selected,
            output_dir=output_dir,
        )
    _plot_top_dimensions(df, output_dir=output_dir, top_n=int(args.top_n))
    print(f"Input: {input_csv}")
    print(f"Output directory: {output_dir}")


if __name__ == "__main__":
    main()