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

Defer attribution output until after GPU pass

Browse Source
This commit is contained in:
Jiarui Li
2026-06-27 16:11:49 +08:00
parent 980dd7c282
commit fe826c4052
1 changed files with 118 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files

152
evaluate_single_disease_mortality_attribution.py
View File

@@ -42,7 +42,6 @@ from evaluate_event_free_survival import (
load_eval_sequence_dataset,
load_organ_groups,
make_landmark_ages,
make_occurred_mask,
mortality_hazard_from_risk,
)
from future_risk import death_risk_from_probabilities, probabilities_from_logits
@@ -225,8 +224,7 @@ def build_disease_ablated_slice(
out["readout_mask"] &= ~remove
has_valid = out["padding_mask"].any(dim=1)
if not bool(has_valid.all().item()):
empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten()
empty_rows = ~has_valid
out["event_seq"][empty_rows, 0] = CHECKUP_IDX
out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
dtype=out["time_seq"].dtype
@@ -236,17 +234,18 @@ def build_disease_ablated_slice(
out["landmark_pos"][empty_rows] = 0
has_readout = out["readout_mask"].any(dim=1)
if not bool(has_readout.all().item()):
rows = torch.nonzero(~has_readout, as_tuple=False).flatten()
local_valid = out["padding_mask"][rows]
missing_readout = ~has_readout
local_valid = out["padding_mask"]
last_pos = torch.where(
local_valid,
positions.expand(local_valid.shape[0], -1),
torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
).amax(dim=1)
out["readout_mask"][rows] = False
out["readout_mask"][rows, last_pos] = True
out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype)
out["readout_mask"][missing_readout] = False
out["readout_mask"][missing_readout, last_pos[missing_readout]] = True
out["landmark_pos"][missing_readout] = last_pos[missing_readout].to(
dtype=out["landmark_pos"].dtype
)
repeated_keys = (
"sex",
@@ -508,11 +507,8 @@ def main() -> None:
cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode
device = resolve_eval_device(args.device)
scanned_disease_tensor = torch.as_tensor(
scanned_disease_tokens,
dtype=torch.long,
device=device,
)
selected_token_mask = np.zeros(int(dataset.vocab_size), dtype=bool)
selected_token_mask[np.asarray(scanned_disease_tokens, dtype=np.int64)] = True
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
load_model_state(model, state_dict)
model.eval()
@@ -575,46 +571,10 @@ def main() -> None:
written_rows = 0
shard_index = 0
shards: list[dict[str, Any]] = []
summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
row_base_cache: dict[int, dict[str, Any]] = {}
shard_frames: list[pd.DataFrame] = []
shard_buffer_rows = 0
def flush_shard_buffer() -> None:
nonlocal written_rows, shard_index, shard_frames, shard_buffer_rows
if shard_buffer_rows == 0:
return
table = pd.concat(shard_frames, ignore_index=True).reindex(columns=OUTPUT_COLUMNS)
shard_name = f"part-{shard_index:06d}.npz"
shard_path = output_dir / shard_name
shard_rows = write_compressed_npz_table(shard_path, table)
shards.append({"file": shard_name, "rows": int(shard_rows)})
shard_index += 1
written_rows += int(shard_rows)
shard_frames = []
shard_buffer_rows = 0
def write_result_rows(meta_rows: list[dict[str, Any]], value_block: np.ndarray) -> None:
nonlocal shard_buffer_rows
if not meta_rows:
return
for i, row in enumerate(meta_rows):
row["death_risk"] = float(value_block[i, 0])
row["death_hazard"] = float(value_block[i, 1])
row["ablated_death_risk"] = float(value_block[i, 2])
row["ablated_death_hazard"] = float(value_block[i, 3])
row["mortality_attribution_probability"] = float(value_block[i, 4])
row["mortality_attribution_hazard"] = float(value_block[i, 5])
row["mortality_attribution_probability_ratio"] = float(value_block[i, 6])
row["mortality_attribution_hazard_ratio"] = float(value_block[i, 7])
table = pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)
update_summary_accumulator(summary_accumulator, table)
shard_frames.append(table)
shard_buffer_rows += len(table)
if shard_buffer_rows >= int(args.shard_rows):
flush_shard_buffer()
result_row_idx_chunks: list[np.ndarray] = []
result_disease_token_chunks: list[np.ndarray] = []
result_value_chunks: list[np.ndarray] = []
def get_row_base(row_idx: int) -> dict[str, Any]:
cached = row_base_cache.get(row_idx)
@@ -677,24 +637,22 @@ def main() -> None:
)
death_risk_tensor = death_risk_from_probabilities(probabilities)
death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
occurred = make_occurred_mask(
batch_dev["event_seq"],
vocab_size=int(dataset.vocab_size),
device=device,
)
pair_indices = torch.nonzero(
occurred[:, scanned_disease_tensor],
as_tuple=False,
)
if pair_indices.numel() == 0:
event_np = batch["event_seq"].numpy()
valid_event = (event_np >= 0) & (event_np < int(dataset.vocab_size))
selected_event = np.zeros_like(valid_event, dtype=bool)
selected_event[valid_event] = selected_token_mask[event_np[valid_event]]
pair_row_np, pair_pos_np = np.nonzero(selected_event)
if pair_row_np.size == 0:
continue
pair_disease_np = event_np[pair_row_np, pair_pos_np].astype(np.int64, copy=False)
pair_offset = 0
while pair_offset < int(pair_indices.shape[0]):
pair_stop = min(int(pair_indices.shape[0]), pair_offset + int(attribution_batch_size))
pair_chunk = pair_indices[pair_offset:pair_stop]
local_rows = pair_chunk[:, 0].long()
disease_token_ids = scanned_disease_tensor[pair_chunk[:, 1]].long()
while pair_offset < int(pair_row_np.shape[0]):
pair_stop = min(int(pair_row_np.shape[0]), pair_offset + int(attribution_batch_size))
local_rows_np = pair_row_np[pair_offset:pair_stop].astype(np.int64, copy=False)
disease_tokens_np = pair_disease_np[pair_offset:pair_stop]
local_rows = torch.as_tensor(local_rows_np, dtype=torch.long, device=device)
disease_token_ids = torch.as_tensor(disease_tokens_np, dtype=torch.long, device=device)
ablated_chunk = build_disease_ablated_slice(
batch=batch_dev,
row_indices=local_rows,
@@ -731,12 +689,24 @@ def main() -> None:
],
dim=1,
).detach().cpu().numpy()
row_ids = batch["row_idx"][local_rows_np].numpy().astype(np.int64, copy=False)
disease_tokens_list = disease_tokens_np
result_row_idx_chunks.append(row_ids)
result_disease_token_chunks.append(disease_tokens_list)
result_value_chunks.append(value_block)
pair_offset = pair_stop
meta_chunk: list[dict[str, Any]] = []
row_ids = batch_dev["row_idx"][local_rows].detach().cpu().numpy().astype(np.int64)
disease_tokens_list = disease_token_ids.detach().cpu().numpy().astype(np.int64)
for local_pos, (row_idx, disease_token) in enumerate(
zip(row_ids.tolist(), disease_tokens_list.tolist())
if result_value_chunks:
all_row_ids = np.concatenate(result_row_idx_chunks).astype(np.int64, copy=False)
all_disease_tokens = np.concatenate(result_disease_token_chunks).astype(
np.int64,
copy=False,
)
all_values = np.concatenate(result_value_chunks, axis=0)
rows: list[dict[str, Any]] = []
for i, (row_idx, disease_token) in enumerate(
zip(all_row_ids.tolist(), all_disease_tokens.tolist())
):
disease_token = int(disease_token)
disease_meta = metadata[disease_token]
@@ -749,7 +719,7 @@ def main() -> None:
f"{disease_token} for row {row_idx}, but cached history has count 0"
)
meta_chunk.append(
rows.append(
{
"patient_id": row_base["patient_id"],
"dataset_index": row_base["dataset_index"],
@@ -763,22 +733,44 @@ def main() -> None:
"selected_disease_token_id": int(disease_token),
"selected_disease_code": str(disease_meta.get("code", "")),
"selected_disease_name": str(disease_meta.get("name", "")),
"selected_disease_organ_system": str(
disease_meta.get("organ_system", "")
),
"selected_disease_organ_system": str(disease_meta.get("organ_system", "")),
"selected_disease_organ_system_label": str(
disease_meta.get("organ_system_label", "")
),
"history_count__selected_organ_system": int(
group_counts.get(str(disease_meta.get("organ_system", "")), 0)
),
"death_risk": float(all_values[i, 0]),
"death_hazard": float(all_values[i, 1]),
"ablated_death_risk": float(all_values[i, 2]),
"ablated_death_hazard": float(all_values[i, 3]),
"mortality_attribution_probability": float(all_values[i, 4]),
"mortality_attribution_hazard": float(all_values[i, 5]),
"mortality_attribution_probability_ratio": float(all_values[i, 6]),
"mortality_attribution_hazard_ratio": float(all_values[i, 7]),
}
)
write_result_rows(meta_chunk, value_block)
pair_offset = pair_stop
result_table = pd.DataFrame(rows).reindex(columns=OUTPUT_COLUMNS)
written_rows = int(len(result_table))
summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
update_summary_accumulator(summary_accumulator, result_table)
for start in range(0, written_rows, int(args.shard_rows)):
stop = min(written_rows, start + int(args.shard_rows))
shard_name = f"part-{shard_index:06d}.npz"
shard_path = output_dir / shard_name
shard_rows = write_compressed_npz_table(
shard_path,
result_table.iloc[start:stop],
)
shards.append({"file": shard_name, "rows": int(shard_rows)})
shard_index += 1
else:
result_table = pd.DataFrame(columns=OUTPUT_COLUMNS)
summary_accumulator = {}
flush_shard_buffer()
if not shards:
empty_path = output_dir / "part-000000.npz"
write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))