259 lines
8.0 KiB
Python
259 lines
8.0 KiB
Python
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()
|