DeepHealth/evaluate_extra_info_attribution.py

"""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;
* disease distribution parameters 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,
)

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_PARAMETER_KEY_COLUMNS = [
    *EXTRA_KEY_COLUMNS,
    "target_group",
    "target_group_label",
]

DISEASE_PARAMETER_COLUMNS = [
    "original_disease_lambda",
    "ablated_disease_lambda",
    "original_disease_scale",
    "ablated_disease_scale",
    "original_disease_shape",
    "ablated_disease_shape",
]


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 disease_distribution_parameters(
    model,
    hidden: torch.Tensor,
    *,
    token_ids: Sequence[int],
    dist_mode: str,
    logits: torch.Tensor | None = None,
    rho: torch.Tensor | None = None,
    eps: float = 1e-8,
) -> tuple[str, torch.Tensor]:
    ids = [int(x) for x in token_ids]
    if not ids:
        empty = hidden.new_empty((hidden.shape[0], 0, 3))
        return "none", empty

    all_logits = model.calc_risk(hidden) if logits is None else logits
    disease_lambda = F.softplus(all_logits[:, ids]) + float(eps)

    if dist_mode in {"exponential", "mixed"}:
        nan = torch.full_like(disease_lambda, float("nan"))
        return "exponential", torch.stack([disease_lambda, nan, nan], dim=2)

    if dist_mode == "weibull":
        all_rho = model.calc_weibull_rho(hidden) if rho is None else rho
        shape = all_rho[:, ids].to(dtype=disease_lambda.dtype).clamp_min(float(eps))
        scale = torch.pow(disease_lambda.clamp_min(float(eps)), -1.0 / shape)
        nan = torch.full_like(disease_lambda, float("nan"))
        return "weibull", torch.stack([nan, scale, shape], dim=2)

    raise ValueError(f"Unsupported dist_mode={dist_mode!r}")


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_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=2,
    )


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_parameter_summary(
    summary: dict[tuple[Any, ...], dict[str, float]],
    *,
    key_rows: pd.DataFrame,
    target_group: str,
    target_group_label: str,
    values: np.ndarray,
) -> None:
    if key_rows.empty or values.size == 0:
        return
    table = key_rows.reset_index(drop=True).copy()
    grouped = table.groupby(EXTRA_KEY_COLUMNS, dropna=False, sort=False)
    for key, group in grouped:
        if not isinstance(key, tuple):
            key = (key,)
        full_key = (*key, str(target_group), str(target_group_label))
        idx = group.index.to_numpy(dtype=np.int64)
        vals_2d = values[idx].reshape(-1, values.shape[-1])
        acc = summary.setdefault(
            full_key,
            {
                "n": 0.0,
                **{f"count__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
                **{f"sum__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
                **{f"sumsq__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
            },
        )
        acc["n"] += float(vals_2d.shape[0])
        for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS):
            vals = vals_2d[:, col_idx]
            vals = vals[np.isfinite(vals)]
            acc[f"count__{column}"] += float(vals.size)
            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_parameter_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"])
        row = {column: value for column, value in zip(DISEASE_PARAMETER_KEY_COLUMNS, key)}
        row["n"] = n
        for column in DISEASE_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 = [
        *DISEASE_PARAMETER_KEY_COLUMNS,
        "n",
        *[
            name
            for column in DISEASE_PARAMETER_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("--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}"
    )
    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_parameter_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_logits = model.calc_risk(hidden)
            original_rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
            original_disease_params_by_group = {}
            disease_distribution_name_by_group = {}
            for group, tokens in risk_groups.items():
                disease_distribution, original_disease_params = disease_distribution_parameters(
                    model,
                    hidden,
                    token_ids=tokens,
                    dist_mode=dist_mode,
                    logits=original_logits,
                    rho=original_rho,
                )
                disease_distribution_name_by_group[group] = disease_distribution
                original_disease_params_by_group[group] = original_disease_params

        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_logits = model.calc_risk(ablated_hidden)
                ablated_rho = model.calc_weibull_rho(ablated_hidden) if dist_mode == "weibull" else None

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

            for group, tokens in risk_groups.items():
                disease_distribution, ablated_disease_params = disease_distribution_parameters(
                    model,
                    ablated_hidden,
                    token_ids=tokens,
                    dist_mode=dist_mode,
                    logits=ablated_logits,
                    rho=ablated_rho,
                )
                if disease_distribution_name_by_group[group] != disease_distribution:
                    raise RuntimeError(
                        "Disease distribution changed between original and ablated passes "
                        f"for group {group!r}: {disease_distribution_name_by_group[group]!r} "
                        f"vs {disease_distribution!r}"
                    )
                value_block = disease_parameter_pair_block(
                    original_disease_params_by_group[group][row_tensor],
                    ablated_disease_params,
                ).detach().cpu().numpy()
                update_disease_parameter_summary(
                    disease_parameter_summary,
                    key_rows=key_table,
                    target_group=group,
                    target_group_label=risk_group_labels[group],
                    values=value_block,
                )

    death_summary_path = output_dir / "summary_extra_info_death_parameters.csv"
    disease_summary_path = output_dir / "summary_extra_info_disease_parameters.csv"
    death_rows = write_death_summary_csv(
        death_summary_path,
        death_summary,
        death_distribution=death_distribution_name,
    )
    disease_rows = write_disease_parameter_summary_csv(
        disease_summary_path,
        disease_parameter_summary,
    )
    manifest = {
        "death_summary_file": death_summary_path.name,
        "disease_parameter_summary_file": disease_summary_path.name,
        "death_summary_rows": int(death_rows),
        "disease_parameter_summary_rows": int(disease_rows),
        "eval_split": eval_split,
        "split_source": split_source,
        "dist_mode": dist_mode,
        "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-parameter summary rows to {disease_summary_path}")


if __name__ == "__main__":
    main()