Add t-query logits and hidden export script

export_tquery_logits_hidden.py

@@ -0,0 +1,325 @@
+"""Export landmark risk logits and hidden states for t_query ages.
+
+This script follows evaluate_event_free_survival.py's data loading,
+landmark construction, checkpoint loading, and readout logic, but only exports:
+
+* all token/disease risk logits from ``model.calc_risk(hidden)``;
+* the corresponding landmark hidden state.
+
+The two large arrays are saved separately as .npy files. Row metadata is saved
+as a CSV with matching row order.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+
+from evaluate_auc_v2 import (
+    LandmarkDataset,
+    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 evaluate_event_free_survival import (
+    IndexedLandmarkDataset,
+    build_first_occurrence_maps_for_landmarks,
+    collate_indexed_landmark_fn,
+    infer_landmark_hidden,
+    load_eval_sequence_dataset,
+    make_landmark_ages,
+)
+
+
+def numpy_float_dtype(name: str) -> np.dtype:
+    key = str(name).lower()
+    if key in {"float16", "fp16", "half"}:
+        return np.dtype(np.float16)
+    if key in {"float32", "fp32", "single"}:
+        return np.dtype(np.float32)
+    raise ValueError(f"dtype must be float16 or float32, got {name!r}")
+
+
+def output_paths_for_run(
+    run_path: Path,
+    eval_split: str,
+    landmark_start: float,
+    landmark_stop: float,
+    landmark_step: float,
+    output_dir: Optional[str],
+) -> tuple[Path, Path, Path, Path]:
+    suffix = f"{eval_split}_t{landmark_start:g}-{landmark_stop:g}_step{landmark_step:g}"
+    base_dir = Path(output_dir) if output_dir else run_path
+    return (
+        base_dir / f"tquery_logits_{suffix}.npy",
+        base_dir / f"tquery_hidden_{suffix}.npy",
+        base_dir / f"tquery_metadata_{suffix}.csv",
+        base_dir / f"tquery_export_config_{suffix}.json",
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Export landmark risk logits and hidden states for t_query ages."
+    )
+    parser.add_argument("--run_path", type=str, required=True)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="Directory for output files. Defaults to run_path.",
+    )
+    parser.add_argument("--logits_path", type=str, default=None)
+    parser.add_argument("--hidden_path", type=str, default=None)
+    parser.add_argument("--metadata_path", type=str, default=None)
+    parser.add_argument("--export_config_path", type=str, default=None)
+    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=1.0,
+        help="t_query grid step in years. Default exports every integer age 40..80.",
+    )
+    parser.add_argument("--min_history_events", type=int, default=None)
+    parser.add_argument("--batch_size", type=int, default=None)
+    parser.add_argument("--num_workers", type=int, default=None)
+    parser.add_argument("--device", type=str, default=None)
+    parser.add_argument("--extra_info_types", type=str, default=None)
+    parser.add_argument(
+        "--logits_dtype",
+        type=str,
+        default="float32",
+        choices=["float16", "float32"],
+    )
+    parser.add_argument(
+        "--hidden_dtype",
+        type=str,
+        default="float32",
+        choices=["float16", "float32"],
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    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)
+
+    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],
+    )
+
+    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,
+    )
+    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()
+
+    default_logits_path, default_hidden_path, default_metadata_path, default_config_path = (
+        output_paths_for_run(
+            run_path=run_path,
+            eval_split=eval_split,
+            landmark_start=float(args.landmark_start),
+            landmark_stop=float(args.landmark_stop),
+            landmark_step=float(args.landmark_step),
+            output_dir=args.output_dir,
+        )
+    )
+    logits_path = Path(args.logits_path) if args.logits_path else default_logits_path
+    hidden_path = Path(args.hidden_path) if args.hidden_path else default_hidden_path
+    metadata_path = Path(args.metadata_path) if args.metadata_path else default_metadata_path
+    export_config_path = (
+        Path(args.export_config_path) if args.export_config_path else default_config_path
+    )
+    for path in (logits_path, hidden_path, metadata_path, export_config_path):
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+    n_rows = len(landmark_dataset)
+    vocab_size = int(dataset.vocab_size)
+    hidden_dim = int(getattr(model, "n_embd", cfg_get(args, cfg_model, "n_embd", 120)))
+    logits_dtype = numpy_float_dtype(args.logits_dtype)
+    hidden_dtype = numpy_float_dtype(args.hidden_dtype)
+
+    logits_memmap = np.lib.format.open_memmap(
+        logits_path,
+        mode="w+",
+        dtype=logits_dtype,
+        shape=(n_rows, vocab_size),
+    )
+    hidden_memmap = np.lib.format.open_memmap(
+        hidden_path,
+        mode="w+",
+        dtype=hidden_dtype,
+        shape=(n_rows, hidden_dim),
+    )
+
+    batch_size = int(cfg_get(args, cfg, "batch_size", 128))
+    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,
+    )
+
+    print(f"Eval split: {eval_split}")
+    print(f"Split source: {split_source}")
+    print(f"Selected patients: {len(subset_indices)}")
+    print(f"t_query ages: {landmark_ages.tolist()}")
+    print(f"Dist mode: {dist_mode}")
+    print(f"Device: {device}")
+    print(f"Landmark rows: {n_rows}")
+    print(f"Logits: {logits_path} shape={(n_rows, vocab_size)} dtype={logits_dtype}")
+    print(f"Hidden: {hidden_path} shape={(n_rows, hidden_dim)} dtype={hidden_dtype}")
+    print(f"Metadata: {metadata_path}")
+
+    meta_rows: list[dict[str, Any]] = []
+    written = 0
+    with torch.no_grad():
+        for batch in tqdm(loader, desc="Export logits/hidden", dynamic_ncols=True):
+            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)
+            row_indices = batch["row_idx"].detach().cpu().numpy().astype(np.int64)
+            if not np.array_equal(row_indices, np.arange(written, written + len(row_indices))):
+                raise RuntimeError("DataLoader row order changed; export requires shuffle=False.")
+
+            batch_n = int(logits.shape[0])
+            logits_memmap[written : written + batch_n] = (
+                logits.detach().cpu().numpy().astype(logits_dtype, copy=False)
+            )
+            hidden_memmap[written : written + batch_n] = (
+                hidden.detach().cpu().numpy().astype(hidden_dtype, copy=False)
+            )
+
+            for row_idx in row_indices.tolist():
+                meta = landmark_dataset.rows[int(row_idx)]
+                sample = dataset.samples[int(meta["dataset_index"])]
+                meta_rows.append(
+                    {
+                        "row_index": int(row_idx),
+                        "patient_id": int(meta["patient_id"]),
+                        "dataset_index": int(meta["dataset_index"]),
+                        "eid": int(sample.get("eid", -1)),
+                        "sex": int(meta["sex"]),
+                        "t_query": float(meta["t_query"]),
+                        "landmark_age": float(meta["landmark_age"]),
+                        "followup_end_time": float(meta["followup_end_time"]),
+                        "death_time": float(meta["death_time"]),
+                    }
+                )
+            written += batch_n
+
+    logits_memmap.flush()
+    hidden_memmap.flush()
+    pd.DataFrame(meta_rows).to_csv(metadata_path, index=False)
+
+    export_config = {
+        "run_path": str(run_path),
+        "eval_split": eval_split,
+        "split_source": split_source,
+        "model_target_mode": model_target_mode,
+        "target_mode": target_mode,
+        "attn_mask_mode": attn_mask_mode,
+        "readout_name": readout_name,
+        "readout_reduce": readout_reduce,
+        "dist_mode": dist_mode,
+        "landmark_ages": [float(x) for x in landmark_ages.tolist()],
+        "n_rows": int(n_rows),
+        "vocab_size": int(vocab_size),
+        "hidden_dim": int(hidden_dim),
+        "death_token": int(death_idx),
+        "logits_path": str(logits_path),
+        "hidden_path": str(hidden_path),
+        "metadata_path": str(metadata_path),
+        "logits_dtype": str(logits_dtype),
+        "hidden_dtype": str(hidden_dtype),
+    }
+    with export_config_path.open("w", encoding="utf-8") as f:
+        json.dump(export_config, f, indent=2)
+
+    print(f"Wrote {written} rows.")
+    print(f"Wrote export config: {export_config_path}")
+
+
+if __name__ == "__main__":
+    main()