Add landmark attribution dependencies

.gitattributes

@@ -1 +1,2 @@
 *.sh text eol=lf
+*.py text eol=lf

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()

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,

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