diff --git a/plot_burden_index_trajectories.py b/plot_burden_index_trajectories.py new file mode 100644 index 0000000..ecbb118 --- /dev/null +++ b/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()