diff --git a/.gitattributes b/.gitattributes index dfdb8b7..1bbd695 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,2 @@ *.sh text eol=lf +*.py text eol=lf diff --git a/evaluate_extra_info_attribution.py b/evaluate_extra_info_attribution.py new file mode 100644 index 0000000..dbccc22 --- /dev/null +++ b/evaluate_extra_info_attribution.py @@ -0,0 +1,791 @@ +"""Evaluate extra-info attribution to death parameters and future disease risks. + +For each landmark query, this script scans selected extra-info types that are +available at or before the query age. For each such type it re-runs the model +with that extra-info type removed and summarizes: + +* death distribution parameters before and after ablation; +* tau-year future incident disease risk before and after ablation, by ICD-10 + chapter-derived organ/system groups. + +Death is always token vocab_size - 1. +""" +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +import pandas as pd +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader +from tqdm.auto import tqdm + +from evaluate_auc_v2 import ( + build_model_from_dataset, + cfg_get, + load_checkpoint_state_dict, + load_json_config, + load_model_state, + resolve_dist_mode_for_checkpoint, + resolve_eval_device, + validate_dataset_metadata, +) +from landmark_eval_utils import ( + IndexedLandmarkDataset, + LandmarkDataset, + build_first_occurrence_maps_for_landmarks, + collate_indexed_landmark_fn, + infer_landmark_hidden, + load_eval_sequence_dataset, + load_organ_groups, + make_landmark_ages, + make_occurred_mask, +) +from future_risk import new_disease_risk_from_probabilities, probabilities_from_logits + + +EXTRA_KEY_COLUMNS = [ + "selected_extra_info_type_id", + "selected_extra_info_var_name", + "selected_extra_info_full_name", + "landmark_age", + "sex", +] + +DEATH_PARAMETER_COLUMNS = [ + "original_death_lambda", + "ablated_death_lambda", + "original_death_scale", + "ablated_death_scale", + "original_death_shape", + "ablated_death_shape", +] + +DISEASE_RISK_KEY_COLUMNS = [ + *EXTRA_KEY_COLUMNS, + "target_group", + "target_group_label", +] + +DISEASE_RISK_COLUMNS = [ + "original_future_disease_risk", + "ablated_future_disease_risk", +] + + +def parse_int_list(value: Any) -> list[int] | None: + if value is None: + return None + if isinstance(value, (list, tuple, np.ndarray)): + return [int(x) for x in value] + text = str(value).strip() + if text == "": + return None + if text.startswith("["): + raw = json.loads(text) + if not isinstance(raw, list): + raise ValueError("Expected JSON list for integer list") + return [int(x) for x in raw] + return [int(x.strip()) for x in re.split(r"[,;\s]+", text) if x.strip()] + + +def load_extra_info_metadata( + *, + dataset_extra_info_types: Sequence[int], + search_root: Path = Path("."), +) -> dict[int, dict[str, Any]]: + metadata: dict[int, dict[str, Any]] = { + int(type_id): { + "type_id": int(type_id), + "var_name": f"extra_info_{int(type_id)}", + "full_name": f"extra-info type {int(type_id)}", + } + for type_id in dataset_extra_info_types + } + + line_re = re.compile(r"^\s*(\d+)\s*#\s*([^|#]+?)(?:\s*\|\s*(.*?))?\s*$") + for path in sorted(search_root.glob("extra_info_types*.txt")): + for line in path.read_text(encoding="utf-8").splitlines(): + match = line_re.match(line) + if not match: + continue + type_id = int(match.group(1)) + if type_id not in metadata: + continue + var_name = match.group(2).strip() + full_name = (match.group(3) or var_name).strip() + metadata[type_id] = { + "type_id": type_id, + "var_name": var_name, + "full_name": full_name, + } + + return metadata + + +def resolve_extra_info_types( + value: str | None, + *, + dataset_extra_info_types: Sequence[int], + metadata: dict[int, dict[str, Any]], +) -> list[int]: + available = [int(x) for x in dataset_extra_info_types] + if value is None or str(value).strip() == "": + return available + + out: list[int] = [] + seen: set[int] = set() + by_var = { + str(meta.get("var_name", "")).lower(): int(type_id) + for type_id, meta in metadata.items() + } + by_full = { + str(meta.get("full_name", "")).lower(): int(type_id) + for type_id, meta in metadata.items() + } + for part in re.split(r"[,;]+", str(value)): + text = part.strip() + if not text: + continue + if text.isdigit() or (text.startswith("-") and text[1:].isdigit()): + type_id = int(text) + else: + lower = text.lower() + if lower in by_var: + type_id = by_var[lower] + elif lower in by_full: + type_id = by_full[lower] + else: + matches = [ + int(t) + for t, meta in metadata.items() + if lower in str(meta.get("var_name", "")).lower() + or lower in str(meta.get("full_name", "")).lower() + ] + if len(matches) != 1: + raise ValueError( + f"--extra_info={text!r} matched {len(matches)} types; " + "use a type id or exact variable name." + ) + type_id = matches[0] + if type_id not in available: + raise ValueError( + f"extra-info type {type_id} is not available in this dataset/run" + ) + if type_id not in seen: + out.append(type_id) + seen.add(type_id) + return out + + +def death_distribution_parameters( + model, + hidden: torch.Tensor, + *, + dist_mode: str, + eps: float = 1e-8, +) -> tuple[str, torch.Tensor]: + logits = model.calc_risk(hidden) + death_idx = int(logits.shape[1]) - 1 + death_lambda = F.softplus(logits[:, death_idx]) + float(eps) + + if dist_mode == "exponential": + nan = torch.full_like(death_lambda, float("nan")) + return "exponential", torch.stack([death_lambda, nan, nan], dim=1) + + if dist_mode == "weibull": + rho = model.calc_weibull_rho(hidden)[:, death_idx].to(dtype=death_lambda.dtype) + elif dist_mode == "mixed": + rho = model.calc_death_rho(hidden).to(dtype=death_lambda.dtype) + else: + raise ValueError(f"Unsupported dist_mode={dist_mode!r}") + + shape = rho.clamp_min(float(eps)) + scale = torch.pow(death_lambda.clamp_min(float(eps)), -1.0 / shape) + nan = torch.full_like(death_lambda, float("nan")) + return "weibull", torch.stack([nan, scale, shape], dim=1) + + +def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor: + return torch.stack( + [ + original[:, 0], + ablated[:, 0], + original[:, 1], + ablated[:, 1], + original[:, 2], + ablated[:, 2], + ], + dim=1, + ) + + +def disease_probabilities( + model, + hidden: torch.Tensor, + *, + dist_mode: str, + tau: float, +) -> torch.Tensor: + logits = model.calc_risk(hidden) + rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None + death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None + return probabilities_from_logits( + logits, + tau, + dist_mode=dist_mode, + rho=rho, + death_rho=death_rho, + ) + + +def build_extra_info_ablated_slice( + batch: dict[str, torch.Tensor], + *, + row_indices: torch.Tensor, + extra_info_type_id: int, +) -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + repeated_keys = ( + "event_seq", + "time_seq", + "padding_mask", + "readout_mask", + "sex", + "landmark_pos", + "t_query", + "patient_id", + "landmark_age", + "followup_end_time", + "death_time", + "row_idx", + ) + for key in repeated_keys: + out[key] = batch[key][row_indices] + + out["other_type"] = batch["other_type"][row_indices].clone() + out["other_value"] = batch["other_value"][row_indices].clone() + out["other_value_kind"] = batch["other_value_kind"][row_indices].clone() + out["other_time"] = batch["other_time"][row_indices].clone() + + remove = out["other_type"] == int(extra_info_type_id) + out["other_type"][remove] = 0 + out["other_value"][remove] = 0 + out["other_value_kind"][remove] = 0 + out["other_time"][remove] = 0 + return out + + +def concat_tensor_batches(chunks: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]: + return {key: torch.cat([chunk[key] for chunk in chunks], dim=0) for key in chunks[0]} + + +def iter_extra_info_ablated_batches( + batch: dict[str, torch.Tensor], + *, + selected_extra_info_types: Sequence[int], + max_batch_size: int, +): + pending_batches: list[dict[str, torch.Tensor]] = [] + pending_types: list[int] = [] + pending_rows: list[int] = [] + pending_n = 0 + + other_type = batch["other_type"] + visible = other_type > 0 + visible &= batch["other_time"] <= batch["t_query"][:, None].to(batch["other_time"].dtype) + + for type_id in selected_extra_info_types: + active_rows = torch.nonzero( + ((other_type == int(type_id)) & visible).any(dim=1), + as_tuple=False, + ).flatten() + if active_rows.numel() == 0: + continue + + row_offset = 0 + while row_offset < int(active_rows.numel()): + capacity = int(max_batch_size) - pending_n + row_stop = min(int(active_rows.numel()), row_offset + capacity) + row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device) + chunk = build_extra_info_ablated_slice( + batch, + row_indices=row_indices, + extra_info_type_id=int(type_id), + ) + chunk_n = int(row_indices.numel()) + pending_batches.append(chunk) + pending_types.extend([int(type_id)] * chunk_n) + pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist()) + pending_n += chunk_n + row_offset = row_stop + + if pending_n >= int(max_batch_size): + yield concat_tensor_batches(pending_batches), pending_types, pending_rows + pending_batches = [] + pending_types = [] + pending_rows = [] + pending_n = 0 + + if pending_batches: + yield concat_tensor_batches(pending_batches), pending_types, pending_rows + + +def update_death_summary( + summary: dict[tuple[Any, ...], dict[str, float]], + *, + key_rows: pd.DataFrame, + values: np.ndarray, +) -> None: + if key_rows.empty: + return + table = key_rows.copy() + for idx, column in enumerate(DEATH_PARAMETER_COLUMNS): + table[column] = values[:, idx] + + for key, group in table.groupby(EXTRA_KEY_COLUMNS, dropna=False, sort=False): + if not isinstance(key, tuple): + key = (key,) + acc = summary.setdefault( + key, + { + "n": 0.0, + **{f"count__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS}, + **{f"sum__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS}, + **{f"sumsq__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS}, + }, + ) + acc["n"] += float(len(group)) + for column in DEATH_PARAMETER_COLUMNS: + vals = pd.to_numeric(group[column], errors="coerce").dropna() + acc[f"count__{column}"] += float(len(vals)) + acc[f"sum__{column}"] += float(vals.sum()) + acc[f"sumsq__{column}"] += float((vals * vals).sum()) + + +def update_disease_risk_summary( + summary: dict[tuple[Any, ...], dict[str, float]], + *, + key_rows: pd.DataFrame, + target_group: str, + target_group_label: str, + original_risk: np.ndarray, + ablated_risk: np.ndarray, +) -> None: + if key_rows.empty: + return + table = key_rows.copy() + table["target_group"] = str(target_group) + table["target_group_label"] = str(target_group_label) + table["original_future_disease_risk"] = original_risk + table["ablated_future_disease_risk"] = ablated_risk + + for key, group in table.groupby(DISEASE_RISK_KEY_COLUMNS, dropna=False, sort=False): + if not isinstance(key, tuple): + key = (key,) + acc = summary.setdefault( + key, + { + "n": 0.0, + **{f"sum__{col}": 0.0 for col in DISEASE_RISK_COLUMNS}, + **{f"sumsq__{col}": 0.0 for col in DISEASE_RISK_COLUMNS}, + }, + ) + acc["n"] += float(len(group)) + for column in DISEASE_RISK_COLUMNS: + vals = pd.to_numeric(group[column], errors="coerce") + acc[f"sum__{column}"] += float(vals.sum()) + acc[f"sumsq__{column}"] += float((vals * vals).sum()) + + +def write_death_summary_csv( + path: Path, + summary: dict[tuple[Any, ...], dict[str, float]], + *, + death_distribution: str, +) -> int: + rows: list[dict[str, Any]] = [] + for key, acc in summary.items(): + n = int(acc["n"]) + row = {column: value for column, value in zip(EXTRA_KEY_COLUMNS, key)} + row["n"] = n + row["death_distribution"] = death_distribution + for column in DEATH_PARAMETER_COLUMNS: + count = int(acc[f"count__{column}"]) + mean = acc[f"sum__{column}"] / count if count > 0 else np.nan + second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan + row[f"mean__{column}"] = mean + row[f"var__{column}"] = second - mean * mean if count > 0 else np.nan + rows.append(row) + columns = [ + *EXTRA_KEY_COLUMNS, + "n", + "death_distribution", + *[ + name + for column in DEATH_PARAMETER_COLUMNS + for name in (f"mean__{column}", f"var__{column}") + ], + ] + pd.DataFrame(rows, columns=columns).sort_values( + ["selected_extra_info_type_id", "landmark_age", "sex"], + kind="mergesort", + ).to_csv(path, index=False) + return len(rows) + + +def write_disease_risk_summary_csv( + path: Path, + summary: dict[tuple[Any, ...], dict[str, float]], + *, + tau: float, +) -> int: + rows: list[dict[str, Any]] = [] + for key, acc in summary.items(): + n = int(acc["n"]) + row = {column: value for column, value in zip(DISEASE_RISK_KEY_COLUMNS, key)} + row["n"] = n + row["tau_years"] = float(tau) + for column in DISEASE_RISK_COLUMNS: + mean = acc[f"sum__{column}"] / n if n > 0 else np.nan + second = acc[f"sumsq__{column}"] / n if n > 0 else np.nan + row[f"mean__{column}"] = mean + row[f"var__{column}"] = second - mean * mean if n > 0 else np.nan + rows.append(row) + columns = [ + *DISEASE_RISK_KEY_COLUMNS, + "n", + "tau_years", + *[ + name + for column in DISEASE_RISK_COLUMNS + for name in (f"mean__{column}", f"var__{column}") + ], + ] + pd.DataFrame(rows, columns=columns).sort_values( + ["selected_extra_info_type_id", "target_group", "landmark_age", "sex"], + kind="mergesort", + ).to_csv(path, index=False) + return len(rows) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Compute extra-info ablation attribution for death parameters and future disease risks." + ) + parser.add_argument("--run_path", type=str, required=True) + parser.add_argument( + "--extra_info", + type=str, + default=None, + help=( + "Optional type id, variable name, exact full name, or comma-separated list. " + "If omitted, scan all extra-info types available in the run." + ), + ) + parser.add_argument("--output_dir", type=str, default=None) + parser.add_argument("--organ_mapping_path", type=str, default="icd10_chapter_organ_mapping.csv") + parser.add_argument("--eval_split", type=str, default=None) + parser.add_argument("--dataset_subset_size", type=int, default=None) + parser.add_argument("--train_eid_file", type=str, default=None) + parser.add_argument("--val_eid_file", type=str, default=None) + parser.add_argument("--test_eid_file", 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("--tau", type=float, default=5.0) + parser.add_argument("--min_history_events", type=int, default=None) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument( + "--attribution_batch_size", + type=int, + default=None, + help="Forward batch size for expanded extra-info ablation queries.", + ) + parser.add_argument("--num_workers", type=int, default=None) + parser.add_argument("--device", type=str, default=None) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + # Dataset extra-info types must reproduce the checkpoint training config. + # --extra_info only filters which already-trained types are ablated. + args.extra_info_types = None + run_path = Path(args.run_path) + config_path = run_path / "train_config.json" + checkpoint_path = run_path / "best_model.pt" + if not config_path.exists(): + raise FileNotFoundError(f"train_config.json not found: {config_path}") + if not checkpoint_path.exists(): + raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}") + + cfg = load_json_config(config_path) + model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower() + if model_target_mode not in {"next_token", "all_future"}: + raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}") + + target_mode = str(cfg.get("target_mode", "uts")) + attn_mask_mode = str( + cfg.get("attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware") + ) + readout_name = str( + cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token") + ) + readout_reduce = str(cfg.get("readout_reduce", "mean")) + + dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(args, cfg) + validate_dataset_metadata(dataset, cfg) + + extra_metadata = load_extra_info_metadata( + dataset_extra_info_types=dataset.extra_info_types, + search_root=Path("."), + ) + selected_extra_info_types = resolve_extra_info_types( + args.extra_info, + dataset_extra_info_types=dataset.extra_info_types, + metadata=extra_metadata, + ) + if not selected_extra_info_types: + raise ValueError("No extra-info types selected for attribution") + + landmark_ages = make_landmark_ages( + float(args.landmark_start), + float(args.landmark_stop), + float(args.landmark_step), + ) + first_occurrence_by_token = build_first_occurrence_maps_for_landmarks( + dataset, + subset_indices, + ) + death_idx = int(dataset.vocab_size) - 1 + landmark_dataset = LandmarkDataset( + dataset=dataset, + subset_indices=subset_indices, + landmark_ages=landmark_ages, + attn_mask_mode=attn_mask_mode, + model_target_mode=model_target_mode, + min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)), + first_occurrence_by_token=first_occurrence_by_token, + death_token_ids=[death_idx], + ) + + organ_groups, organ_labels, _token_to_group = load_organ_groups( + Path(args.organ_mapping_path), + vocab_size=int(dataset.vocab_size), + ) + all_disease_tokens = sorted( + { + int(token) + for tokens in organ_groups.values() + for token in tokens + if int(token) != death_idx + } + ) + risk_groups = { + "all_modeled_diseases": all_disease_tokens, + **{group: tokens for group, tokens in sorted(organ_groups.items())}, + } + risk_group_labels = { + "all_modeled_diseases": "All modeled diseases", + **organ_labels, + } + + state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu") + dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict) + death_distribution_name = "exponential" if dist_mode == "exponential" else "weibull" + cfg_model = dict(cfg) + cfg_model["dist_mode"] = dist_mode + device = resolve_eval_device(args.device) + model = build_model_from_dataset(args, cfg_model, dataset).to(device) + load_model_state(model, state_dict) + model.eval() + + batch_size = int(cfg_get(args, cfg, "batch_size", 128)) + attribution_batch_size = int( + cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096)) + ) + if attribution_batch_size <= 0: + raise ValueError("attribution_batch_size must be positive") + + num_workers = int(cfg_get(args, cfg, "num_workers", 4)) + loader = DataLoader( + IndexedLandmarkDataset(landmark_dataset), + batch_size=batch_size, + shuffle=False, + collate_fn=collate_indexed_landmark_fn, + num_workers=num_workers, + pin_memory=device.type == "cuda", + persistent_workers=num_workers > 0, + prefetch_factor=2 if num_workers > 0 else None, + ) + + output_dir = ( + Path(args.output_dir) + if args.output_dir + else run_path / f"extra_info_attribution_{eval_split}_tau{float(args.tau):g}y" + ) + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"Eval split: {eval_split}") + print(f"Split source: {split_source}") + print(f"Selected patients: {len(subset_indices)}") + print(f"Landmark ages: {landmark_ages.tolist()}") + print(f"Dist mode: {dist_mode}") + print(f"Device: {device}") + print(f"Death token: {death_idx}") + print(f"Extra-info types: {selected_extra_info_types}") + print(f"Landmark rows: {len(landmark_dataset)}") + print(f"Attribution batch size: {attribution_batch_size}") + print(f"Output directory: {output_dir}") + + death_summary: dict[tuple[Any, ...], dict[str, float]] = {} + disease_risk_summary: dict[tuple[Any, ...], dict[str, float]] = {} + + for batch in tqdm(loader, desc="Extra-info attribution", dynamic_ncols=True): + batch_dev = { + k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) + for k, v in batch.items() + } + with torch.no_grad(): + hidden = infer_landmark_hidden( + model=model, + batch=batch_dev, + device=device, + model_target_mode=model_target_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + ) + _death_distribution, original_death_params = death_distribution_parameters( + model, + hidden, + dist_mode=dist_mode, + ) + original_probabilities = disease_probabilities( + model, + hidden, + dist_mode=dist_mode, + tau=float(args.tau), + ) + occurred = make_occurred_mask( + batch_dev["event_seq"].long(), + vocab_size=int(dataset.vocab_size), + device=device, + ) + original_risk_by_group = { + group: new_disease_risk_from_probabilities( + original_probabilities, + occurred, + tokens, + ) + for group, tokens in risk_groups.items() + } + + for ablated_batch, type_ids, local_rows in iter_extra_info_ablated_batches( + batch_dev, + selected_extra_info_types=selected_extra_info_types, + max_batch_size=attribution_batch_size, + ): + row_tensor = torch.as_tensor(local_rows, dtype=torch.long, device=device) + with torch.no_grad(): + ablated_hidden = infer_landmark_hidden( + model=model, + batch=ablated_batch, + device=device, + model_target_mode=model_target_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + ) + _ablated_distribution, ablated_death_params = death_distribution_parameters( + model, + ablated_hidden, + dist_mode=dist_mode, + ) + ablated_probabilities = disease_probabilities( + model, + ablated_hidden, + dist_mode=dist_mode, + tau=float(args.tau), + ) + + key_rows = [] + for type_id, local_row in zip(type_ids, local_rows): + meta = extra_metadata[int(type_id)] + key_rows.append( + { + "selected_extra_info_type_id": int(type_id), + "selected_extra_info_var_name": str(meta.get("var_name", "")), + "selected_extra_info_full_name": str(meta.get("full_name", "")), + "landmark_age": float(batch["landmark_age"][int(local_row)].item()), + "sex": int(batch["sex"][int(local_row)].item()), + } + ) + key_table = pd.DataFrame(key_rows, columns=EXTRA_KEY_COLUMNS) + value_block = parameter_pair_block( + original_death_params[row_tensor], + ablated_death_params, + ).detach().cpu().numpy() + update_death_summary( + death_summary, + key_rows=key_table, + values=value_block, + ) + + ablated_occurred = occurred[row_tensor] + for group, tokens in risk_groups.items(): + ablated_risk = new_disease_risk_from_probabilities( + ablated_probabilities, + ablated_occurred, + tokens, + ) + update_disease_risk_summary( + disease_risk_summary, + key_rows=key_table, + target_group=group, + target_group_label=risk_group_labels[group], + original_risk=original_risk_by_group[group][row_tensor].detach().cpu().numpy(), + ablated_risk=ablated_risk.detach().cpu().numpy(), + ) + + death_summary_path = output_dir / "summary_extra_info_death_parameters.csv" + disease_summary_path = output_dir / "summary_extra_info_future_disease_risk.csv" + death_rows = write_death_summary_csv( + death_summary_path, + death_summary, + death_distribution=death_distribution_name, + ) + disease_rows = write_disease_risk_summary_csv( + disease_summary_path, + disease_risk_summary, + tau=float(args.tau), + ) + manifest = { + "death_summary_file": death_summary_path.name, + "disease_risk_summary_file": disease_summary_path.name, + "death_summary_rows": int(death_rows), + "disease_risk_summary_rows": int(disease_rows), + "eval_split": eval_split, + "split_source": split_source, + "dist_mode": dist_mode, + "tau_years": float(args.tau), + "landmark_start": float(args.landmark_start), + "landmark_stop": float(args.landmark_stop), + "landmark_step": float(args.landmark_step), + "selected_extra_info_types": [ + extra_metadata[int(type_id)] for type_id in selected_extra_info_types + ], + } + with (output_dir / "manifest.json").open("w", encoding="utf-8") as f: + json.dump(manifest, f, ensure_ascii=False, indent=2) + + print(f"Wrote {death_rows} death summary rows to {death_summary_path}") + print(f"Wrote {disease_rows} disease-risk summary rows to {disease_summary_path}") + + +if __name__ == "__main__": + main() diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py index 61ff7fc..1a111b0 100644 --- a/evaluate_single_disease_mortality_attribution.py +++ b/evaluate_single_disease_mortality_attribution.py @@ -32,7 +32,7 @@ from evaluate_auc_v2 import ( resolve_eval_device, validate_dataset_metadata, ) -from evaluate_event_free_survival import ( +from landmark_eval_utils import ( IndexedLandmarkDataset, LandmarkDataset, build_first_occurrence_maps_for_landmarks, diff --git a/landmark_eval_utils.py b/landmark_eval_utils.py new file mode 100644 index 0000000..5e91607 --- /dev/null +++ b/landmark_eval_utils.py @@ -0,0 +1,511 @@ +"""Shared landmark evaluation helpers for attribution scripts.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence + +import numpy as np +import pandas as pd +import torch +from torch.nn.utils.rnn import pad_sequence +from torch.utils.data import Dataset + +from dataset import HealthDataset +from eval_data import load_sequence_eval_dataset +from evaluate_auc_v2 import ( + LandmarkDataset, + build_model_from_dataset, + cfg_get, + make_eval_indices, +) +from models import DeepHealth +from readouts import build_readout +from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX +from train_util import load_eid_file, load_extra_info_types_file + + +SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX} + + +def parse_int_list(value: Any) -> Optional[List[int]]: + if value is None: + return None + if isinstance(value, (list, tuple, np.ndarray)): + return [int(x) for x in value] + text = str(value).strip() + if text == "": + return None + if text.startswith("["): + values = json.loads(text) + if not isinstance(values, list): + raise ValueError(f"Expected a JSON list, got {type(values).__name__}") + return [int(x) for x in values] + return [int(x.strip()) for x in text.split(",") if x.strip()] + + +def load_extra_info_types(value: Any) -> Optional[List[int]]: + if value is None: + return None + text = str(value) + path = Path(text) + if path.exists(): + return load_extra_info_types_file(text) + return parse_int_list(value) + + +def make_landmark_ages(start: float, stop: float, step: float) -> np.ndarray: + if step <= 0: + raise ValueError("landmark_step must be positive") + if stop < start: + raise ValueError("landmark_stop must be >= landmark_start") + # Include stop when it lands on the grid, e.g. 40,45,...,80. + return np.arange(start, stop + step * 0.5, step, dtype=np.float32) + + +def build_first_occurrence_maps_for_landmarks( + dataset: HealthDataset, + subset_indices: np.ndarray, +) -> Dict[int, tuple[np.ndarray, np.ndarray]]: + first_lists: Dict[int, list[tuple[int, float]]] = {} + for patient_id, dataset_index in enumerate(np.asarray(subset_indices, dtype=np.int64).tolist()): + s = dataset.samples[int(dataset_index)] + seq_event = np.asarray(s["event_seq"], dtype=np.int64) + seq_time = np.asarray(s["time_seq"], dtype=np.float32) + tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64) + tgt_time = np.asarray(s["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:]]) + uniq_tokens, first_idx = np.unique(full_event, return_index=True) + for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()): + token = int(token) + if token in SPECIAL_TOKENS: + continue + first_lists.setdefault(token, []).append((patient_id, float(full_time[int(idx)]))) + + return { + int(token): ( + np.asarray([p for p, _ in pairs], dtype=np.int32), + np.asarray([t for _, t in pairs], dtype=np.float32), + ) + for token, pairs in first_lists.items() + if pairs + } + + +def normalize_eval_split(args: argparse.Namespace, cfg: Dict[str, Any]) -> str: + eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower() + if eval_split in {"valid", "validation"}: + return "val" + if eval_split not in {"train", "val", "test", "all"}: + raise ValueError(f"Unsupported eval_split={eval_split!r}") + return eval_split + + +def load_eval_sequence_dataset( + args: argparse.Namespace, + cfg: Dict[str, Any], +) -> tuple[Any, np.ndarray, str, str]: + eval_split = normalize_eval_split(args, cfg) + model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower() + data_prefix = str(cfg.get("data_prefix", "ukb")) + labels_file = str(cfg.get("labels_file", "labels.csv")) + no_event_interval_years = float(cfg.get("no_event_interval_years", 5.0)) + include_no_event_in_uts_target = bool(cfg.get("include_no_event_in_uts_target", False)) + extra_info_types = load_extra_info_types(args.extra_info_types) + if extra_info_types is None: + extra_info_types = parse_int_list(cfg.get("extra_info_types", None)) + + print("Loading one sequence eval dataset...") + dataset = load_sequence_eval_dataset( + model_target_mode=model_target_mode, + data_prefix=data_prefix, + labels_file=labels_file, + no_event_interval_years=no_event_interval_years, + include_no_event_in_uts_target=include_no_event_in_uts_target, + min_history_events=int(cfg.get("all_future_min_history_events", 1)), + min_future_events=int(cfg.get("all_future_min_future_events", 1)), + extra_info_types=extra_info_types, + ) + + train_eid_file = cfg_get(args, cfg, "train_eid_file", "ukb_train_eid.csv") + val_eid_file = cfg_get(args, cfg, "val_eid_file", "ukb_val_eid.csv") + test_eid_file = cfg_get(args, cfg, "test_eid_file", "ukb_test_eid.csv") + split_files_exist = all( + Path(str(path)).exists() + for path in (train_eid_file, val_eid_file, test_eid_file) + ) + + if eval_split != "all" and split_files_exist: + split_files = { + "train": train_eid_file, + "val": val_eid_file, + "test": test_eid_file, + } + selected_eids = load_eid_file(split_files[eval_split]) + out = np.asarray( + [ + idx + for idx, sample in enumerate(dataset.samples) + if int(sample["eid"]) in selected_eids + ], + dtype=np.int64, + ) + if out.size == 0: + raise ValueError( + f"No samples found for eval_split={eval_split!r} using {split_files[eval_split]}" + ) + split_source = "eid_files" + else: + if eval_split == "all": + out = np.arange(len(dataset.samples), dtype=np.int64) + split_source = "all" + else: + out = make_eval_indices(dataset, args, cfg) + split_source = "ratio_split" + + subset_size = cfg_get(args, cfg, "dataset_subset_size", None) + if subset_size is not None and int(subset_size) > 0: + out = out[: int(subset_size)] + return dataset, np.asarray(out, dtype=np.int64), eval_split, split_source + + +def load_organ_groups( + path: Path, + *, + vocab_size: int, +) -> tuple[dict[str, list[int]], dict[str, str], dict[int, str]]: + table = pd.read_csv(path) + required = {"token_id", "organ_system", "organ_system_label", "is_death"} + missing = required - set(table.columns) + if missing: + raise ValueError(f"{path} is missing columns: {sorted(missing)}") + + death_idx = int(vocab_size) - 1 + groups: dict[str, list[int]] = {} + labels: dict[str, str] = {} + token_to_group: dict[int, str] = {} + for row in table.itertuples(index=False): + token = int(getattr(row, "token_id")) + if token in SPECIAL_TOKENS or token == death_idx: + continue + if token < 0 or token >= int(vocab_size): + continue + if int(getattr(row, "is_death")) == 1: + continue + group = str(getattr(row, "organ_system")) + label = str(getattr(row, "organ_system_label")) + groups.setdefault(group, []).append(token) + labels[group] = label + token_to_group[token] = group + + groups = {k: sorted(set(v)) for k, v in groups.items() if v} + return groups, labels, token_to_group + + +class IndexedLandmarkDataset(Dataset): + def __init__(self, base: LandmarkDataset) -> None: + self.base = base + + def __len__(self) -> int: + return len(self.base) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + item = dict(self.base[idx]) + item["row_idx"] = torch.tensor(int(idx), dtype=torch.long) + return item + + +def collate_indexed_landmark_fn(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 + ) + time_seq = pad_sequence( + [x["time_seq"] for x in batch], batch_first=True, padding_value=0.0 + ) + readout_mask = pad_sequence( + [x["readout_mask"] for x in batch], batch_first=True, padding_value=False + ) + 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 + ) + return { + "event_seq": event_seq, + "time_seq": time_seq, + "padding_mask": event_seq > PAD_IDX, + "readout_mask": readout_mask, + "sex": torch.stack([x["sex"] for x in batch]), + "other_type": other_type, + "other_value": other_value, + "other_value_kind": other_value_kind, + "other_time": other_time, + "landmark_pos": torch.stack([x["landmark_pos"] for x in batch]), + "t_query": torch.stack([x["t_query"] for x in batch]), + "patient_id": torch.stack([x["patient_id"] for x in batch]), + "landmark_age": torch.stack([x["landmark_age"] for x in batch]), + "followup_end_time": torch.stack([x["followup_end_time"] for x in batch]), + "death_time": torch.stack([x["death_time"] for x in batch]), + "row_idx": torch.stack([x["row_idx"] for x in batch]), + } + + +def build_group_ablated_slice( + batch: Dict[str, torch.Tensor], + token_ids: Sequence[int], + row_indices: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """Build one fixed-width ablated slice without rebuilding variable-length rows.""" + event_seq = batch["event_seq"] + + out: Dict[str, torch.Tensor] = {} + out["event_seq"] = event_seq[row_indices].clone() + out["time_seq"] = batch["time_seq"][row_indices] + out["readout_mask"] = batch["readout_mask"][row_indices].clone() + out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone() + out["landmark_pos"] = batch["landmark_pos"][row_indices].clone() + + seq_len = int(event_seq.shape[1]) + positions = torch.arange(seq_len, device=event_seq.device)[None, :] + ids = torch.as_tensor(token_ids, dtype=event_seq.dtype, device=event_seq.device) + remove = torch.isin(out["event_seq"], ids) & out["padding_mask"] + out["event_seq"] = torch.where( + remove, + torch.full_like(out["event_seq"], PAD_IDX), + out["event_seq"], + ) + out["padding_mask"] &= ~remove + out["readout_mask"] &= ~remove + + has_valid = out["padding_mask"].any(dim=1) + if not bool(has_valid.all().item()): + empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten() + out["event_seq"][empty_rows, 0] = CHECKUP_IDX + out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to( + dtype=out["time_seq"].dtype + ) + out["padding_mask"][empty_rows, 0] = True + out["readout_mask"][empty_rows, 0] = True + out["landmark_pos"][empty_rows] = 0 + + has_readout = out["readout_mask"].any(dim=1) + if not bool(has_readout.all().item()): + rows = torch.nonzero(~has_readout, as_tuple=False).flatten() + local_valid = out["padding_mask"][rows] + last_pos = torch.where( + local_valid, + positions.expand(local_valid.shape[0], -1), + torch.zeros_like(positions.expand(local_valid.shape[0], -1)), + ).amax(dim=1) + out["readout_mask"][rows] = False + out["readout_mask"][rows, last_pos] = True + out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype) + + repeated_keys = ( + "sex", + "other_type", + "other_value", + "other_value_kind", + "other_time", + "t_query", + "patient_id", + "landmark_age", + "followup_end_time", + "death_time", + "row_idx", + ) + for key in repeated_keys: + out[key] = batch[key][row_indices] + return out + + +def concat_tensor_batches(chunks: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: + return { + key: torch.cat([chunk[key] for chunk in chunks], dim=0) + for key in chunks[0] + } + + +def iter_group_ablated_batches( + batch: Dict[str, torch.Tensor], + group_names: Sequence[str], + organ_groups: dict[str, list[int]], + occurred: torch.Tensor, + max_batch_size: int, +): + """Yield ablated chunks as soon as enough rows are available for a forward pass.""" + pending_batches: list[Dict[str, torch.Tensor]] = [] + pending_groups: list[str] = [] + pending_rows: list[int] = [] + pending_n = 0 + + for group in group_names: + ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device) + if ids.numel() == 0: + continue + active_rows = torch.nonzero(occurred[:, ids].any(dim=1), as_tuple=False).flatten() + if active_rows.numel() == 0: + continue + + row_offset = 0 + while row_offset < int(active_rows.numel()): + capacity = int(max_batch_size) - pending_n + row_stop = min(int(active_rows.numel()), row_offset + capacity) + row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device) + chunk = build_group_ablated_slice( + batch=batch, + token_ids=organ_groups[group], + row_indices=row_indices, + ) + chunk_n = int(row_indices.numel()) + pending_batches.append(chunk) + pending_groups.extend([group] * chunk_n) + pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist()) + pending_n += chunk_n + row_offset = row_stop + + if pending_n >= int(max_batch_size): + yield concat_tensor_batches(pending_batches), pending_groups, pending_rows + pending_batches = [] + pending_groups = [] + pending_rows = [] + pending_n = 0 + + if pending_batches: + yield concat_tensor_batches(pending_batches), pending_groups, pending_rows + + +@torch.no_grad() +def infer_landmark_hidden( + *, + model: DeepHealth, + batch: Dict[str, torch.Tensor], + device: torch.device, + model_target_mode: str, + readout_name: str, + readout_reduce: str, +) -> torch.Tensor: + batch_dev = { + k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) + for k, v in batch.items() + } + if model_target_mode == "all_future": + return model( + event_seq=batch_dev["event_seq"].long(), + time_seq=batch_dev["time_seq"].float(), + sex=batch_dev["sex"].long(), + padding_mask=batch_dev["padding_mask"].bool(), + t_query=batch_dev["t_query"].float(), + other_type=batch_dev["other_type"].long(), + other_value=batch_dev["other_value"].float(), + other_value_kind=batch_dev["other_value_kind"].long(), + other_time=batch_dev["other_time"].float(), + target_mode="all_future", + ) + + hidden = model( + event_seq=batch_dev["event_seq"].long(), + time_seq=batch_dev["time_seq"].float(), + sex=batch_dev["sex"].long(), + padding_mask=batch_dev["padding_mask"].bool(), + other_type=batch_dev["other_type"].long(), + other_value=batch_dev["other_value"].float(), + other_value_kind=batch_dev["other_value_kind"].long(), + other_time=batch_dev["other_time"].float(), + target_mode="next_token", + ) + readout = build_readout(readout_name, reduce=readout_reduce) + readout_out = readout( + hidden=hidden, + time_seq=batch_dev["time_seq"].float(), + padding_mask=batch_dev["padding_mask"].bool(), + readout_mask=batch_dev["readout_mask"].bool(), + ) + return readout_out.hidden.gather( + 1, + batch_dev["landmark_pos"].long()[:, None, None].expand( + -1, 1, readout_out.hidden.shape[-1] + ), + ).squeeze(1) + + +def make_occurred_mask( + event_seq: torch.Tensor, + *, + vocab_size: int, + device: torch.device, +) -> torch.Tensor: + occurred = torch.zeros(event_seq.shape[0], int(vocab_size), dtype=torch.bool, device=device) + valid = (event_seq >= 0) & (event_seq < int(vocab_size)) + safe = event_seq.clamp(min=0, max=int(vocab_size) - 1).to(device) + occurred.scatter_(1, safe, valid.to(device)) + return occurred + + +def mortality_hazard_from_risk(risk: torch.Tensor, eps: float = 1e-7) -> torch.Tensor: + return -torch.log1p(-risk.clamp(0.0, 1.0 - float(eps))) + + +def death_risk_for_batch( + *, + model: DeepHealth, + batch: Dict[str, torch.Tensor], + device: torch.device, + model_target_mode: str, + readout_name: str, + readout_reduce: str, + dist_mode: str, + tau: float, +) -> torch.Tensor: + hidden = infer_landmark_hidden( + model=model, + batch=batch, + device=device, + model_target_mode=model_target_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + ) + logits = model.calc_risk(hidden) + rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None + death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None + probabilities = probabilities_from_logits( + logits, + tau, + dist_mode=dist_mode, + rho=rho, + death_rho=death_rho, + ) + return death_risk_from_probabilities(probabilities) + + +def historical_counts_by_group( + tokens: np.ndarray, + *, + death_idx: int, + token_to_group: dict[int, str], + group_names: Sequence[str], +) -> tuple[int, dict[str, int]]: + unique_tokens = { + int(token) + for token in np.asarray(tokens, dtype=np.int64).tolist() + if int(token) not in SPECIAL_TOKENS and int(token) != int(death_idx) + } + total = len(unique_tokens) + out = {group: 0 for group in group_names} + for token in unique_tokens: + group = token_to_group.get(token) + if group in out: + out[group] += 1 + return total, out