diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py index 29c1672..e9a7527 100644 --- a/evaluate_single_disease_mortality_attribution.py +++ b/evaluate_single_disease_mortality_attribution.py @@ -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,28 +224,28 @@ 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() - 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 - ) - out["padding_mask"][empty_rows, 0] = True - out["readout_mask"][empty_rows, 0] = True - out["landmark_pos"][empty_rows] = 0 + 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 + ) + out["padding_mask"][empty_rows, 0] = True + out["readout_mask"][empty_rows, 0] = True + 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] - 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) + 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"][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,54 +689,88 @@ def main() -> None: ], dim=1, ).detach().cpu().numpy() - - 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()) - ): - 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) + 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 - 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: empty_path = output_dir / "part-000000.npz" write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))