diff --git a/evaluate_event_free_survival.py b/evaluate_event_free_survival.py index 5e9af6e..64b2043 100644 --- a/evaluate_event_free_survival.py +++ b/evaluate_event_free_survival.py @@ -288,14 +288,21 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st } -def ablate_event_history_for_tokens( +def build_group_ablated_batch( batch: Dict[str, torch.Tensor], - token_ids: Sequence[int], -) -> Dict[str, torch.Tensor]: - """Return a batch with selected disease tokens removed from event history.""" - selected = {int(token) for token in token_ids} - if not selected: - return batch + group_names: Sequence[str], + organ_groups: dict[str, list[int]], + occurred: torch.Tensor, +) -> tuple[Dict[str, torch.Tensor], list[str]]: + """Build a group-major batch for organ/system mortality attribution.""" + active_groups: list[str] = [] + for group in group_names: + ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device) + if ids.numel() > 0 and bool(occurred[:, ids].any().item()): + active_groups.append(group) + + if not active_groups: + return {}, [] event_rows: list[torch.Tensor] = [] time_rows: list[torch.Tensor] = [] @@ -306,47 +313,50 @@ def ablate_event_history_for_tokens( time_seq = batch["time_seq"] readout_mask = batch["readout_mask"] padding_mask = batch["padding_mask"].bool() - for i in range(event_seq.shape[0]): - valid = padding_mask[i] - events = event_seq[i, valid] - times = time_seq[i, valid] - reads = readout_mask[i, valid] - keep = torch.ones_like(events, dtype=torch.bool) - for token in selected: - keep &= events != int(token) + for group in active_groups: + selected = {int(token) for token in organ_groups[group]} + for i in range(event_seq.shape[0]): + valid = padding_mask[i] + events = event_seq[i, valid] + times = time_seq[i, valid] + reads = readout_mask[i, valid] + keep = torch.ones_like(events, dtype=torch.bool) + for token in selected: + keep &= events != int(token) - kept_events = events[keep] - kept_times = times[keep] - kept_reads = reads[keep] - if kept_events.numel() == 0: - kept_events = torch.tensor( - [CHECKUP_IDX], - dtype=event_seq.dtype, - device=event_seq.device, - ) - kept_times = batch["t_query"][i : i + 1].to( - dtype=time_seq.dtype, - device=time_seq.device, - ) - kept_reads = torch.ones(1, dtype=torch.bool, device=readout_mask.device) + kept_events = events[keep] + kept_times = times[keep] + kept_reads = reads[keep] + if kept_events.numel() == 0: + kept_events = torch.tensor( + [CHECKUP_IDX], + dtype=event_seq.dtype, + device=event_seq.device, + ) + kept_times = batch["t_query"][i : i + 1].to( + dtype=time_seq.dtype, + device=time_seq.device, + ) + kept_reads = torch.ones(1, dtype=torch.bool, device=readout_mask.device) - if bool(kept_reads.any()): - landmark_pos = torch.nonzero(kept_reads, as_tuple=False)[-1, 0] - else: - landmark_pos = torch.tensor( - int(kept_events.numel() - 1), - dtype=batch["landmark_pos"].dtype, - device=batch["landmark_pos"].device, - ) - kept_reads = torch.zeros_like(kept_events, dtype=torch.bool) - kept_reads[int(landmark_pos.item())] = True + if bool(kept_reads.any()): + landmark_pos = torch.nonzero(kept_reads, as_tuple=False)[-1, 0] + else: + landmark_pos = torch.tensor( + int(kept_events.numel() - 1), + dtype=batch["landmark_pos"].dtype, + device=batch["landmark_pos"].device, + ) + kept_reads = torch.zeros_like(kept_events, dtype=torch.bool) + kept_reads[int(landmark_pos.item())] = True - event_rows.append(kept_events) - time_rows.append(kept_times) - readout_rows.append(kept_reads) - landmark_positions.append(landmark_pos.to(dtype=batch["landmark_pos"].dtype)) + event_rows.append(kept_events) + time_rows.append(kept_times) + readout_rows.append(kept_reads) + landmark_positions.append(landmark_pos.to(dtype=batch["landmark_pos"].dtype)) - out = dict(batch) + repeat_count = len(active_groups) + out: Dict[str, torch.Tensor] = {} out["event_seq"] = pad_sequence(event_rows, batch_first=True, padding_value=PAD_IDX) out["time_seq"] = pad_sequence(time_rows, batch_first=True, padding_value=0.0) out["readout_mask"] = pad_sequence( @@ -354,7 +364,33 @@ def ablate_event_history_for_tokens( ) out["padding_mask"] = out["event_seq"] > PAD_IDX out["landmark_pos"] = torch.stack(landmark_positions) - return out + + repeated_keys = ( + "sex", + "other_type", + "other_value", + "other_value_kind", + "other_time", + "t_query", + "patient_id", + "landmark_age", + "followup_end_time", + "death_time", + "row_idx", + ) + for key in repeated_keys: + value = batch[key] + repeats = (repeat_count,) + (1,) * (value.ndim - 1) + out[key] = value.repeat(repeats) + return out, active_groups + + +def slice_tensor_batch( + batch: Dict[str, torch.Tensor], + start: int, + stop: int, +) -> Dict[str, torch.Tensor]: + return {key: value[start:stop] for key, value in batch.items()} @torch.no_grad() @@ -503,6 +539,12 @@ def parse_args() -> argparse.Namespace: 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( + "--attribution_batch_size", + type=int, + default=None, + help="Forward batch size for expanded organ/system ablation queries.", + ) parser.add_argument("--num_workers", type=int, default=None) parser.add_argument("--device", type=str, default=None) parser.add_argument("--extra_info_types", type=str, default=None) @@ -578,6 +620,11 @@ def main() -> None: model.eval() batch_size = int(cfg_get(args, cfg, "batch_size", 128)) + attribution_batch_size = int( + cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 4, batch_size)) + ) + if attribution_batch_size <= 0: + raise ValueError("attribution_batch_size must be positive") num_workers = int(cfg_get(args, cfg, "num_workers", 4)) loader = DataLoader( IndexedLandmarkDataset(landmark_dataset), @@ -602,6 +649,7 @@ def main() -> None: print(f"Death token: {death_idx}") print(f"Organ/system groups: {len(group_names)}") print(f"Landmark rows: {len(landmark_dataset)}") + print(f"Attribution batch size: {attribution_batch_size}") print(f"Output: {output_path}") rows: list[dict[str, Any]] = [] @@ -644,32 +692,53 @@ def main() -> None: group_mortality_attr_prob: dict[str, np.ndarray] = {} group_mortality_attr_hazard: dict[str, np.ndarray] = {} + batch_n = int(batch["event_seq"].shape[0]) + zeros = np.zeros(batch_n, dtype=np.float32) for group in group_names: - ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=device) - if ids.numel() == 0 or not bool(occurred[:, ids].any().item()): - zeros = np.zeros(batch["event_seq"].shape[0], dtype=np.float32) - group_mortality_attr_prob[group] = zeros - group_mortality_attr_hazard[group] = zeros - continue + group_mortality_attr_prob[group] = zeros.copy() + group_mortality_attr_hazard[group] = zeros.copy() - ablated_batch = ablate_event_history_for_tokens(batch, organ_groups[group]) - ablated_death_risk = death_risk_for_batch( - model=model, - batch=ablated_batch, - device=device, - model_target_mode=model_target_mode, - readout_name=readout_name, - readout_reduce=readout_reduce, - dist_mode=dist_mode, - tau=tau, + ablated_batch, active_groups = build_group_ablated_batch( + batch=batch, + group_names=group_names, + organ_groups=organ_groups, + occurred=occurred, + ) + if active_groups: + ablated_risk_chunks: list[torch.Tensor] = [] + expanded_n = int(ablated_batch["event_seq"].shape[0]) + for start in range(0, expanded_n, attribution_batch_size): + chunk = slice_tensor_batch( + ablated_batch, + start, + min(start + attribution_batch_size, expanded_n), + ) + ablated_risk_chunks.append( + death_risk_for_batch( + model=model, + batch=chunk, + device=device, + model_target_mode=model_target_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + dist_mode=dist_mode, + tau=tau, + ) + ) + + ablated_death_risk = torch.cat(ablated_risk_chunks, dim=0).view( + len(active_groups), batch_n ) ablated_death_hazard = mortality_hazard_from_risk(ablated_death_risk) - group_mortality_attr_prob[group] = ( - death_risk_tensor - ablated_death_risk + active_attr_hazard = ( + death_hazard_tensor[None, :] - ablated_death_hazard ).detach().cpu().numpy() - group_mortality_attr_hazard[group] = ( - death_hazard_tensor - ablated_death_hazard + active_attr_prob = ( + death_risk_tensor[None, :] - ablated_death_risk ).detach().cpu().numpy() + for group_idx, group in enumerate(active_groups): + group_mortality_attr_prob[group] = active_attr_prob[group_idx] + group_mortality_attr_hazard[group] = active_attr_hazard[group_idx] row_indices = batch["row_idx"].cpu().numpy().astype(np.int64) for j, row_idx in enumerate(row_indices.tolist()):