JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Report disease parameter attribution

Browse Source
This commit is contained in:
Jiarui Li
2026-07-01 10:16:51 +08:00
parent 3cf756ccb0
commit d9dc115e09
2 changed files with 127 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

214
evaluate_extra_info_attribution.py
View File

@@ -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(
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,
tau=float(args.tau),
logits=original_logits,
rho=original_rho,
)
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()
}
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__":

8
run_missing_evaluations.sh
View File

@@ -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