from __future__ import annotations import argparse import math from pathlib import Path from typing import Any, Iterable import numpy as np import pandas as pd from tqdm.auto import tqdm from burden_index import compute_burden_index, load_burden_context from evaluate_auc_v2 import ( make_eval_indices, parse_float_list, split_indices, ) from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX def _parse_landmark_ages(args: argparse.Namespace) -> np.ndarray: explicit = parse_float_list(args.landmark_ages) if explicit: ages = np.asarray(explicit, dtype=np.float32) else: ages = np.arange( float(args.landmark_start), float(args.landmark_stop) + 1e-6, float(args.landmark_step), dtype=np.float32, ) if ages.size == 0: raise ValueError("No landmark ages were provided.") return ages def _parse_horizons(value: Any) -> np.ndarray: horizons = np.asarray(parse_float_list(value) or [5.0], dtype=np.float32) if horizons.size == 0: raise ValueError("horizons must contain at least one value.") if np.any(horizons < 0): raise ValueError(f"horizons must be non-negative, got {horizons}") return horizons def _build_burden_matrix_from_mapping( mapping_csv: Path, *, token_col: str, category_id_col: str, category_col: str, key_area_col: str, weight_col: str, ) -> tuple[np.ndarray, np.ndarray, pd.DataFrame]: df = pd.read_csv(mapping_csv) required = {token_col, category_id_col, category_col, weight_col} missing = sorted(required - set(df.columns)) if missing: raise ValueError(f"{mapping_csv} is missing required columns: {missing}") df = df.copy() df[token_col] = pd.to_numeric(df[token_col], errors="raise").astype(int) raw_category = df[category_id_col] numeric_category = pd.to_numeric(raw_category, errors="coerce") if numeric_category.notna().all(): df["_burden_category_key"] = numeric_category.astype(int) else: df["_burden_category_key"] = raw_category.astype(str) df[weight_col] = pd.to_numeric(df[weight_col], errors="raise").astype(float) df = df[df[weight_col] != 0].copy() if df.empty: raise ValueError(f"{mapping_csv} has no non-zero burden weights.") disease_ids = np.asarray(sorted(df[token_col].unique().tolist()), dtype=np.int64) category_keys = sorted(df["_burden_category_key"].unique().tolist()) disease_pos = {int(token): j for j, token in enumerate(disease_ids.tolist())} category_pos = {cat: i for i, cat in enumerate(category_keys)} A = np.zeros((len(category_keys), disease_ids.size), dtype=np.float64) for _, row in df.iterrows(): token = int(row[token_col]) cat = row["_burden_category_key"] weight = float(row[weight_col]) A[category_pos[cat], disease_pos[token]] += weight meta_cols = list(dict.fromkeys(["_burden_category_key", category_id_col, category_col])) if key_area_col in df.columns: meta_cols.append(key_area_col) category_meta = ( df[meta_cols] .drop_duplicates(subset=["_burden_category_key"]) .sort_values("_burden_category_key") .reset_index(drop=True) ) category_meta = category_meta.rename( columns={ "_burden_category_key": "burden_dimension_id", category_col: "burden_dimension", key_area_col: "burden_key_area", } ) if category_id_col in category_meta.columns: category_meta = category_meta.drop(columns=[category_id_col]) if "burden_key_area" not in category_meta.columns: category_meta["burden_key_area"] = "" return A, disease_ids, category_meta def _parse_burden_types(value: str) -> list[str]: out = [x.strip().lower() for x in str(value).split(",") if x.strip()] if not out: raise ValueError("burden_types must contain at least one value.") valid = {"functional", "organ"} unknown = sorted(set(out) - valid) if unknown: raise ValueError(f"Unknown burden_types {unknown}; valid values are {sorted(valid)}") return list(dict.fromkeys(out)) def _load_mapping_specs(args: argparse.Namespace) -> list[dict[str, Any]]: specs: list[dict[str, Any]] = [] for burden_type in _parse_burden_types(args.burden_types): if burden_type == "functional": path = Path(args.functional_mapping_csv) specs.append( { "burden_type": "functional", "mapping_csv": path, "token_col": "token_id", "category_id_col": "hfrm_category_id", "category_col": "hfrm_category", "key_area_col": "hfrm_key_area", "weight_col": args.functional_weight_col, } ) elif burden_type == "organ": path = Path(args.organ_mapping_csv) specs.append( { "burden_type": "organ", "mapping_csv": path, "token_col": "token_id", "category_id_col": "organ_system", "category_col": "organ_system", "key_area_col": "icd10_chapter_name", "weight_col": "organ_weight", } ) for spec in specs: if not spec["mapping_csv"].exists(): raise FileNotFoundError( f"{spec['burden_type']} mapping_csv not found: {spec['mapping_csv']}" ) return specs def _eligible_landmark_rows( dataset: Any, subset_indices: np.ndarray, landmark_ages: np.ndarray, *, min_history_events: int, ) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] special = np.asarray([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64) for patient_id, dataset_index in enumerate(subset_indices.tolist()): sample = dataset.samples[int(dataset_index)] seq_event = np.asarray(sample["event_seq"], dtype=np.int64) seq_time = np.asarray(sample["time_seq"], dtype=np.float32) tgt_event = np.asarray(sample["target_event_seq"], dtype=np.int64) tgt_time = np.asarray(sample["target_time_seq"], dtype=np.float32) if seq_event.size == 0 or tgt_event.size == 0: continue full_event = np.concatenate([seq_event, tgt_event[-1:]]) full_time = np.concatenate([seq_time, tgt_time[-1:]]) followup_end = float(np.max(full_time)) for landmark_age in landmark_ages.tolist(): landmark_age = float(landmark_age) if not (followup_end > landmark_age): continue prefix_mask = full_time <= np.float32(landmark_age) if not np.any(prefix_mask): continue prefix_events = full_event[prefix_mask].astype(np.int64, copy=False) prefix_times = full_time[prefix_mask].astype(np.float32, copy=False) valid_history = ~np.isin(prefix_events, special) if int(valid_history.sum()) < int(min_history_events): continue rows.append( { "patient_id": int(patient_id), "dataset_index": int(dataset_index), "sex": int(sample["sex"]), "landmark_age": np.float32(landmark_age), "t_query": np.float32(landmark_age), "followup_end_time": np.float32(followup_end), "event_seq": prefix_events, "time_seq": prefix_times, "other_type": np.asarray(sample["other_type"], dtype=np.int64), "other_value": np.asarray(sample["other_value"], dtype=np.float32), "other_value_kind": np.asarray( sample["other_value_kind"], dtype=np.int64), "other_time": np.asarray(sample["other_time"], dtype=np.float32), } ) return rows def _config_split_indices( n: int, cfg: dict[str, Any], eval_split: str, subset_size: int, ) -> np.ndarray: args = argparse.Namespace( train_ratio=None, val_ratio=None, test_ratio=None, seed=None, eval_split=eval_split, dataset_subset_size=subset_size if subset_size > 0 else None, ) # make_eval_indices reads len(dataset), so a small shim is enough. class _Sized: def __len__(self) -> int: return n return make_eval_indices(_Sized(), args, cfg) def _result_rows_for_sample( *, sample_row: dict[str, Any], horizons: Iterable[float], A: np.ndarray, disease_ids: np.ndarray, category_meta: pd.DataFrame, burden_type: str, formed_mode: str, ctx: Any, run_path: Path, ) -> list[dict[str, Any]]: out: list[dict[str, Any]] = [] for horizon in horizons: result = compute_burden_index( run_path=run_path, burden_matrix=A, disease_ids=disease_ids, event_seq=sample_row["event_seq"], time_seq=sample_row["time_seq"], sex=int(sample_row["sex"]), other_type=sample_row["other_type"], other_value=sample_row["other_value"], other_value_kind=sample_row["other_value_kind"], other_time=sample_row["other_time"], t_query=float(sample_row["t_query"]), horizon=float(horizon), formed_mode=formed_mode, context=ctx, ) for dim_idx, meta in category_meta.iterrows(): out.append( { "patient_id": sample_row["patient_id"], "dataset_index": sample_row["dataset_index"], "sex": sample_row["sex"], "landmark_age": float(sample_row["landmark_age"]), "t_query": float(sample_row["t_query"]), "followup_end_time": float(sample_row["followup_end_time"]), "horizon": float(horizon), "formed_mode": formed_mode, "burden_type": burden_type, "burden_dimension_id": meta["burden_dimension_id"], "burden_dimension": str(meta["burden_dimension"]), "burden_key_area": str(meta.get("burden_key_area", "")), "bi_historical": float(result.historical[int(dim_idx)]), "bi_future": float(result.future[int(dim_idx)]), "bi_total": float(result.total[int(dim_idx)]), } ) return out def main() -> None: parser = argparse.ArgumentParser( description="Compute DeepHealth Burden Indices at landmark ages." ) parser.add_argument("--run_path", type=str, required=True) parser.add_argument("--burden_types", type=str, default="functional,organ", help="Comma-separated burden types to compute: functional,organ.") parser.add_argument("--functional_mapping_csv", type=str, default="cihi_hfrm_label_mapping.csv") parser.add_argument("--organ_mapping_csv", type=str, default="icd10_organ_label_mapping.csv") parser.add_argument("--output_path", type=str, default=None) parser.add_argument("--eval_split", type=str, default="test", choices=["train", "val", "valid", "validation", "test", "all"]) parser.add_argument("--formed_mode", type=str, default="model_weighted", choices=["observed", "model_weighted"]) parser.add_argument("--horizons", type=str, default="5") parser.add_argument("--landmark_ages", type=str, default=None) parser.add_argument("--landmark_start", type=float, default=40.0) parser.add_argument("--landmark_stop", type=float, default=80.0) parser.add_argument("--landmark_step", type=float, default=5.0) parser.add_argument("--min_history_events", type=int, default=1) parser.add_argument("--dataset_subset_size", type=int, default=0) parser.add_argument("--device", type=str, default=None) parser.add_argument("--functional_weight_col", type=str, default="hfrm_normalized_weight") args = parser.parse_args() run_path = Path(args.run_path) mapping_specs = _load_mapping_specs(args) ctx = load_burden_context(run_path, device=args.device) matrices = [] for spec in mapping_specs: A, disease_ids, category_meta = _build_burden_matrix_from_mapping( spec["mapping_csv"], token_col=spec["token_col"], category_id_col=spec["category_id_col"], category_col=spec["category_col"], key_area_col=spec["key_area_col"], weight_col=spec["weight_col"], ) matrices.append( { "burden_type": spec["burden_type"], "A": A, "disease_ids": disease_ids, "category_meta": category_meta, } ) landmark_ages = _parse_landmark_ages(args) horizons = _parse_horizons(args.horizons) eval_split = str(args.eval_split).lower() if eval_split in {"valid", "validation"}: eval_split = "val" subset_indices = _config_split_indices( len(ctx.dataset), ctx.cfg, eval_split, int(args.dataset_subset_size), ) rows = _eligible_landmark_rows( ctx.dataset, subset_indices, landmark_ages, min_history_events=int(args.min_history_events), ) if not rows: raise RuntimeError( "No eligible landmark rows. Check eval split, landmark ages, and min_history_events." ) output_path = Path(args.output_path) if args.output_path else ( run_path / f"burden_index_{eval_split}_{args.formed_mode}.csv" ) output_path.parent.mkdir(parents=True, exist_ok=True) all_rows: list[dict[str, Any]] = [] for row in tqdm(rows, desc="Computing BI", dynamic_ncols=True): for matrix in matrices: all_rows.extend( _result_rows_for_sample( sample_row=row, horizons=horizons.tolist(), A=matrix["A"], disease_ids=matrix["disease_ids"], category_meta=matrix["category_meta"], burden_type=matrix["burden_type"], formed_mode=args.formed_mode, ctx=ctx, run_path=run_path, ) ) out_df = pd.DataFrame(all_rows) out_df.to_csv(output_path, index=False) print(f"Run path: {run_path}") print(f"Eval split: {eval_split}") print(f"Landmark rows: {len(rows)}") for matrix in matrices: print( f"{matrix['burden_type']} dimensions: {matrix['A'].shape[0]}, " f"mapped disease tokens: {matrix['A'].shape[1]}" ) print(f"Output: {output_path}") if __name__ == "__main__": main()