JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix mortality attribution active rows

Browse Source
This commit is contained in:
Jiarui Li
2026-06-27 14:29:32 +08:00
parent 2ae92da0ec
commit 01629cec56
1 changed files with 25 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
evaluate_event_free_survival.py
View File

@@ -291,18 +291,17 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st
def build_group_ablated_slice(
batch: Dict[str, torch.Tensor],
token_ids: Sequence[int],
row_start: int,
row_stop: int,
row_indices: torch.Tensor,
) -> Dict[str, torch.Tensor]:
"""Build one fixed-width ablated slice without rebuilding variable-length rows."""
event_seq = batch["event_seq"]
out: Dict[str, torch.Tensor] = {}
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()
out["event_seq"] = event_seq[row_indices].clone()
out["time_seq"] = batch["time_seq"][row_indices]
out["readout_mask"] = batch["readout_mask"][row_indices].clone()
out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
seq_len = int(event_seq.shape[1])
positions = torch.arange(seq_len, device=event_seq.device)[None, :]
@@ -320,9 +319,9 @@ def build_group_ablated_slice(
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["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
@@ -354,7 +353,7 @@ def build_group_ablated_slice(
"row_idx",
)
for key in repeated_keys:
out[key] = batch[key][row_start:row_stop]
out[key] = batch[key][row_indices]
return out
@@ -373,7 +372,6 @@ def iter_group_ablated_batches(
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] = []
@@ -381,25 +379,28 @@ def iter_group_ablated_batches(
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()):
if ids.numel() == 0:
continue
active_rows = torch.nonzero(occurred[:, ids].any(dim=1), as_tuple=False).flatten()
if active_rows.numel() == 0:
continue
row_start = 0
while row_start < batch_n:
row_offset = 0
while row_offset < int(active_rows.numel()):
capacity = int(max_batch_size) - pending_n
row_stop = min(batch_n, row_start + capacity)
row_stop = min(int(active_rows.numel()), row_offset + capacity)
row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
chunk = build_group_ablated_slice(
batch=batch,
token_ids=organ_groups[group],
row_start=row_start,
row_stop=row_stop,
row_indices=row_indices,
)
chunk_n = int(row_stop - row_start)
chunk_n = int(row_indices.numel())
pending_batches.append(chunk)
pending_groups.extend([group] * chunk_n)
pending_rows.extend(range(row_start, row_stop))
pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist())
pending_n += chunk_n
row_start = row_stop
row_offset = row_stop
if pending_n >= int(max_batch_size):
yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
@@ -774,6 +775,9 @@ def main() -> None:
for group in group_names:
out[f"history_count__{group}"] = int(group_counts[group])
out[f"new_disease_risk__{group}"] = float(group_risk[group][j])
if int(group_counts[group]) == 0:
group_mortality_attr_prob[group][j] = 0.0
group_mortality_attr_hazard[group][j] = 0.0
out[f"mortality_attribution_probability__{group}"] = float(
group_mortality_attr_prob[group][j]
)