from __future__ import annotations import argparse import multiprocessing as mp import time from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path from typing import Any, Iterable import numpy as np import pandas as pd import torch from torch.nn.utils.rnn import pad_sequence from torch.utils.data import DataLoader, IterableDataset, get_worker_info from tqdm.auto import tqdm from burden_index import ( build_readout_grid, load_deephealth_context, probabilities_from_hidden, ) from evaluate_auc_v2 import make_eval_indices, parse_float_list 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_devices(args: argparse.Namespace) -> list[str | None]: if args.devices is not None and str(args.devices).strip(): devices = [x.strip() for x in str(args.devices).split(",") if x.strip()] if not devices: raise ValueError("--devices was provided but no devices were parsed.") return devices return [args.device] def _load_index_matrices( *, organ_mapping_csv: Path, hfrs_mapping_csv: Path, ) -> tuple[np.ndarray, list[dict[str, Any]], dict[str, Any]]: organ_df = pd.read_csv(organ_mapping_csv) organ_required = {"token_id", "organ_id", "organ_label", "organ_weight"} missing = sorted(organ_required - set(organ_df.columns)) if missing: raise ValueError(f"{organ_mapping_csv} is missing required columns: {missing}") organ_df = organ_df.copy() organ_df["token_id"] = pd.to_numeric(organ_df["token_id"], errors="raise").astype(int) organ_df["organ_weight"] = pd.to_numeric( organ_df["organ_weight"], errors="raise" ).astype(float) organ_df = organ_df[(organ_df["organ_id"].astype(str) != "") & (organ_df["organ_weight"] > 0)] if organ_df.empty: raise ValueError(f"{organ_mapping_csv} has no mapped organ rows.") hfrs_df = pd.read_csv(hfrs_mapping_csv) hfrs_required = {"token_id", "hfrs_weight"} missing = sorted(hfrs_required - set(hfrs_df.columns)) if missing: raise ValueError(f"{hfrs_mapping_csv} is missing required columns: {missing}") hfrs_df = hfrs_df.copy() hfrs_df["token_id"] = pd.to_numeric(hfrs_df["token_id"], errors="raise").astype(int) hfrs_df["hfrs_weight"] = pd.to_numeric( hfrs_df["hfrs_weight"], errors="raise" ).astype(float) hfrs_df = hfrs_df[hfrs_df["hfrs_weight"] > 0] if hfrs_df.empty: raise ValueError(f"{hfrs_mapping_csv} has no non-zero HFRS weights.") union_disease_ids = np.asarray( sorted( set(organ_df["token_id"].astype(int).tolist()) | set(hfrs_df["token_id"].astype(int).tolist()) ), dtype=np.int64, ) union_pos = {int(token): i for i, token in enumerate(union_disease_ids.tolist())} organ_ids = sorted(organ_df["organ_id"].astype(str).unique().tolist()) organ_pos = {organ_id: i for i, organ_id in enumerate(organ_ids)} organ_matrix = np.zeros((len(organ_ids), union_disease_ids.size), dtype=np.float32) organ_meta_by_id = {} for _, row in organ_df.iterrows(): organ_id = str(row["organ_id"]) token = int(row["token_id"]) organ_matrix[organ_pos[organ_id], union_pos[token]] = 1.0 organ_meta_by_id.setdefault( organ_id, { "index_type": "organ_involvement", "index_id": organ_id, "index_label": str(row["organ_label"]), }, ) organ_meta = [organ_meta_by_id[organ_id] for organ_id in organ_ids] hfrs_weights = np.zeros(union_disease_ids.size, dtype=np.float32) for _, row in hfrs_df.iterrows(): hfrs_weights[union_pos[int(row["token_id"])]] = float(row["hfrs_weight"]) hfrs_meta = { "index_type": "frailty_risk", "index_id": "deephealth_hfrs", "index_label": "DeepHealth-HFRS frailty risk index", } matrices = [ { "kind": "organ_involvement", "matrix": organ_matrix, "meta": organ_meta, }, { "kind": "frailty_risk", "weights": hfrs_weights, "meta": hfrs_meta, }, ] return union_disease_ids, matrices, { "organ_mapped_tokens": int(organ_df["token_id"].nunique()), "hfrs_mapped_tokens": int(hfrs_df["token_id"].nunique()), } 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, ) class _Sized: def __len__(self) -> int: return n return make_eval_indices(_Sized(), args, cfg) 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(): t_query = np.float32(float(landmark_age)) if not (followup_end > float(t_query)): continue prefix_mask = full_time <= t_query if not np.any(prefix_mask): continue prefix_events = full_event[prefix_mask].astype(np.int64, 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": t_query, "t_query": t_query, "followup_end_time": np.float32(followup_end), "event_seq": prefix_events, "time_seq": full_time[prefix_mask].astype(np.float32, copy=False), "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 _row_to_worker_spec(row: dict[str, Any]) -> dict[str, Any]: return { "patient_id": int(row["patient_id"]), "dataset_index": int(row["dataset_index"]), "landmark_age": float(row["landmark_age"]), "followup_end_time": float(row["followup_end_time"]), } def _materialize_worker_rows( dataset: Any, row_specs: list[dict[str, Any]], ) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] for spec in row_specs: sample = dataset.samples[int(spec["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) full_event = np.concatenate([seq_event, tgt_event[-1:]]) full_time = np.concatenate([seq_time, tgt_time[-1:]]) t_query = np.float32(float(spec["landmark_age"])) prefix_mask = full_time <= t_query rows.append( { "patient_id": int(spec["patient_id"]), "dataset_index": int(spec["dataset_index"]), "sex": int(sample["sex"]), "landmark_age": t_query, "t_query": t_query, "followup_end_time": np.float32(float(spec["followup_end_time"])), "event_seq": full_event[prefix_mask].astype(np.int64, copy=False), "time_seq": full_time[prefix_mask].astype(np.float32, copy=False), "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 class HistoricalReadoutDataset(IterableDataset): def __init__(self, rows: list[dict[str, Any]]) -> None: super().__init__() self.rows = rows def __iter__(self) -> Iterable[dict[str, torch.Tensor]]: worker = get_worker_info() if worker is None: start, step = 0, 1 else: start, step = int(worker.id), int(worker.num_workers) for row_idx in range(start, len(self.rows), step): row = self.rows[row_idx] grid = build_readout_grid( event_seq=row["event_seq"], time_seq=row["time_seq"], other_type=row["other_type"], other_time=row["other_time"], t_query=float(row["t_query"]), ) if grid.size == 0: continue end_times = np.concatenate([grid[1:], np.asarray([row["t_query"]], dtype=np.float32)]) deltas = np.maximum(end_times - grid, 0.0).astype(np.float32) valid = deltas > 0 for query_time, delta in zip(grid[valid].tolist(), deltas[valid].tolist()): yield _make_readout_job(row, row_idx, query_time, delta) def _make_readout_job( row: dict[str, Any], row_idx: int, query_time: float, delta: float, ) -> dict[str, torch.Tensor]: return { "event_seq": torch.from_numpy(np.asarray(row["event_seq"], dtype=np.int64)).long(), "time_seq": torch.from_numpy(np.asarray(row["time_seq"], dtype=np.float32)).float(), "sex": torch.tensor(int(row["sex"]), dtype=torch.long), "other_type": torch.from_numpy(np.asarray(row["other_type"], dtype=np.int64)).long(), "other_value": torch.from_numpy(np.asarray(row["other_value"], dtype=np.float32)).float(), "other_value_kind": torch.from_numpy( np.asarray(row["other_value_kind"], dtype=np.int64) ).long(), "other_time": torch.from_numpy(np.asarray(row["other_time"], dtype=np.float32)).float(), "query_time": torch.tensor(float(query_time), dtype=torch.float32), "delta": torch.tensor(float(delta), dtype=torch.float32), "row_idx": torch.tensor(int(row_idx), dtype=torch.long), } def _collate_readout_jobs(batch: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]: event_seq = pad_sequence( [x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX ) return { "event_seq": event_seq, "time_seq": pad_sequence( [x["time_seq"] for x in batch], batch_first=True, padding_value=0.0 ), "padding_mask": event_seq > PAD_IDX, "sex": torch.stack([x["sex"] for x in batch]), "other_type": pad_sequence( [x["other_type"] for x in batch], batch_first=True, padding_value=0 ), "other_value": pad_sequence( [x["other_value"] for x in batch], batch_first=True, padding_value=0.0 ), "other_value_kind": pad_sequence( [x["other_value_kind"] for x in batch], batch_first=True, padding_value=0 ), "other_time": pad_sequence( [x["other_time"] for x in batch], batch_first=True, padding_value=0.0 ), "query_time": torch.stack([x["query_time"] for x in batch]), "delta": torch.stack([x["delta"] for x in batch]), "row_idx": torch.stack([x["row_idx"] for x in batch]), } @torch.inference_mode() def _readout_probabilities( *, ctx: Any, batch: dict[str, torch.Tensor], disease_ids: np.ndarray, ) -> torch.Tensor: event = batch["event_seq"].long().to(ctx.device, non_blocking=True) hidden = ctx.model( event_seq=event, time_seq=batch["time_seq"].float().to(ctx.device, non_blocking=True), sex=batch["sex"].long().to(ctx.device, non_blocking=True), padding_mask=event > PAD_IDX, t_query=batch["query_time"].float().to(ctx.device, non_blocking=True), other_type=batch["other_type"].long().to(ctx.device, non_blocking=True), other_value=batch["other_value"].float().to(ctx.device, non_blocking=True), other_value_kind=batch["other_value_kind"].long().to(ctx.device, non_blocking=True), other_time=batch["other_time"].float().to(ctx.device, non_blocking=True), target_mode="all_future", ) deltas = batch["delta"].detach().cpu().numpy().astype(np.float32, copy=False) prob = probabilities_from_hidden( ctx=ctx, hidden=hidden, disease_ids=disease_ids, deltas=deltas, ) return torch.as_tensor(prob, dtype=torch.float32, device=ctx.device) def _project_rows( *, rows: list[dict[str, Any]], survival_by_row: torch.Tensor, matrices: list[dict[str, Any]], ctx: Any, ) -> list[dict[str, Any]]: disease_expression = 1.0 - survival_by_row.clamp(0.0, 1.0) disease_intensity = -torch.log(survival_by_row.clamp(1e-7, 1.0)) out: list[dict[str, Any]] = [] organ_matrix = torch.as_tensor( matrices[0]["matrix"], dtype=torch.float32, device=ctx.device ) organ_values = -torch.expm1(-torch.matmul(disease_intensity, organ_matrix.T)) organ_values_np = organ_values.detach().cpu().numpy() hfrs_weights = torch.as_tensor( matrices[1]["weights"], dtype=torch.float32, device=ctx.device ) hfrs_values = torch.matmul(disease_expression, hfrs_weights) hfrs_values_np = hfrs_values.detach().cpu().numpy() for row_idx, row in enumerate(rows): base = { "patient_id": row["patient_id"], "dataset_index": row["dataset_index"], "sex": row["sex"], "landmark_age": float(row["landmark_age"]), "t_query": float(row["t_query"]), "followup_end_time": float(row["followup_end_time"]), } for dim_idx, meta in enumerate(matrices[0]["meta"]): out.append( { **base, "index_type": meta["index_type"], "index_id": meta["index_id"], "index_label": meta["index_label"], "index_value": float(organ_values_np[row_idx, dim_idx]), } ) out.append( { **base, "index_type": matrices[1]["meta"]["index_type"], "index_id": matrices[1]["meta"]["index_id"], "index_label": matrices[1]["meta"]["index_label"], "index_value": float(hfrs_values_np[row_idx]), } ) return out def _compute_rows( *, rows: list[dict[str, Any]], disease_ids: np.ndarray, matrices: list[dict[str, Any]], readout_batch_size: int, num_workers: int, ctx: Any, log_prefix: str, ) -> tuple[list[dict[str, Any]], int, dict[str, float]]: survival_by_row = torch.ones( (len(rows), disease_ids.size), dtype=torch.float32, device=ctx.device ) loader = DataLoader( HistoricalReadoutDataset(rows), batch_size=max(1, int(readout_batch_size)), collate_fn=_collate_readout_jobs, num_workers=max(0, int(num_workers)), pin_memory=ctx.device.type == "cuda", persistent_workers=int(num_workers) > 0, prefetch_factor=2 if int(num_workers) > 0 else None, ) readout_jobs = 0 n_batches = 0 forward_sec = 0.0 reduce_sec = 0.0 for batch in loader: t0 = time.perf_counter() prob = _readout_probabilities(ctx=ctx, batch=batch, disease_ids=disease_ids) if ctx.device.type == "cuda": torch.cuda.synchronize(ctx.device) forward_sec += time.perf_counter() - t0 t1 = time.perf_counter() row_indices = batch["row_idx"].long().to(ctx.device, non_blocking=True) interval_survival = 1.0 - prob.clamp(0.0, 1.0) if hasattr(survival_by_row, "scatter_reduce_"): survival_by_row.scatter_reduce_( dim=0, index=row_indices[:, None].expand_as(interval_survival), src=interval_survival, reduce="prod", include_self=True, ) else: for job_idx in range(interval_survival.shape[0]): survival_by_row[int(row_indices[job_idx].item())] *= interval_survival[job_idx] if ctx.device.type == "cuda": torch.cuda.synchronize(ctx.device) reduce_sec += time.perf_counter() - t1 n_batches += 1 readout_jobs += int(batch["row_idx"].numel()) if n_batches == 1 or n_batches % 50 == 0: print( f"{log_prefix} processed {readout_jobs} readout jobs in {n_batches} batches", flush=True, ) t2 = time.perf_counter() out = _project_rows( rows=rows, survival_by_row=survival_by_row, matrices=matrices, ctx=ctx, ) if ctx.device.type == "cuda": torch.cuda.synchronize(ctx.device) reduce_sec += time.perf_counter() - t2 return out, readout_jobs, {"forward_sec": forward_sec, "reduce_sec": reduce_sec} def _write_rows_csv(rows: list[dict[str, Any]], output_path: Path) -> int: df = pd.DataFrame(rows) df.to_csv(output_path, index=False) return int(len(df)) def _concat_csv_shards(shard_paths: list[Path], output_path: Path) -> None: wrote_header = False with output_path.open("w", encoding="utf-8", newline="") as out_f: for shard_path in shard_paths: with shard_path.open("r", encoding="utf-8", newline="") as in_f: header = in_f.readline() if not wrote_header: out_f.write(header) wrote_header = True for line in in_f: out_f.write(line) shard_path.unlink(missing_ok=True) def _estimate_jobs(row: dict[str, Any]) -> int: grid = build_readout_grid( event_seq=row["event_seq"], time_seq=row["time_seq"], other_type=row["other_type"], other_time=row["other_time"], t_query=float(row["t_query"]), ) if grid.size == 0: return 1 end_times = np.concatenate([grid[1:], np.asarray([row["t_query"]], dtype=np.float32)]) return max(int(np.sum(np.maximum(end_times - grid, 0.0) > 0)), 1) def _split_rows(rows: list[dict[str, Any]], devices: list[str | None]) -> list[tuple[str | None, list[dict[str, Any]]]]: if len(devices) <= 1: return [(devices[0], rows)] buckets: list[list[dict[str, Any]]] = [[] for _ in devices] loads = np.zeros(len(devices), dtype=np.int64) for row in sorted(rows, key=_estimate_jobs, reverse=True): idx = int(np.argmin(loads)) buckets[idx].append(row) loads[idx] += _estimate_jobs(row) return [(device, bucket) for device, bucket in zip(devices, buckets) if bucket] def _worker(payload: dict[str, Any]) -> dict[str, Any]: device = payload["device"] shard_path = Path(payload["shard_path"]) print(f"[Index worker {device}] starting with {len(payload['row_specs'])} rows", flush=True) ctx = load_deephealth_context(payload["run_path"], device=device) rows = _materialize_worker_rows(ctx.dataset, payload["row_specs"]) out, readout_jobs, timings = _compute_rows( rows=rows, disease_ids=payload["disease_ids"], matrices=payload["matrices"], readout_batch_size=int(payload["readout_batch_size"]), num_workers=int(payload["num_workers"]), ctx=ctx, log_prefix=f"[Index worker {device}]", ) t0 = time.perf_counter() row_count = _write_rows_csv(out, shard_path) timings["write_csv_sec"] = time.perf_counter() - t0 print(f"[Index worker {device}] wrote {row_count} rows to {shard_path}", flush=True) return { "shard_path": str(shard_path), "row_count": row_count, "readout_jobs": readout_jobs, "timings": timings, } def main() -> None: parser = argparse.ArgumentParser( description="Compute DeepHealth organ involvement and frailty risk indices." ) parser.add_argument("--run_path", type=str, required=True) parser.add_argument( "--organ_mapping_csv", type=str, default="organ_involvement_label_mapping.csv", ) parser.add_argument("--hfrs_mapping_csv", type=str, default="uk_hfrs_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("--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("--readout_batch_size", type=int, default=8192) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--device", type=str, default=None) parser.add_argument("--devices", type=str, default=None) parser.add_argument( "--mp_start_method", type=str, default="fork", choices=["fork", "forkserver"], ) args = parser.parse_args() run_path = Path(args.run_path) devices = _parse_devices(args) initial_device = "cpu" if len(devices) > 1 else devices[0] ctx = load_deephealth_context(run_path, device=initial_device) disease_ids, matrices, mapping_counts = _load_index_matrices( organ_mapping_csv=Path(args.organ_mapping_csv), hfrs_mapping_csv=Path(args.hfrs_mapping_csv), ) landmark_ages = _parse_landmark_ages(args) 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.") output_path = Path(args.output_path) if args.output_path else ( run_path / f"deephealth_indices_{eval_split}.csv" ) output_path.parent.mkdir(parents=True, exist_ok=True) chunks = _split_rows(rows, devices) for device, chunk in chunks: print( f"Assigned {len(chunk)} rows / ~{sum(_estimate_jobs(r) for r in chunk)} " f"readout jobs to {device}", flush=True, ) total_readout_jobs = 0 timings = {"forward_sec": 0.0, "reduce_sec": 0.0, "write_csv_sec": 0.0} if len(chunks) == 1: out, total_readout_jobs, chunk_timings = _compute_rows( rows=rows, disease_ids=disease_ids, matrices=matrices, readout_batch_size=int(args.readout_batch_size), num_workers=int(args.num_workers), ctx=ctx, log_prefix="[Index main]", ) for key, value in chunk_timings.items(): timings[key] += float(value) t0 = time.perf_counter() _write_rows_csv(out, output_path) timings["write_csv_sec"] = time.perf_counter() - t0 else: del ctx payloads = [ { "device": device, "run_path": str(run_path), "shard_path": str( output_path.with_name( f"{output_path.stem}.part{part_idx:03d}{output_path.suffix}" ) ), "row_specs": [_row_to_worker_spec(row) for row in chunk], "disease_ids": disease_ids, "matrices": matrices, "readout_batch_size": int(args.readout_batch_size), "num_workers": int(args.num_workers), } for part_idx, (device, chunk) in enumerate(chunks) ] shard_paths = [] with ProcessPoolExecutor( max_workers=len(payloads), mp_context=mp.get_context(args.mp_start_method), ) as executor: futures = [executor.submit(_worker, payload) for payload in payloads] for future in tqdm( as_completed(futures), total=len(futures), desc="Computing DeepHealth index chunks", dynamic_ncols=True, ): result = future.result() shard_paths.append(Path(result["shard_path"])) total_readout_jobs += int(result["readout_jobs"]) for key, value in result["timings"].items(): timings[key] += float(value) t0 = time.perf_counter() _concat_csv_shards(sorted(shard_paths), output_path) timings["write_csv_sec"] += time.perf_counter() - t0 print(f"Run path: {run_path}") print(f"Eval split: {eval_split}") print(f"Landmark rows: {len(rows)}") print(f"Readout jobs: {total_readout_jobs}") print(f"Union disease tokens: {disease_ids.size}") print(f"Organ mapped tokens: {mapping_counts['organ_mapped_tokens']}") print(f"HFRS mapped tokens: {mapping_counts['hfrs_mapped_tokens']}") print( "Timing seconds: " f"forward={timings['forward_sec']:.2f}, " f"reduce={timings['reduce_sec']:.2f}, " f"write_csv={timings['write_csv_sec']:.2f}" ) print(f"Devices: {', '.join(str(d) for d, _ in chunks)}") print(f"Output: {output_path}") if __name__ == "__main__": main()