951 lines
34 KiB
Python
951 lines
34 KiB
Python
"""Evaluate extra-info attribution to death and disease distribution parameters.
|
|
|
|
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 concurrent.futures import ProcessPoolExecutor, as_completed
|
|
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 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 all_disease_parameter_pair_block(
|
|
*,
|
|
original_logits: torch.Tensor,
|
|
ablated_logits: torch.Tensor,
|
|
dist_mode: str,
|
|
original_rho: torch.Tensor | None = None,
|
|
ablated_rho: torch.Tensor | None = None,
|
|
eps: float = 1e-8,
|
|
) -> torch.Tensor:
|
|
original_lambda = F.softplus(original_logits) + float(eps)
|
|
ablated_lambda = F.softplus(ablated_logits) + float(eps)
|
|
|
|
if dist_mode in {"exponential", "mixed"}:
|
|
nan = torch.full_like(original_lambda, float("nan"))
|
|
return torch.stack(
|
|
[
|
|
original_lambda,
|
|
ablated_lambda,
|
|
nan,
|
|
nan,
|
|
nan,
|
|
nan,
|
|
],
|
|
dim=2,
|
|
)
|
|
|
|
if dist_mode == "weibull":
|
|
if original_rho is None or ablated_rho is None:
|
|
raise ValueError("rho tensors are required for weibull disease parameters")
|
|
original_shape = original_rho.to(dtype=original_lambda.dtype).clamp_min(float(eps))
|
|
ablated_shape = ablated_rho.to(dtype=ablated_lambda.dtype).clamp_min(float(eps))
|
|
original_scale = torch.pow(original_lambda.clamp_min(float(eps)), -1.0 / original_shape)
|
|
ablated_scale = torch.pow(ablated_lambda.clamp_min(float(eps)), -1.0 / ablated_shape)
|
|
nan = torch.full_like(original_lambda, float("nan"))
|
|
return torch.stack(
|
|
[
|
|
nan,
|
|
nan,
|
|
original_scale,
|
|
ablated_scale,
|
|
original_shape,
|
|
ablated_shape,
|
|
],
|
|
dim=2,
|
|
)
|
|
|
|
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
|
|
|
|
|
|
def grouped_parameter_stats(
|
|
values: torch.Tensor,
|
|
group_token_mask: torch.Tensor,
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
finite = torch.isfinite(values)
|
|
values64 = values.to(dtype=torch.float64)
|
|
safe_values = torch.where(finite, values64, torch.zeros_like(values64))
|
|
mask = group_token_mask.to(device=values.device, dtype=torch.float64)
|
|
sums = torch.einsum("nvc,gv->ngc", safe_values, mask)
|
|
sumsq = torch.einsum("nvc,gv->ngc", safe_values * safe_values, mask)
|
|
counts = torch.einsum("nvc,gv->ngc", finite.to(dtype=torch.float64), mask)
|
|
return (
|
|
sums.detach().cpu().numpy().astype(np.float64, copy=False),
|
|
sumsq.detach().cpu().numpy().astype(np.float64, copy=False),
|
|
counts.detach().cpu().numpy().astype(np.float64, copy=False),
|
|
)
|
|
|
|
|
|
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 finite_float64(values: Any) -> np.ndarray:
|
|
arr = np.asarray(values, dtype=np.float64)
|
|
return arr[np.isfinite(arr)]
|
|
|
|
|
|
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 = finite_float64(pd.to_numeric(group[column], errors="coerce"))
|
|
acc[f"count__{column}"] += float(vals.size)
|
|
acc[f"sum__{column}"] += float(vals.sum())
|
|
acc[f"sumsq__{column}"] += float(np.square(vals).sum())
|
|
|
|
|
|
def update_disease_parameter_summary_from_group_stats(
|
|
summary: dict[tuple[Any, ...], dict[str, float]],
|
|
*,
|
|
key_rows: pd.DataFrame,
|
|
group_names: Sequence[str],
|
|
group_labels: Sequence[str],
|
|
sums: np.ndarray,
|
|
sumsq: np.ndarray,
|
|
counts: np.ndarray,
|
|
) -> None:
|
|
if key_rows.empty or sums.size == 0:
|
|
return
|
|
rows = key_rows.reset_index(drop=True)
|
|
for row_idx, row in rows.iterrows():
|
|
base_key = tuple(row[column] for column in EXTRA_KEY_COLUMNS)
|
|
for group_idx, (group, label) in enumerate(zip(group_names, group_labels)):
|
|
count_row = counts[int(row_idx), int(group_idx)]
|
|
n_add = float(np.nanmax(count_row)) if count_row.size else 0.0
|
|
if n_add <= 0:
|
|
continue
|
|
full_key = (*base_key, str(group), str(label))
|
|
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"] += n_add
|
|
for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS):
|
|
count = float(counts[int(row_idx), int(group_idx), int(col_idx)])
|
|
if count <= 0:
|
|
continue
|
|
acc[f"count__{column}"] += count
|
|
acc[f"sum__{column}"] += float(sums[int(row_idx), int(group_idx), int(col_idx)])
|
|
acc[f"sumsq__{column}"] += float(sumsq[int(row_idx), int(group_idx), int(col_idx)])
|
|
|
|
|
|
def merge_summary_dict(
|
|
dst: dict[tuple[Any, ...], dict[str, float]],
|
|
src: dict[tuple[Any, ...], dict[str, float]],
|
|
) -> None:
|
|
for key, src_acc in src.items():
|
|
dst_acc = dst.setdefault(key, {name: 0.0 for name in src_acc})
|
|
for name, value in src_acc.items():
|
|
dst_acc[name] = dst_acc.get(name, 0.0) + float(value)
|
|
|
|
|
|
def reduce_attribution_chunk_bundle(
|
|
payload: tuple[
|
|
list[tuple[pd.DataFrame, np.ndarray]],
|
|
list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
|
|
list[str],
|
|
list[str],
|
|
],
|
|
) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
|
|
death_items, disease_items, group_names, group_labels = payload
|
|
death_summary: dict[tuple[Any, ...], dict[str, float]] = {}
|
|
disease_summary: dict[tuple[Any, ...], dict[str, float]] = {}
|
|
|
|
for key_rows, values in death_items:
|
|
update_death_summary(
|
|
death_summary,
|
|
key_rows=key_rows,
|
|
values=values,
|
|
)
|
|
|
|
for key_rows, sums, sumsq, counts in disease_items:
|
|
update_disease_parameter_summary_from_group_stats(
|
|
disease_summary,
|
|
key_rows=key_rows,
|
|
group_names=group_names,
|
|
group_labels=group_labels,
|
|
sums=sums,
|
|
sumsq=sumsq,
|
|
counts=counts,
|
|
)
|
|
|
|
return death_summary, disease_summary
|
|
|
|
|
|
def reduce_attribution_chunks(
|
|
*,
|
|
death_key_chunks: list[pd.DataFrame],
|
|
death_value_chunks: list[np.ndarray],
|
|
disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
|
|
group_names: list[str],
|
|
group_labels: list[str],
|
|
cpu_reduce_workers: int,
|
|
) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
|
|
n_chunks = max(len(death_key_chunks), len(disease_stat_chunks))
|
|
if n_chunks == 0:
|
|
return {}, {}
|
|
|
|
worker_count = max(1, min(int(cpu_reduce_workers), n_chunks))
|
|
if worker_count == 1:
|
|
return reduce_attribution_chunk_bundle(
|
|
(
|
|
list(zip(death_key_chunks, death_value_chunks)),
|
|
disease_stat_chunks,
|
|
group_names,
|
|
group_labels,
|
|
)
|
|
)
|
|
|
|
bundles = []
|
|
for worker_idx in range(worker_count):
|
|
start = worker_idx * n_chunks // worker_count
|
|
stop = (worker_idx + 1) * n_chunks // worker_count
|
|
if start >= stop:
|
|
continue
|
|
death_items = [
|
|
(death_key_chunks[i], death_value_chunks[i])
|
|
for i in range(start, min(stop, len(death_key_chunks)))
|
|
]
|
|
disease_items = disease_stat_chunks[start:min(stop, len(disease_stat_chunks))]
|
|
bundles.append((death_items, disease_items, group_names, group_labels))
|
|
|
|
merged_death: dict[tuple[Any, ...], dict[str, float]] = {}
|
|
merged_disease: dict[tuple[Any, ...], dict[str, float]] = {}
|
|
with ProcessPoolExecutor(max_workers=len(bundles)) as executor:
|
|
futures = [executor.submit(reduce_attribution_chunk_bundle, bundle) for bundle in bundles]
|
|
for future in tqdm(as_completed(futures), total=len(futures), desc="CPU summary reduction", dynamic_ncols=True):
|
|
death_part, disease_part = future.result()
|
|
merge_summary_dict(merged_death, death_part)
|
|
merge_summary_dict(merged_disease, disease_part)
|
|
|
|
return merged_death, merged_disease
|
|
|
|
|
|
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 and disease distribution parameters."
|
|
)
|
|
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(
|
|
"--cpu_reduce_workers",
|
|
type=int,
|
|
default=None,
|
|
help="Worker processes for CPU-side summary reduction. Defaults to --num_workers.",
|
|
)
|
|
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,
|
|
}
|
|
group_names = list(risk_groups.keys())
|
|
group_labels = [str(risk_group_labels[group]) for group in group_names]
|
|
|
|
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()
|
|
|
|
group_token_mask = torch.zeros(
|
|
(len(group_names), int(dataset.vocab_size)),
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
for group_idx, group in enumerate(group_names):
|
|
valid_tokens = [
|
|
int(token)
|
|
for token in risk_groups[group]
|
|
if 0 <= int(token) < int(dataset.vocab_size) and int(token) != death_idx
|
|
]
|
|
if valid_tokens:
|
|
group_token_mask[group_idx, torch.as_tensor(valid_tokens, dtype=torch.long, device=device)] = 1.0
|
|
|
|
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))
|
|
cpu_reduce_workers = int(
|
|
args.cpu_reduce_workers
|
|
if args.cpu_reduce_workers is not None
|
|
else max(1, num_workers)
|
|
)
|
|
if cpu_reduce_workers <= 0:
|
|
raise ValueError("--cpu_reduce_workers must be positive")
|
|
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"CPU reduce workers: {cpu_reduce_workers}")
|
|
print(f"Output directory: {output_dir}")
|
|
|
|
death_key_chunks: list[pd.DataFrame] = []
|
|
death_value_chunks: list[np.ndarray] = []
|
|
disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]] = []
|
|
|
|
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
|
|
|
|
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()
|
|
death_key_chunks.append(key_table)
|
|
death_value_chunks.append(value_block)
|
|
|
|
disease_values = all_disease_parameter_pair_block(
|
|
original_logits=original_logits[row_tensor],
|
|
ablated_logits=ablated_logits,
|
|
dist_mode=dist_mode,
|
|
original_rho=None if original_rho is None else original_rho[row_tensor],
|
|
ablated_rho=ablated_rho,
|
|
)
|
|
sums, sumsq, counts = grouped_parameter_stats(
|
|
disease_values,
|
|
group_token_mask,
|
|
)
|
|
disease_stat_chunks.append((key_table, sums, sumsq, counts))
|
|
|
|
death_summary, disease_parameter_summary = reduce_attribution_chunks(
|
|
death_key_chunks=death_key_chunks,
|
|
death_value_chunks=death_value_chunks,
|
|
disease_stat_chunks=disease_stat_chunks,
|
|
group_names=group_names,
|
|
group_labels=group_labels,
|
|
cpu_reduce_workers=cpu_reduce_workers,
|
|
)
|
|
|
|
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()
|