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()