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