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

244
evaluate_single_disease_mortality_attribution.py
View File

@@ -42,7 +42,6 @@ from evaluate_event_free_survival import (
load_eval_sequence_dataset, load_eval_sequence_dataset,
load_organ_groups, load_organ_groups,
make_landmark_ages, make_landmark_ages,
make_occurred_mask,
mortality_hazard_from_risk, mortality_hazard_from_risk,
) )
from future_risk import death_risk_from_probabilities, probabilities_from_logits from future_risk import death_risk_from_probabilities, probabilities_from_logits
@@ -225,28 +224,28 @@ def build_disease_ablated_slice(
out["readout_mask"] &= ~remove out["readout_mask"] &= ~remove
has_valid = out["padding_mask"].any(dim=1) has_valid = out["padding_mask"].any(dim=1)
if not bool(has_valid.all().item()): empty_rows = ~has_valid
empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten() out["event_seq"][empty_rows, 0] = CHECKUP_IDX
out["event_seq"][empty_rows, 0] = CHECKUP_IDX out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to( dtype=out["time_seq"].dtype
dtype=out["time_seq"].dtype )
) out["padding_mask"][empty_rows, 0] = True
out["padding_mask"][empty_rows, 0] = True out["readout_mask"][empty_rows, 0] = True
out["readout_mask"][empty_rows, 0] = True out["landmark_pos"][empty_rows] = 0
out["landmark_pos"][empty_rows] = 0
has_readout = out["readout_mask"].any(dim=1) has_readout = out["readout_mask"].any(dim=1)
if not bool(has_readout.all().item()): missing_readout = ~has_readout
rows = torch.nonzero(~has_readout, as_tuple=False).flatten() local_valid = out["padding_mask"]
local_valid = out["padding_mask"][rows] last_pos = torch.where(
last_pos = torch.where( local_valid,
local_valid, positions.expand(local_valid.shape[0], -1),
positions.expand(local_valid.shape[0], -1), torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
torch.zeros_like(positions.expand(local_valid.shape[0], -1)), ).amax(dim=1)
).amax(dim=1) out["readout_mask"][missing_readout] = False
out["readout_mask"][rows] = False out["readout_mask"][missing_readout, last_pos[missing_readout]] = True
out["readout_mask"][rows, last_pos] = True out["landmark_pos"][missing_readout] = last_pos[missing_readout].to(
out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype) dtype=out["landmark_pos"].dtype
)
repeated_keys = ( repeated_keys = (
"sex", "sex",
@@ -508,11 +507,8 @@ def main() -> None:
cfg_model = dict(cfg) cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode cfg_model["dist_mode"] = dist_mode
device = resolve_eval_device(args.device) device = resolve_eval_device(args.device)
scanned_disease_tensor = torch.as_tensor( selected_token_mask = np.zeros(int(dataset.vocab_size), dtype=bool)
scanned_disease_tokens, selected_token_mask[np.asarray(scanned_disease_tokens, dtype=np.int64)] = True
dtype=torch.long,
device=device,
)
model = build_model_from_dataset(args, cfg_model, dataset).to(device) model = build_model_from_dataset(args, cfg_model, dataset).to(device)
load_model_state(model, state_dict) load_model_state(model, state_dict)
model.eval() model.eval()
@@ -575,46 +571,10 @@ def main() -> None:
written_rows = 0 written_rows = 0
shard_index = 0 shard_index = 0
shards: list[dict[str, Any]] = [] shards: list[dict[str, Any]] = []
summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
row_base_cache: dict[int, dict[str, Any]] = {} row_base_cache: dict[int, dict[str, Any]] = {}
shard_frames: list[pd.DataFrame] = [] result_row_idx_chunks: list[np.ndarray] = []
shard_buffer_rows = 0 result_disease_token_chunks: list[np.ndarray] = []
result_value_chunks: list[np.ndarray] = []
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()
def get_row_base(row_idx: int) -> dict[str, Any]: def get_row_base(row_idx: int) -> dict[str, Any]:
cached = row_base_cache.get(row_idx) cached = row_base_cache.get(row_idx)
@@ -677,24 +637,22 @@ def main() -> None:
) )
death_risk_tensor = death_risk_from_probabilities(probabilities) death_risk_tensor = death_risk_from_probabilities(probabilities)
death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor) death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
occurred = make_occurred_mask( event_np = batch["event_seq"].numpy()
batch_dev["event_seq"], valid_event = (event_np >= 0) & (event_np < int(dataset.vocab_size))
vocab_size=int(dataset.vocab_size), selected_event = np.zeros_like(valid_event, dtype=bool)
device=device, selected_event[valid_event] = selected_token_mask[event_np[valid_event]]
) pair_row_np, pair_pos_np = np.nonzero(selected_event)
pair_indices = torch.nonzero( if pair_row_np.size == 0:
occurred[:, scanned_disease_tensor],
as_tuple=False,
)
if pair_indices.numel() == 0:
continue continue
pair_disease_np = event_np[pair_row_np, pair_pos_np].astype(np.int64, copy=False)
pair_offset = 0 pair_offset = 0
while pair_offset < int(pair_indices.shape[0]): while pair_offset < int(pair_row_np.shape[0]):
pair_stop = min(int(pair_indices.shape[0]), pair_offset + int(attribution_batch_size)) pair_stop = min(int(pair_row_np.shape[0]), pair_offset + int(attribution_batch_size))
pair_chunk = pair_indices[pair_offset:pair_stop] local_rows_np = pair_row_np[pair_offset:pair_stop].astype(np.int64, copy=False)
local_rows = pair_chunk[:, 0].long() disease_tokens_np = pair_disease_np[pair_offset:pair_stop]
disease_token_ids = scanned_disease_tensor[pair_chunk[:, 1]].long() 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( ablated_chunk = build_disease_ablated_slice(
batch=batch_dev, batch=batch_dev,
row_indices=local_rows, row_indices=local_rows,
@@ -731,54 +689,88 @@ def main() -> None:
], ],
dim=1, dim=1,
).detach().cpu().numpy() ).detach().cpu().numpy()
row_ids = batch["row_idx"][local_rows_np].numpy().astype(np.int64, copy=False)
meta_chunk: list[dict[str, Any]] = [] disease_tokens_list = disease_tokens_np
row_ids = batch_dev["row_idx"][local_rows].detach().cpu().numpy().astype(np.int64) result_row_idx_chunks.append(row_ids)
disease_tokens_list = disease_token_ids.detach().cpu().numpy().astype(np.int64) result_disease_token_chunks.append(disease_tokens_list)
for local_pos, (row_idx, disease_token) in enumerate( result_value_chunks.append(value_block)
zip(row_ids.tolist(), disease_tokens_list.tolist())
):
disease_token = int(disease_token)
disease_meta = metadata[disease_token]
row_base = get_row_base(int(row_idx))
group_counts = row_base["_group_counts"]
disease_history_count = int(row_base["_token_counts"].get(disease_token, 0))
if disease_history_count <= 0:
raise RuntimeError(
"Internal mismatch: occurred mask selected disease "
f"{disease_token} for row {row_idx}, but cached history has count 0"
)
meta_chunk.append(
{
"patient_id": row_base["patient_id"],
"dataset_index": row_base["dataset_index"],
"eid": row_base["eid"],
"sex": row_base["sex"],
"landmark_age": row_base["landmark_age"],
"tau": row_base["tau"],
"followup_end_time": row_base["followup_end_time"],
"history_disease_count": row_base["history_disease_count"],
"selected_disease_history_count": disease_history_count,
"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_label": str(
disease_meta.get("organ_system_label", "")
),
"history_count__selected_organ_system": int(
group_counts.get(str(disease_meta.get("organ_system", "")), 0)
),
}
)
write_result_rows(meta_chunk, value_block)
pair_offset = pair_stop pair_offset = pair_stop
flush_shard_buffer() 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]
row_base = get_row_base(int(row_idx))
group_counts = row_base["_group_counts"]
disease_history_count = int(row_base["_token_counts"].get(disease_token, 0))
if disease_history_count <= 0:
raise RuntimeError(
"Internal mismatch: occurred mask selected disease "
f"{disease_token} for row {row_idx}, but cached history has count 0"
)
rows.append(
{
"patient_id": row_base["patient_id"],
"dataset_index": row_base["dataset_index"],
"eid": row_base["eid"],
"sex": row_base["sex"],
"landmark_age": row_base["landmark_age"],
"tau": row_base["tau"],
"followup_end_time": row_base["followup_end_time"],
"history_disease_count": row_base["history_disease_count"],
"selected_disease_history_count": disease_history_count,
"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_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]),
}
)
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 = {}
if not shards: if not shards:
empty_path = output_dir / "part-000000.npz" empty_path = output_dir / "part-000000.npz"
write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS)) write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))