Stream mortality attribution batches
This commit is contained in:
@@ -288,82 +288,57 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st
|
||||
}
|
||||
|
||||
|
||||
def build_group_ablated_batch(
|
||||
def build_group_ablated_slice(
|
||||
batch: Dict[str, torch.Tensor],
|
||||
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] = []
|
||||
readout_rows: list[torch.Tensor] = []
|
||||
landmark_positions: list[torch.Tensor] = []
|
||||
|
||||
token_ids: Sequence[int],
|
||||
row_start: int,
|
||||
row_stop: int,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Build one fixed-width ablated slice without rebuilding variable-length rows."""
|
||||
event_seq = batch["event_seq"]
|
||||
time_seq = batch["time_seq"]
|
||||
readout_mask = batch["readout_mask"]
|
||||
padding_mask = batch["padding_mask"].bool()
|
||||
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)
|
||||
|
||||
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))
|
||||
|
||||
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(
|
||||
readout_rows, batch_first=True, padding_value=False
|
||||
out["event_seq"] = event_seq[row_start:row_stop].clone()
|
||||
out["time_seq"] = batch["time_seq"][row_start:row_stop]
|
||||
out["readout_mask"] = batch["readout_mask"][row_start:row_stop].clone()
|
||||
out["padding_mask"] = batch["padding_mask"][row_start:row_stop].bool().clone()
|
||||
out["landmark_pos"] = batch["landmark_pos"][row_start:row_stop].clone()
|
||||
|
||||
seq_len = int(event_seq.shape[1])
|
||||
positions = torch.arange(seq_len, device=event_seq.device)[None, :]
|
||||
ids = torch.as_tensor(token_ids, dtype=event_seq.dtype, device=event_seq.device)
|
||||
remove = torch.isin(out["event_seq"], ids) & out["padding_mask"]
|
||||
out["event_seq"] = torch.where(
|
||||
remove,
|
||||
torch.full_like(out["event_seq"], PAD_IDX),
|
||||
out["event_seq"],
|
||||
)
|
||||
out["padding_mask"] = out["event_seq"] > PAD_IDX
|
||||
out["landmark_pos"] = torch.stack(landmark_positions)
|
||||
out["padding_mask"] &= ~remove
|
||||
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_start + 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)
|
||||
|
||||
repeated_keys = (
|
||||
"sex",
|
||||
@@ -379,18 +354,62 @@ def build_group_ablated_batch(
|
||||
"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
|
||||
out[key] = batch[key][row_start:row_stop]
|
||||
return out
|
||||
|
||||
|
||||
def slice_tensor_batch(
|
||||
def concat_tensor_batches(chunks: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
||||
return {
|
||||
key: torch.cat([chunk[key] for chunk in chunks], dim=0)
|
||||
for key in chunks[0]
|
||||
}
|
||||
|
||||
|
||||
def iter_group_ablated_batches(
|
||||
batch: Dict[str, torch.Tensor],
|
||||
start: int,
|
||||
stop: int,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
return {key: value[start:stop] for key, value in batch.items()}
|
||||
group_names: Sequence[str],
|
||||
organ_groups: dict[str, list[int]],
|
||||
occurred: torch.Tensor,
|
||||
max_batch_size: int,
|
||||
):
|
||||
"""Yield ablated chunks as soon as enough rows are available for a forward pass."""
|
||||
batch_n = int(batch["event_seq"].shape[0])
|
||||
pending_batches: list[Dict[str, torch.Tensor]] = []
|
||||
pending_groups: list[str] = []
|
||||
pending_rows: list[int] = []
|
||||
pending_n = 0
|
||||
|
||||
for group in group_names:
|
||||
ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device)
|
||||
if ids.numel() == 0 or not bool(occurred[:, ids].any().item()):
|
||||
continue
|
||||
|
||||
row_start = 0
|
||||
while row_start < batch_n:
|
||||
capacity = int(max_batch_size) - pending_n
|
||||
row_stop = min(batch_n, row_start + capacity)
|
||||
chunk = build_group_ablated_slice(
|
||||
batch=batch,
|
||||
token_ids=organ_groups[group],
|
||||
row_start=row_start,
|
||||
row_stop=row_stop,
|
||||
)
|
||||
chunk_n = int(row_stop - row_start)
|
||||
pending_batches.append(chunk)
|
||||
pending_groups.extend([group] * chunk_n)
|
||||
pending_rows.extend(range(row_start, row_stop))
|
||||
pending_n += chunk_n
|
||||
row_start = row_stop
|
||||
|
||||
if pending_n >= int(max_batch_size):
|
||||
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
|
||||
pending_batches = []
|
||||
pending_groups = []
|
||||
pending_rows = []
|
||||
pending_n = 0
|
||||
|
||||
if pending_batches:
|
||||
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -646,6 +665,7 @@ def main() -> None:
|
||||
print(f"Landmark ages: {landmark_ages.tolist()}")
|
||||
print(f"Tau: {tau:g} years")
|
||||
print(f"Dist mode: {dist_mode}")
|
||||
print(f"Device: {device}")
|
||||
print(f"Death token: {death_idx}")
|
||||
print(f"Organ/system groups: {len(group_names)}")
|
||||
print(f"Landmark rows: {len(landmark_dataset)}")
|
||||
@@ -698,25 +718,16 @@ def main() -> None:
|
||||
group_mortality_attr_prob[group] = zeros.copy()
|
||||
group_mortality_attr_hazard[group] = zeros.copy()
|
||||
|
||||
ablated_batch, active_groups = build_group_ablated_batch(
|
||||
for ablated_chunk, chunk_groups, chunk_rows in iter_group_ablated_batches(
|
||||
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(
|
||||
max_batch_size=attribution_batch_size,
|
||||
):
|
||||
ablated_death_risk = death_risk_for_batch(
|
||||
model=model,
|
||||
batch=chunk,
|
||||
batch=ablated_chunk,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
@@ -724,21 +735,17 @@ def main() -> None:
|
||||
dist_mode=dist_mode,
|
||||
tau=tau,
|
||||
)
|
||||
)
|
||||
|
||||
ablated_death_risk = torch.cat(ablated_risk_chunks, dim=0).view(
|
||||
len(active_groups), batch_n
|
||||
)
|
||||
row_tensor = torch.as_tensor(chunk_rows, dtype=torch.long, device=device)
|
||||
ablated_death_hazard = mortality_hazard_from_risk(ablated_death_risk)
|
||||
active_attr_hazard = (
|
||||
death_hazard_tensor[None, :] - ablated_death_hazard
|
||||
attr_prob = (
|
||||
death_risk_tensor[row_tensor] - ablated_death_risk
|
||||
).detach().cpu().numpy()
|
||||
active_attr_prob = (
|
||||
death_risk_tensor[None, :] - ablated_death_risk
|
||||
attr_hazard = (
|
||||
death_hazard_tensor[row_tensor] - ablated_death_hazard
|
||||
).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]
|
||||
for local_idx, (group, row_idx) in enumerate(zip(chunk_groups, chunk_rows)):
|
||||
group_mortality_attr_prob[group][row_idx] = attr_prob[local_idx]
|
||||
group_mortality_attr_hazard[group][row_idx] = attr_hazard[local_idx]
|
||||
|
||||
row_indices = batch["row_idx"].cpu().numpy().astype(np.int64)
|
||||
for j, row_idx in enumerate(row_indices.tolist()):
|
||||
|
||||
Reference in New Issue
Block a user