Summarize disease mortality attribution by age and sex

evaluate_single_disease_mortality_attribution.py

@@ -75,6 +75,30 @@ OUTPUT_COLUMNS = [
     "mortality_attribution_hazard_ratio",
 ]
 
+SUMMARY_KEY_COLUMNS = [
+    "selected_disease_token_id",
+    "selected_disease_code",
+    "selected_disease_name",
+    "selected_disease_organ_system",
+    "selected_disease_organ_system_label",
+    "landmark_age",
+    "sex",
+]
+
+SUMMARY_MEAN_COLUMNS = [
+    "history_disease_count",
+    "selected_disease_history_count",
+    "history_count__selected_organ_system",
+    "death_risk",
+    "death_hazard",
+    "ablated_death_risk",
+    "ablated_death_hazard",
+    "mortality_attribution_probability",
+    "mortality_attribution_hazard",
+    "mortality_attribution_probability_ratio",
+    "mortality_attribution_hazard_ratio",
+]
+
 
 def write_compressed_npz_table(path: Path, table: pd.DataFrame) -> int:
     table = table.reindex(columns=OUTPUT_COLUMNS)
@@ -102,6 +126,7 @@ def write_manifest(
     *,
     rows: int,
     shards: list[dict[str, Any]],
+    summary_file: str,
     scanned_diseases: list[dict[str, Any]],
     eval_split: str,
     tau: float,
@@ -114,6 +139,7 @@ def write_manifest(
         "columns": OUTPUT_COLUMNS,
         "rows": int(rows),
         "shards": shards,
+        "summary_file": summary_file,
         "scanned_diseases": scanned_diseases,
         "eval_split": eval_split,
         "tau": float(tau),
@@ -125,6 +151,51 @@ def write_manifest(
         json.dump(payload, f, ensure_ascii=False, indent=2)
 
 
+def update_summary_accumulator(
+    summary: dict[tuple[Any, ...], dict[str, float]],
+    table: pd.DataFrame,
+) -> None:
+    if table.empty:
+        return
+    grouped = table.groupby(SUMMARY_KEY_COLUMNS, dropna=False, sort=False)
+    for key, group in grouped:
+        if not isinstance(key, tuple):
+            key = (key,)
+        acc = summary.setdefault(
+            key,
+            {"n": 0.0, **{column: 0.0 for column in SUMMARY_MEAN_COLUMNS}},
+        )
+        n = int(len(group))
+        acc["n"] += float(n)
+        for column in SUMMARY_MEAN_COLUMNS:
+            acc[column] += float(pd.to_numeric(group[column], errors="coerce").sum())
+
+
+def write_summary_csv(
+    path: Path,
+    summary: dict[tuple[Any, ...], dict[str, float]],
+) -> int:
+    rows: list[dict[str, Any]] = []
+    for key, acc in summary.items():
+        n = int(acc["n"])
+        out = {column: value for column, value in zip(SUMMARY_KEY_COLUMNS, key)}
+        out["n"] = n
+        for column in SUMMARY_MEAN_COLUMNS:
+            out[f"mean__{column}"] = acc[column] / n if n > 0 else np.nan
+        rows.append(out)
+
+    columns = [
+        *SUMMARY_KEY_COLUMNS,
+        "n",
+        *[f"mean__{column}" for column in SUMMARY_MEAN_COLUMNS],
+    ]
+    pd.DataFrame(rows, columns=columns).sort_values(
+        ["selected_disease_token_id", "landmark_age", "sex"],
+        kind="mergesort",
+    ).to_csv(path, index=False)
+    return len(rows)
+
+
 def load_disease_metadata(
     mapping_path: Path,
     *,
@@ -422,6 +493,7 @@ def main() -> None:
     written_rows = 0
     shard_index = 0
     shards: list[dict[str, Any]] = []
+    summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
     pending_batch_chunks: list[Dict[str, torch.Tensor]] = []
     pending_meta_chunks: list[list[dict[str, Any]]] = []
     pending_n = 0
@@ -479,6 +551,7 @@ def main() -> None:
             )
 
         table = pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)
+        update_summary_accumulator(summary_accumulator, table)
         shard_name = f"part-{shard_index:06d}.npz"
         shard_path = output_dir / shard_name
         shard_rows = write_compressed_npz_table(shard_path, table)
@@ -607,10 +680,13 @@ def main() -> None:
         empty_path = output_dir / "part-000000.npz"
         write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))
         shards.append({"file": empty_path.name, "rows": 0})
+    summary_path = output_dir / "summary_by_disease_age_sex.csv"
+    summary_rows = write_summary_csv(summary_path, summary_accumulator)
     write_manifest(
         output_dir,
         rows=written_rows,
         shards=shards,
+        summary_file=summary_path.name,
         scanned_diseases=[
             {"token_id": int(token), **{k: v for k, v in meta.items() if k != "token_id"}}
             for token, meta in scanned_disease_items
@@ -622,6 +698,7 @@ def main() -> None:
         landmark_step=float(args.landmark_step),
     )
     print(f"Wrote {written_rows} rows in {len(shards)} shard(s) to {output_dir}")
+    print(f"Wrote {summary_rows} summary rows to {summary_path}")
 
 
 if __name__ == "__main__":