Batch mortality attribution queries
This commit is contained in:
@@ -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],
|
batch: Dict[str, torch.Tensor],
|
||||||
token_ids: Sequence[int],
|
group_names: Sequence[str],
|
||||||
) -> Dict[str, torch.Tensor]:
|
organ_groups: dict[str, list[int]],
|
||||||
"""Return a batch with selected disease tokens removed from event history."""
|
occurred: torch.Tensor,
|
||||||
selected = {int(token) for token in token_ids}
|
) -> tuple[Dict[str, torch.Tensor], list[str]]:
|
||||||
if not selected:
|
"""Build a group-major batch for organ/system mortality attribution."""
|
||||||
return batch
|
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] = []
|
event_rows: list[torch.Tensor] = []
|
||||||
time_rows: list[torch.Tensor] = []
|
time_rows: list[torch.Tensor] = []
|
||||||
@@ -306,47 +313,50 @@ def ablate_event_history_for_tokens(
|
|||||||
time_seq = batch["time_seq"]
|
time_seq = batch["time_seq"]
|
||||||
readout_mask = batch["readout_mask"]
|
readout_mask = batch["readout_mask"]
|
||||||
padding_mask = batch["padding_mask"].bool()
|
padding_mask = batch["padding_mask"].bool()
|
||||||
for i in range(event_seq.shape[0]):
|
for group in active_groups:
|
||||||
valid = padding_mask[i]
|
selected = {int(token) for token in organ_groups[group]}
|
||||||
events = event_seq[i, valid]
|
for i in range(event_seq.shape[0]):
|
||||||
times = time_seq[i, valid]
|
valid = padding_mask[i]
|
||||||
reads = readout_mask[i, valid]
|
events = event_seq[i, valid]
|
||||||
keep = torch.ones_like(events, dtype=torch.bool)
|
times = time_seq[i, valid]
|
||||||
for token in selected:
|
reads = readout_mask[i, valid]
|
||||||
keep &= events != int(token)
|
keep = torch.ones_like(events, dtype=torch.bool)
|
||||||
|
for token in selected:
|
||||||
|
keep &= events != int(token)
|
||||||
|
|
||||||
kept_events = events[keep]
|
kept_events = events[keep]
|
||||||
kept_times = times[keep]
|
kept_times = times[keep]
|
||||||
kept_reads = reads[keep]
|
kept_reads = reads[keep]
|
||||||
if kept_events.numel() == 0:
|
if kept_events.numel() == 0:
|
||||||
kept_events = torch.tensor(
|
kept_events = torch.tensor(
|
||||||
[CHECKUP_IDX],
|
[CHECKUP_IDX],
|
||||||
dtype=event_seq.dtype,
|
dtype=event_seq.dtype,
|
||||||
device=event_seq.device,
|
device=event_seq.device,
|
||||||
)
|
)
|
||||||
kept_times = batch["t_query"][i : i + 1].to(
|
kept_times = batch["t_query"][i : i + 1].to(
|
||||||
dtype=time_seq.dtype,
|
dtype=time_seq.dtype,
|
||||||
device=time_seq.device,
|
device=time_seq.device,
|
||||||
)
|
)
|
||||||
kept_reads = torch.ones(1, dtype=torch.bool, device=readout_mask.device)
|
kept_reads = torch.ones(1, dtype=torch.bool, device=readout_mask.device)
|
||||||
|
|
||||||
if bool(kept_reads.any()):
|
if bool(kept_reads.any()):
|
||||||
landmark_pos = torch.nonzero(kept_reads, as_tuple=False)[-1, 0]
|
landmark_pos = torch.nonzero(kept_reads, as_tuple=False)[-1, 0]
|
||||||
else:
|
else:
|
||||||
landmark_pos = torch.tensor(
|
landmark_pos = torch.tensor(
|
||||||
int(kept_events.numel() - 1),
|
int(kept_events.numel() - 1),
|
||||||
dtype=batch["landmark_pos"].dtype,
|
dtype=batch["landmark_pos"].dtype,
|
||||||
device=batch["landmark_pos"].device,
|
device=batch["landmark_pos"].device,
|
||||||
)
|
)
|
||||||
kept_reads = torch.zeros_like(kept_events, dtype=torch.bool)
|
kept_reads = torch.zeros_like(kept_events, dtype=torch.bool)
|
||||||
kept_reads[int(landmark_pos.item())] = True
|
kept_reads[int(landmark_pos.item())] = True
|
||||||
|
|
||||||
event_rows.append(kept_events)
|
event_rows.append(kept_events)
|
||||||
time_rows.append(kept_times)
|
time_rows.append(kept_times)
|
||||||
readout_rows.append(kept_reads)
|
readout_rows.append(kept_reads)
|
||||||
landmark_positions.append(landmark_pos.to(dtype=batch["landmark_pos"].dtype))
|
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["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["time_seq"] = pad_sequence(time_rows, batch_first=True, padding_value=0.0)
|
||||||
out["readout_mask"] = pad_sequence(
|
out["readout_mask"] = pad_sequence(
|
||||||
@@ -354,7 +364,33 @@ def ablate_event_history_for_tokens(
|
|||||||
)
|
)
|
||||||
out["padding_mask"] = out["event_seq"] > PAD_IDX
|
out["padding_mask"] = out["event_seq"] > PAD_IDX
|
||||||
out["landmark_pos"] = torch.stack(landmark_positions)
|
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()
|
@torch.no_grad()
|
||||||
@@ -503,6 +539,12 @@ def parse_args() -> argparse.Namespace:
|
|||||||
parser.add_argument("--tau", 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("--min_history_events", type=int, default=None)
|
||||||
parser.add_argument("--batch_size", 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("--num_workers", type=int, default=None)
|
||||||
parser.add_argument("--device", type=str, default=None)
|
parser.add_argument("--device", type=str, default=None)
|
||||||
parser.add_argument("--extra_info_types", type=str, default=None)
|
parser.add_argument("--extra_info_types", type=str, default=None)
|
||||||
@@ -578,6 +620,11 @@ def main() -> None:
|
|||||||
model.eval()
|
model.eval()
|
||||||
|
|
||||||
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
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))
|
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
|
||||||
loader = DataLoader(
|
loader = DataLoader(
|
||||||
IndexedLandmarkDataset(landmark_dataset),
|
IndexedLandmarkDataset(landmark_dataset),
|
||||||
@@ -602,6 +649,7 @@ def main() -> None:
|
|||||||
print(f"Death token: {death_idx}")
|
print(f"Death token: {death_idx}")
|
||||||
print(f"Organ/system groups: {len(group_names)}")
|
print(f"Organ/system groups: {len(group_names)}")
|
||||||
print(f"Landmark rows: {len(landmark_dataset)}")
|
print(f"Landmark rows: {len(landmark_dataset)}")
|
||||||
|
print(f"Attribution batch size: {attribution_batch_size}")
|
||||||
print(f"Output: {output_path}")
|
print(f"Output: {output_path}")
|
||||||
|
|
||||||
rows: list[dict[str, Any]] = []
|
rows: list[dict[str, Any]] = []
|
||||||
@@ -644,32 +692,53 @@ def main() -> None:
|
|||||||
|
|
||||||
group_mortality_attr_prob: dict[str, np.ndarray] = {}
|
group_mortality_attr_prob: dict[str, np.ndarray] = {}
|
||||||
group_mortality_attr_hazard: 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:
|
for group in group_names:
|
||||||
ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=device)
|
group_mortality_attr_prob[group] = zeros.copy()
|
||||||
if ids.numel() == 0 or not bool(occurred[:, ids].any().item()):
|
group_mortality_attr_hazard[group] = zeros.copy()
|
||||||
zeros = np.zeros(batch["event_seq"].shape[0], dtype=np.float32)
|
|
||||||
group_mortality_attr_prob[group] = zeros
|
|
||||||
group_mortality_attr_hazard[group] = zeros
|
|
||||||
continue
|
|
||||||
|
|
||||||
ablated_batch = ablate_event_history_for_tokens(batch, organ_groups[group])
|
ablated_batch, active_groups = build_group_ablated_batch(
|
||||||
ablated_death_risk = death_risk_for_batch(
|
batch=batch,
|
||||||
model=model,
|
group_names=group_names,
|
||||||
batch=ablated_batch,
|
organ_groups=organ_groups,
|
||||||
device=device,
|
occurred=occurred,
|
||||||
model_target_mode=model_target_mode,
|
)
|
||||||
readout_name=readout_name,
|
if active_groups:
|
||||||
readout_reduce=readout_reduce,
|
ablated_risk_chunks: list[torch.Tensor] = []
|
||||||
dist_mode=dist_mode,
|
expanded_n = int(ablated_batch["event_seq"].shape[0])
|
||||||
tau=tau,
|
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)
|
ablated_death_hazard = mortality_hazard_from_risk(ablated_death_risk)
|
||||||
group_mortality_attr_prob[group] = (
|
active_attr_hazard = (
|
||||||
death_risk_tensor - ablated_death_risk
|
death_hazard_tensor[None, :] - ablated_death_hazard
|
||||||
).detach().cpu().numpy()
|
).detach().cpu().numpy()
|
||||||
group_mortality_attr_hazard[group] = (
|
active_attr_prob = (
|
||||||
death_hazard_tensor - ablated_death_hazard
|
death_risk_tensor[None, :] - ablated_death_risk
|
||||||
).detach().cpu().numpy()
|
).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)
|
row_indices = batch["row_idx"].cpu().numpy().astype(np.int64)
|
||||||
for j, row_idx in enumerate(row_indices.tolist()):
|
for j, row_idx in enumerate(row_indices.tolist()):
|
||||||
|
|||||||
Reference in New Issue
Block a user