DeepHealth/evaluate_single_disease_mortality_attribution.py

"""Compute per-disease attribution to predicted mortality distribution parameters.

For each selected patient and landmark age, this script keeps only rows where
each scanned disease token has already occurred in the history. It then deletes
that historical disease token, re-queries the model, and reports the original
and ablated fitted death distribution parameters. If --disease is omitted, all
disease tokens in the mapping are scanned.

Death is always token vocab_size - 1.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict

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 evaluate_event_free_survival import (
    IndexedLandmarkDataset,
    LandmarkDataset,
    build_first_occurrence_maps_for_landmarks,
    collate_indexed_landmark_fn,
    historical_counts_by_group,
    infer_landmark_hidden,
    load_eval_sequence_dataset,
    load_organ_groups,
    make_landmark_ages,
)
from targets import CHECKUP_IDX, PAD_IDX


OUTPUT_COLUMNS = [
    "patient_id",
    "dataset_index",
    "eid",
    "sex",
    "landmark_age",
    "followup_end_time",
    "history_disease_count",
    "selected_disease_history_count",
    "selected_disease_token_id",
    "selected_disease_code",
    "selected_disease_name",
    "selected_disease_organ_system",
    "selected_disease_organ_system_label",
    "history_count__selected_organ_system",
    "death_distribution",
    "original_death_lambda",
    "ablated_death_lambda",
    "original_death_scale",
    "ablated_death_scale",
    "original_death_shape",
    "ablated_death_shape",
]

SUMMARY_KEY_COLUMNS = [
    "selected_disease_token_id",
    "selected_disease_code",
    "selected_disease_name",
    "selected_disease_organ_system",
    "selected_disease_organ_system_label",
    "landmark_age",
    "sex",
]

SUMMARY_MEAN_COLUMNS = [
    "history_disease_count",
    "selected_disease_history_count",
    "history_count__selected_organ_system",
]

SUMMARY_PARAMETER_COLUMNS = [
    "original_death_lambda",
    "ablated_death_lambda",
    "original_death_scale",
    "ablated_death_scale",
    "original_death_shape",
    "ablated_death_shape",
]


def write_compressed_npz_table(path: Path, table: pd.DataFrame) -> int:
    table = table.reindex(columns=OUTPUT_COLUMNS)
    arrays: dict[str, np.ndarray] = {
        "__columns__": np.asarray(OUTPUT_COLUMNS, dtype="U"),
    }
    for column in OUTPUT_COLUMNS:
        values = table[column] if column in table else pd.Series([], dtype=object)
        if values.dtype == object:
            arrays[column] = values.fillna("").astype(str).to_numpy(dtype="U")
        else:
            arrays[column] = values.to_numpy()
    np.savez_compressed(path, **arrays)
    return int(len(table))


def normalize_output_dir(path: Path) -> Path:
    if path.suffix:
        return path.with_suffix(path.suffix + "_shards")
    return path


def write_manifest(
    output_dir: Path,
    *,
    rows: int,
    shards: list[dict[str, Any]],
    summary_file: str,
    scanned_diseases: list[dict[str, Any]],
    eval_split: str,
    dist_mode: str,
    landmark_start: float,
    landmark_stop: float,
    landmark_step: float,
) -> None:
    payload = {
        "format": "compressed_npz_shards",
        "columns": OUTPUT_COLUMNS,
        "rows": int(rows),
        "shards": shards,
        "summary_file": summary_file,
        "scanned_diseases": scanned_diseases,
        "eval_split": eval_split,
        "dist_mode": str(dist_mode),
        "landmark_start": float(landmark_start),
        "landmark_stop": float(landmark_stop),
        "landmark_step": float(landmark_step),
    }
    with (output_dir / "manifest.json").open("w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)


def update_summary_accumulator(
    summary: dict[tuple[Any, ...], dict[str, float]],
    table: pd.DataFrame,
) -> None:
    if table.empty:
        return
    grouped = table.groupby(SUMMARY_KEY_COLUMNS, dropna=False, sort=False)
    for key, group in grouped:
        if not isinstance(key, tuple):
            key = (key,)
        acc = summary.setdefault(
            key,
            {
                "n": 0.0,
                **{column: 0.0 for column in SUMMARY_MEAN_COLUMNS},
                **{f"count__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
                **{f"sum__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
                **{f"sumsq__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
            },
        )
        n = int(len(group))
        acc["n"] += float(n)
        for column in SUMMARY_MEAN_COLUMNS:
            acc[column] += float(pd.to_numeric(group[column], errors="coerce").sum())
        for column in SUMMARY_PARAMETER_COLUMNS:
            values = pd.to_numeric(group[column], errors="coerce").dropna()
            acc[f"count__{column}"] += float(len(values))
            acc[f"sum__{column}"] += float(values.sum())
            acc[f"sumsq__{column}"] += float((values * values).sum())


def write_summary_csv(
    path: Path,
    summary: dict[tuple[Any, ...], dict[str, float]],
) -> int:
    rows: list[dict[str, Any]] = []
    for key, acc in summary.items():
        n = int(acc["n"])
        out = {column: value for column, value in zip(SUMMARY_KEY_COLUMNS, key)}
        out["n"] = n
        for column in SUMMARY_MEAN_COLUMNS:
            out[f"mean__{column}"] = acc[column] / n if n > 0 else np.nan
        for column in SUMMARY_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
            out[f"mean__{column}"] = mean
            out[f"var__{column}"] = second - mean * mean if count > 0 else np.nan
        rows.append(out)

    columns = [
        *SUMMARY_KEY_COLUMNS,
        "n",
        *[f"mean__{column}" for column in SUMMARY_MEAN_COLUMNS],
        *[
            name
            for column in SUMMARY_PARAMETER_COLUMNS
            for name in (f"mean__{column}", f"var__{column}")
        ],
    ]
    pd.DataFrame(rows, columns=columns).sort_values(
        ["selected_disease_token_id", "landmark_age", "sex"],
        kind="mergesort",
    ).to_csv(path, index=False)
    return len(rows)


def build_disease_ablated_slice(
    batch: Dict[str, torch.Tensor],
    row_indices: torch.Tensor,
    token_ids: torch.Tensor,
) -> Dict[str, torch.Tensor]:
    """Build an ablated slice for aligned (row, disease_token) pairs."""
    event_seq = batch["event_seq"]
    row_indices = row_indices.to(device=event_seq.device, dtype=torch.long)
    token_ids = token_ids.to(device=event_seq.device, dtype=event_seq.dtype)

    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, :]
    remove = (out["event_seq"] == token_ids[:, None]) & 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)
    empty_rows = ~has_valid
    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)
    missing_readout = ~has_readout
    local_valid = out["padding_mask"]
    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"][missing_readout] = False
    out["readout_mask"][missing_readout, last_pos[missing_readout]] = True
    out["landmark_pos"][missing_readout] = last_pos[missing_readout].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 load_disease_metadata(
    mapping_path: Path,
    *,
    vocab_size: int,
) -> dict[int, dict[str, Any]]:
    if not mapping_path.exists():
        raise FileNotFoundError(f"Disease mapping file not found: {mapping_path}")
    table = pd.read_csv(mapping_path)
    required = {"token_id", "code", "name", "is_death"}
    missing = required - set(table.columns)
    if missing:
        raise ValueError(f"{mapping_path} is missing columns: {sorted(missing)}")

    death_idx = int(vocab_size) - 1
    out: dict[int, dict[str, Any]] = {}
    for row in table.itertuples(index=False):
        token = int(getattr(row, "token_id"))
        if token < 0 or token >= int(vocab_size) or token == death_idx:
            continue
        if int(getattr(row, "is_death")) == 1:
            continue
        meta = {
            "token_id": token,
            "code": str(getattr(row, "code")),
            "name": str(getattr(row, "name")),
        }
        for column in (
            "icd10_chapter",
            "icd10_chapter_title",
            "organ_system",
            "organ_system_label",
        ):
            if hasattr(row, column):
                meta[column] = str(getattr(row, column))
        out[token] = meta
    return out


def resolve_disease_token(
    value: str,
    metadata: dict[int, dict[str, Any]],
) -> tuple[int, dict[str, Any]]:
    text = str(value).strip()
    if text == "":
        raise ValueError("--disease must not be empty")

    if text.isdigit() or (text.startswith("-") and text[1:].isdigit()):
        token = int(text)
        if token not in metadata:
            raise ValueError(f"Disease token_id {token} was not found in the mapping")
        return token, metadata[token]

    lower = text.lower()
    exact = [
        (token, meta)
        for token, meta in metadata.items()
        if str(meta.get("code", "")).lower() == lower
        or str(meta.get("name", "")).lower() == lower
    ]
    if len(exact) == 1:
        return exact[0]
    if len(exact) > 1:
        raise ValueError(f"--disease={value!r} matched multiple diseases exactly")

    contains = [
        (token, meta)
        for token, meta in metadata.items()
        if lower in str(meta.get("code", "")).lower()
        or lower in str(meta.get("name", "")).lower()
    ]
    if len(contains) == 1:
        return contains[0]
    if not contains:
        raise ValueError(f"--disease={value!r} did not match any disease token")
    preview = ", ".join(
        f"{token}:{meta.get('code')} ({meta.get('name')})"
        for token, meta in contains[:10]
    )
    raise ValueError(
        f"--disease={value!r} matched {len(contains)} diseases; use token_id or code. "
        f"First matches: {preview}"
    )


def resolve_disease_tokens(
    value: str | None,
    metadata: dict[int, dict[str, Any]],
) -> list[tuple[int, dict[str, Any]]]:
    if value is None or str(value).strip() == "":
        return [(token, metadata[token]) for token in sorted(metadata)]
    out: list[tuple[int, dict[str, Any]]] = []
    seen: set[int] = set()
    for part in str(value).split(","):
        token, meta = resolve_disease_token(part, metadata)
        if token not in seen:
            out.append((token, meta))
            seen.add(token)
    return out


def death_distribution_parameters(
    model,
    hidden: torch.Tensor,
    *,
    dist_mode: str,
    eps: float = 1e-8,
) -> tuple[str, torch.Tensor]:
    """Return death distribution parameters with columns matching PARAMETER_VALUE_COLUMNS."""
    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 output_name_for_run(run_path: Path, eval_split: str, *, all_diseases: bool) -> Path:
    scope = "all_diseases" if all_diseases else "selected_diseases"
    return run_path / f"single_disease_mortality_parameters_{eval_split}_{scope}"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Compute per-disease model attribution to mortality distribution parameters."
    )
    parser.add_argument("--run_path", type=str, required=True)
    parser.add_argument(
        "--disease",
        type=str,
        default=None,
        help=(
            "Optional disease token_id, ICD-10 code, exact name, unambiguous name "
            "substring, or comma-separated list. If omitted, scan all disease tokens."
        ),
    )
    parser.add_argument(
        "--output_path",
        type=str,
        default=None,
        help="Output directory for compressed .npz shards.",
    )
    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("--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 disease-token ablation queries.",
    )
    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(
        "--shard_rows",
        type=int,
        default=200_000,
        help="Approximate number of detailed rows to buffer before writing one .npz shard.",
    )
    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)

    metadata = load_disease_metadata(
        Path(args.organ_mapping_path),
        vocab_size=int(dataset.vocab_size),
    )
    scanned_disease_items = resolve_disease_tokens(args.disease, metadata)
    if not scanned_disease_items:
        raise ValueError("No diseases selected for attribution")
    scanned_disease_tokens = [token for token, _meta in scanned_disease_items]

    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),
    )
    group_names = sorted(organ_groups)

    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)
    selected_token_mask = np.zeros(int(dataset.vocab_size), dtype=bool)
    selected_token_mask[np.asarray(scanned_disease_tokens, dtype=np.int64)] = True
    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")
    if int(args.shard_rows) <= 0:
        raise ValueError("--shard_rows 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_path = (
        Path(args.output_path)
        if args.output_path
        else output_name_for_run(
            run_path,
            eval_split,
            all_diseases=args.disease is None or str(args.disease).strip() == "",
        )
    )
    output_dir = normalize_output_dir(output_path)
    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}")
    if len(scanned_disease_items) == len(metadata):
        print(f"Diseases: all mapped diseases ({len(scanned_disease_items)})")
    else:
        preview = ", ".join(
            f"{token}:{meta.get('code')}" for token, meta in scanned_disease_items[:10]
        )
        print(f"Diseases: {len(scanned_disease_items)} selected ({preview})")
    print(f"Landmark rows: {len(landmark_dataset)}")
    print(f"Attribution batch size: {attribution_batch_size}")
    print(f"Output directory: {output_dir}")

    written_rows = 0
    shard_index = 0
    shards: list[dict[str, Any]] = []
    row_base_cache: dict[int, dict[str, Any]] = {}
    result_row_idx_chunks: list[np.ndarray] = []
    result_disease_token_chunks: list[np.ndarray] = []
    result_value_chunks: list[np.ndarray] = []

    def get_row_base(row_idx: int) -> dict[str, Any]:
        cached = row_base_cache.get(row_idx)
        if cached is not None:
            return cached

        meta = landmark_dataset.rows[int(row_idx)]
        dataset_index = int(meta["dataset_index"])
        sample = dataset.samples[dataset_index]
        hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
        unique_tokens, token_counts = np.unique(hist_tokens, return_counts=True)
        total_count, group_counts = historical_counts_by_group(
            hist_tokens,
            death_idx=death_idx,
            token_to_group=token_to_group,
            group_names=group_names,
        )
        cached = {
            "patient_id": int(meta["patient_id"]),
            "dataset_index": dataset_index,
            "eid": int(sample.get("eid", -1)),
            "sex": int(meta["sex"]),
            "landmark_age": float(meta["landmark_age"]),
            "followup_end_time": float(meta["followup_end_time"]),
            "history_disease_count": int(total_count),
            "_hist_tokens": hist_tokens,
            "_token_counts": {
                int(token): int(count)
                for token, count in zip(unique_tokens.tolist(), token_counts.tolist())
            },
            "_group_counts": group_counts,
        }
        row_base_cache[row_idx] = cached
        return cached

    for batch in tqdm(loader, desc="Per-disease mortality 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()
        }
        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,
        )
        with torch.no_grad():
            _death_distribution, original_params = death_distribution_parameters(
                model,
                hidden,
                dist_mode=dist_mode,
            )
        event_np = batch["event_seq"].numpy()
        valid_event = (event_np >= 0) & (event_np < int(dataset.vocab_size))
        selected_event = np.zeros_like(valid_event, dtype=bool)
        selected_event[valid_event] = selected_token_mask[event_np[valid_event]]
        pair_row_np, pair_pos_np = np.nonzero(selected_event)
        if pair_row_np.size == 0:
            continue
        pair_disease_np = event_np[pair_row_np, pair_pos_np].astype(np.int64, copy=False)

        pair_offset = 0
        while pair_offset < int(pair_row_np.shape[0]):
            pair_stop = min(int(pair_row_np.shape[0]), pair_offset + int(attribution_batch_size))
            local_rows_np = pair_row_np[pair_offset:pair_stop].astype(np.int64, copy=False)
            disease_tokens_np = pair_disease_np[pair_offset:pair_stop]
            local_rows = torch.as_tensor(local_rows_np, dtype=torch.long, device=device)
            disease_token_ids = torch.as_tensor(disease_tokens_np, dtype=torch.long, device=device)
            ablated_chunk = build_disease_ablated_slice(
                batch=batch_dev,
                row_indices=local_rows,
                token_ids=disease_token_ids,
            )
            with torch.no_grad():
                ablated_hidden = infer_landmark_hidden(
                    model=model,
                    batch=ablated_chunk,
                    device=device,
                    model_target_mode=model_target_mode,
                    readout_name=readout_name,
                    readout_reduce=readout_reduce,
                )
                _ablated_distribution, ablated_params = death_distribution_parameters(
                    model,
                    ablated_hidden,
                    dist_mode=dist_mode,
                )
            value_block = parameter_pair_block(
                original_params[local_rows],
                ablated_params,
            ).detach().cpu().numpy()
            row_ids = batch["row_idx"][local_rows_np].numpy().astype(np.int64, copy=False)
            disease_tokens_list = disease_tokens_np
            result_row_idx_chunks.append(row_ids)
            result_disease_token_chunks.append(disease_tokens_list)
            result_value_chunks.append(value_block)
            pair_offset = pair_stop

    if result_value_chunks:
        all_row_ids = np.concatenate(result_row_idx_chunks).astype(np.int64, copy=False)
        all_disease_tokens = np.concatenate(result_disease_token_chunks).astype(
            np.int64,
            copy=False,
        )
        all_values = np.concatenate(result_value_chunks, axis=0)

        rows: list[dict[str, Any]] = []
        for i, (row_idx, disease_token) in enumerate(
            zip(all_row_ids.tolist(), all_disease_tokens.tolist())
        ):
            disease_token = int(disease_token)
            disease_meta = metadata[disease_token]
            row_base = get_row_base(int(row_idx))
            group_counts = row_base["_group_counts"]
            disease_history_count = int(row_base["_token_counts"].get(disease_token, 0))
            if disease_history_count <= 0:
                raise RuntimeError(
                    "Internal mismatch: occurred mask selected disease "
                    f"{disease_token} for row {row_idx}, but cached history has count 0"
                )

            rows.append(
                {
                    "patient_id": row_base["patient_id"],
                    "dataset_index": row_base["dataset_index"],
                    "eid": row_base["eid"],
                    "sex": row_base["sex"],
                    "landmark_age": row_base["landmark_age"],
                    "followup_end_time": row_base["followup_end_time"],
                    "history_disease_count": row_base["history_disease_count"],
                    "selected_disease_history_count": disease_history_count,
                    "selected_disease_token_id": int(disease_token),
                    "selected_disease_code": str(disease_meta.get("code", "")),
                    "selected_disease_name": str(disease_meta.get("name", "")),
                    "selected_disease_organ_system": str(disease_meta.get("organ_system", "")),
                    "selected_disease_organ_system_label": str(
                        disease_meta.get("organ_system_label", "")
                    ),
                    "history_count__selected_organ_system": int(
                        group_counts.get(str(disease_meta.get("organ_system", "")), 0)
                    ),
                    "death_distribution": death_distribution_name,
                    "original_death_lambda": float(all_values[i, 0]),
                    "ablated_death_lambda": float(all_values[i, 1]),
                    "original_death_scale": float(all_values[i, 2]),
                    "ablated_death_scale": float(all_values[i, 3]),
                    "original_death_shape": float(all_values[i, 4]),
                    "ablated_death_shape": float(all_values[i, 5]),
                }
            )

        result_table = pd.DataFrame(rows).reindex(columns=OUTPUT_COLUMNS)
        written_rows = int(len(result_table))

        summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
        update_summary_accumulator(summary_accumulator, result_table)

        for start in range(0, written_rows, int(args.shard_rows)):
            stop = min(written_rows, start + int(args.shard_rows))
            shard_name = f"part-{shard_index:06d}.npz"
            shard_path = output_dir / shard_name
            shard_rows = write_compressed_npz_table(
                shard_path,
                result_table.iloc[start:stop],
            )
            shards.append({"file": shard_name, "rows": int(shard_rows)})
            shard_index += 1
    else:
        result_table = pd.DataFrame(columns=OUTPUT_COLUMNS)
        summary_accumulator = {}

    if not shards:
        empty_path = output_dir / "part-000000.npz"
        write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))
        shards.append({"file": empty_path.name, "rows": 0})
    summary_path = output_dir / "summary_by_disease_age_sex.csv"
    summary_rows = write_summary_csv(summary_path, summary_accumulator)
    write_manifest(
        output_dir,
        rows=written_rows,
        shards=shards,
        summary_file=summary_path.name,
        scanned_diseases=[
            {"token_id": int(token), **{k: v for k, v in meta.items() if k != "token_id"}}
            for token, meta in scanned_disease_items
        ],
        eval_split=eval_split,
        dist_mode=dist_mode,
        landmark_start=float(args.landmark_start),
        landmark_stop=float(args.landmark_stop),
        landmark_step=float(args.landmark_step),
    )
    print(f"Wrote {written_rows} rows in {len(shards)} shard(s) to {output_dir}")
    print(f"Wrote {summary_rows} summary rows to {summary_path}")


if __name__ == "__main__":
    main()