"""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