diff --git a/evaluate_extra_info_attribution.py b/evaluate_extra_info_attribution.py index dbccc22..cb8c425 100644 --- a/evaluate_extra_info_attribution.py +++ b/evaluate_extra_info_attribution.py @@ -5,7 +5,7 @@ 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; -* tau-year future incident disease risk before and after ablation, by ICD-10 +* disease distribution parameters before and after ablation, by ICD-10 chapter-derived organ/system groups. Death is always token vocab_size - 1. @@ -44,10 +44,7 @@ from landmark_eval_utils import ( load_eval_sequence_dataset, load_organ_groups, make_landmark_ages, - make_occurred_mask, ) -from future_risk import new_disease_risk_from_probabilities, probabilities_from_logits - EXTRA_KEY_COLUMNS = [ "selected_extra_info_type_id", @@ -66,15 +63,19 @@ DEATH_PARAMETER_COLUMNS = [ "ablated_death_shape", ] -DISEASE_RISK_KEY_COLUMNS = [ +DISEASE_PARAMETER_KEY_COLUMNS = [ *EXTRA_KEY_COLUMNS, "target_group", "target_group_label", ] -DISEASE_RISK_COLUMNS = [ - "original_future_disease_risk", - "ablated_future_disease_risk", +DISEASE_PARAMETER_COLUMNS = [ + "original_disease_lambda", + "ablated_disease_lambda", + "original_disease_scale", + "ablated_disease_scale", + "original_disease_shape", + "ablated_disease_shape", ] @@ -211,6 +212,38 @@ def death_distribution_parameters( 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( [ @@ -225,22 +258,17 @@ def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch ) -def disease_probabilities( - model, - hidden: torch.Tensor, - *, - dist_mode: str, - tau: float, -) -> torch.Tensor: - logits = model.calc_risk(hidden) - rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None - death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None - return probabilities_from_logits( - logits, - tau, - dist_mode=dist_mode, - rho=rho, - death_rho=death_rho, +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, ) @@ -368,37 +396,38 @@ def update_death_summary( acc[f"sumsq__{column}"] += float((vals * vals).sum()) -def update_disease_risk_summary( +def update_disease_parameter_summary( summary: dict[tuple[Any, ...], dict[str, float]], *, key_rows: pd.DataFrame, target_group: str, target_group_label: str, - original_risk: np.ndarray, - ablated_risk: np.ndarray, + values: np.ndarray, ) -> None: - if key_rows.empty: + if key_rows.empty or values.size == 0: return - table = key_rows.copy() - table["target_group"] = str(target_group) - table["target_group_label"] = str(target_group_label) - table["original_future_disease_risk"] = original_risk - table["ablated_future_disease_risk"] = ablated_risk - - for key, group in table.groupby(DISEASE_RISK_KEY_COLUMNS, dropna=False, sort=False): + 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( - key, + full_key, { "n": 0.0, - **{f"sum__{col}": 0.0 for col in DISEASE_RISK_COLUMNS}, - **{f"sumsq__{col}": 0.0 for col in DISEASE_RISK_COLUMNS}, + **{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(len(group)) - for column in DISEASE_RISK_COLUMNS: - vals = pd.to_numeric(group[column], errors="coerce") + 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()) @@ -439,31 +468,28 @@ def write_death_summary_csv( return len(rows) -def write_disease_risk_summary_csv( +def write_disease_parameter_summary_csv( path: Path, summary: dict[tuple[Any, ...], dict[str, float]], - *, - tau: 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_RISK_KEY_COLUMNS, key)} + row = {column: value for column, value in zip(DISEASE_PARAMETER_KEY_COLUMNS, key)} row["n"] = n - row["tau_years"] = float(tau) - for column in DISEASE_RISK_COLUMNS: - mean = acc[f"sum__{column}"] / n if n > 0 else np.nan - second = acc[f"sumsq__{column}"] / n if n > 0 else np.nan + 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 n > 0 else np.nan + row[f"var__{column}"] = second - mean * mean if count > 0 else np.nan rows.append(row) columns = [ - *DISEASE_RISK_KEY_COLUMNS, + *DISEASE_PARAMETER_KEY_COLUMNS, "n", - "tau_years", *[ name - for column in DISEASE_RISK_COLUMNS + for column in DISEASE_PARAMETER_COLUMNS for name in (f"mean__{column}", f"var__{column}") ], ] @@ -498,7 +524,6 @@ def parse_args() -> argparse.Namespace: 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("--tau", 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( @@ -628,7 +653,7 @@ def main() -> None: output_dir = ( Path(args.output_dir) if args.output_dir - else run_path / f"extra_info_attribution_{eval_split}_tau{float(args.tau):g}y" + else run_path / f"extra_info_attribution_{eval_split}" ) output_dir.mkdir(parents=True, exist_ok=True) @@ -645,7 +670,7 @@ def main() -> None: print(f"Output directory: {output_dir}") death_summary: dict[tuple[Any, ...], dict[str, float]] = {} - disease_risk_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 = { @@ -666,25 +691,21 @@ def main() -> None: hidden, dist_mode=dist_mode, ) - original_probabilities = disease_probabilities( - model, - hidden, - dist_mode=dist_mode, - tau=float(args.tau), - ) - occurred = make_occurred_mask( - batch_dev["event_seq"].long(), - vocab_size=int(dataset.vocab_size), - device=device, - ) - original_risk_by_group = { - group: new_disease_risk_from_probabilities( - original_probabilities, - occurred, - tokens, - ) - for group, tokens in risk_groups.items() - } + 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, @@ -706,12 +727,8 @@ def main() -> None: ablated_hidden, dist_mode=dist_mode, ) - ablated_probabilities = disease_probabilities( - model, - ablated_hidden, - dist_mode=dist_mode, - tau=float(args.tau), - ) + 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): @@ -736,43 +753,52 @@ def main() -> None: values=value_block, ) - ablated_occurred = occurred[row_tensor] for group, tokens in risk_groups.items(): - ablated_risk = new_disease_risk_from_probabilities( - ablated_probabilities, - ablated_occurred, - tokens, + disease_distribution, ablated_disease_params = disease_distribution_parameters( + model, + ablated_hidden, + token_ids=tokens, + dist_mode=dist_mode, + logits=ablated_logits, + rho=ablated_rho, ) - update_disease_risk_summary( - disease_risk_summary, + 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], - original_risk=original_risk_by_group[group][row_tensor].detach().cpu().numpy(), - ablated_risk=ablated_risk.detach().cpu().numpy(), + values=value_block, ) death_summary_path = output_dir / "summary_extra_info_death_parameters.csv" - disease_summary_path = output_dir / "summary_extra_info_future_disease_risk.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_risk_summary_csv( + disease_rows = write_disease_parameter_summary_csv( disease_summary_path, - disease_risk_summary, - tau=float(args.tau), + disease_parameter_summary, ) manifest = { "death_summary_file": death_summary_path.name, - "disease_risk_summary_file": disease_summary_path.name, + "disease_parameter_summary_file": disease_summary_path.name, "death_summary_rows": int(death_rows), - "disease_risk_summary_rows": int(disease_rows), + "disease_parameter_summary_rows": int(disease_rows), "eval_split": eval_split, "split_source": split_source, "dist_mode": dist_mode, - "tau_years": float(args.tau), "landmark_start": float(args.landmark_start), "landmark_stop": float(args.landmark_stop), "landmark_step": float(args.landmark_step), @@ -784,7 +810,7 @@ def main() -> None: 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-risk summary rows to {disease_summary_path}") + print(f"Wrote {disease_rows} disease-parameter summary rows to {disease_summary_path}") if __name__ == "__main__": diff --git a/run_missing_evaluations.sh b/run_missing_evaluations.sh index 069e4b3..8e03a43 100644 --- a/run_missing_evaluations.sh +++ b/run_missing_evaluations.sh @@ -9,7 +9,6 @@ cd "$(dirname "${BASH_SOURCE[0]}")" PYTHON_BIN="${PYTHON_BIN:-python}" DEVICE="${DEVICE:-cuda}" EVAL_SPLIT="${EVAL_SPLIT:-test}" -TAU="${TAU:-5}" NUM_WORKERS="${NUM_WORKERS:-4}" NUM_WORKERS_AUC="${NUM_WORKERS_AUC:-}" BATCH_SIZE="${BATCH_SIZE:-}" @@ -20,7 +19,6 @@ DRY_RUN="${DRY_RUN:-0}" # surface in this repository. Set either variable to 0 to leave that family out. RUN_EXTRA_INFO_ATTRIBUTION="${RUN_EXTRA_INFO_ATTRIBUTION:-1}" RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION="${RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION:-1}" -TAU_LABEL="$("${PYTHON_BIN}" -c 'import sys; print(f"{float(sys.argv[1]):g}")' "${TAU}")" common_args_base() { printf '%s\n' --run_path "$1" --eval_split "${EVAL_SPLIT}" --num_workers "${NUM_WORKERS}" @@ -130,10 +128,10 @@ for run_path in runs/*; do if run_has_extra_info "${run_path}"; then run_dir_result_if_missing \ "evaluate_extra_info_attribution.py" \ - "${run_path}/extra_info_attribution_${EVAL_SPLIT}_tau${TAU_LABEL}y" \ + "${run_path}/extra_info_attribution_${EVAL_SPLIT}" \ "manifest.json" \ - "summary_extra_info_future_disease_risk.csv" \ - "${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}" --tau "${TAU}" + "summary_extra_info_disease_parameters.csv" \ + "${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}" else echo " skip evaluate_extra_info_attribution.py: run has no extra-info types" fi