Task Contribution: Dynamic Survival Analysis for Early Event Prediction

docs/api/tasks.rst

@@ -230,3 +230,4 @@ Available Tasks
     Mutation Pathogenicity (COSMIC) <tasks/pyhealth.tasks.MutationPathogenicityPrediction>
     Cancer Survival Prediction (TCGA) <tasks/pyhealth.tasks.CancerSurvivalPrediction>
     Cancer Mutation Burden (TCGA) <tasks/pyhealth.tasks.CancerMutationBurden>
+    Dynamic Survival Analysis <tasks/pyhealth.tasks.dynamic_survival>

docs/api/tasks/pyhealth.tasks.dynamic_survival.rst

@@ -0,0 +1,68 @@
+DynamicSurvivalTask
+==================
+
+This module implements a dynamic survival analysis task for early event prediction.
+
+The task follows the anchor-based discrete-time survival formulation proposed in:
+
+Yèche et al. (2024), *Dynamic Survival Analysis for Early Event Prediction*.
+
+Key Features
+------------
+- Multiple anchors per patient
+- Discrete-time hazard prediction
+- Support for censoring
+- Configurable observation windows and anchor strategies
+
+Output Format
+-------------
+Each processed sample contains:
+
+- **patient_id**: unique patient identifier
+- **visit_id**: unique anchor-based visit ID
+- **x**: input features (temporal sequence)
+- **y**: hazard label vector (0/1)
+- **mask**: indicates valid risk set:
+    - 1 = patient is at risk at this timestep
+    - 0 = timestep excluded (post-event or post-censoring)
+
+Usage Example
+-------------
+
+.. code-block:: python
+
+    from pyhealth.tasks.dynamic_survival import DynamicSurvivalTask
+
+    # Minimal dataset wrapper (MockDataset or a real PyHealth dataset)
+    class MockDataset:
+        def __init__(self):
+            self.patients = {}
+
+    dataset = MockDataset()
+
+    task = DynamicSurvivalTask(
+        dataset=dataset,
+        observation_window=24,
+        horizon=24,
+        anchor_strategy="fixed",
+    )
+
+    # Apply to a patient object
+    samples = task(patient)
+
+Example Output
+--------------
+
+Each sample:
+
+- x: shape (T, d)
+- y: shape (horizon,)
+- mask: shape (horizon,)
+
+API Reference
+-------------
+
+.. autoclass:: pyhealth.tasks.dynamic_survival.DynamicSurvivalTask
+   :members:
+   :undoc-members:
+   :show-inheritance:

examples/dynamic_survival_ablation.py

@@ -0,0 +1,164 @@
+# Authors: Skyler Lehto (lehto2@illinois.edu),
+#          Ryan Bradley (ryancb3@illinois.edu),
+#          Weonah Choi (weonahc2@illinois.edu)
+# Paper: Dynamic Survival Analysis for Early Event Prediction (Yèche et al., 2024)
+# Link: https://arxiv.org/abs/2403.12818
+# Description: Ablation study for observation window size on synthetic patients.
+
+"""
+Ablation Study: Effect of Observation Window Length
+
+We vary observation window sizes (12, 24, 48 hours) and
+measure performance using masked BCE and MSE.
+
+This demonstrates how task configuration (NOT model complexity)
+impacts predictive performance.
+"""
+
+import sys
+import os
+from datetime import datetime, timedelta
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from pyhealth.tasks.dynamic_survival import DynamicSurvivalTask
+from examples.synthetic_dataset import generate_synthetic_dataset
+from examples.mock_ehr import MockEvent, MockVisit, MockPatient, MockDataset
+
+
+# Convert synthetic dict → MockPatient
+
+def convert_to_mock_patients(patients_dict):
+    base_time = datetime(2025, 1, 1)
+
+    mock_patients = []
+
+    for p in patients_dict:
+        visits_data = []
+
+        for v in p["visits"]:
+            visits_data.append({
+                "time": base_time + timedelta(days=v["time"]),
+                "diagnosis": ["0000"],  # dummy code for vocab
+            })
+
+        death_time = None
+        if p.get("outcome_time") is not None:
+            death_time = base_time + timedelta(days=p["outcome_time"])
+
+        mock_patients.append(
+            MockPatient(
+                pid=p["patient_id"],
+                visits_data=visits_data,
+                death_time=death_time,
+            )
+        )
+
+    return mock_patients
+
+
+# Model
+
+class SimpleModel(nn.Module):
+    def __init__(self, input_dim=2, hidden_dim=8, horizon=24):
+        super().__init__()
+        self.rnn = nn.GRU(input_dim, hidden_dim, batch_first=True)
+        self.fc = nn.Linear(hidden_dim, horizon)
+
+    def forward(self, x):
+        _, h = self.rnn(x)
+        return torch.sigmoid(self.fc(h.squeeze(0)))
+
+
+# Utils
+
+def prepare_batch(samples):
+    X, Y, M = [], [], []
+
+    for s in samples:
+        X.append(s["x"])
+        Y.append(s["y"])
+        M.append(s["mask"])
+
+    if len(X) == 0:
+        raise ValueError("No valid samples generated.")
+
+    max_len = max(len(x) for x in X)
+
+    X_pad = []
+    for x in X:
+        pad = np.zeros((max_len - len(x), x.shape[1]))
+        X_pad.append(np.vstack([x, pad]))
+
+    return (
+        torch.tensor(np.array(X_pad), dtype=torch.float32),
+        torch.tensor(np.array(Y), dtype=torch.float32),
+        torch.tensor(np.array(M), dtype=torch.float32),
+    )
+
+
+def train_and_eval(samples):
+    model = SimpleModel()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+    X, Y, M = prepare_batch(samples)
+
+    for _ in range(5):
+        pred = model(X)
+        loss = -(Y * torch.log(pred + 1e-8) +
+                 (1 - Y) * torch.log(1 - pred + 1e-8))
+        loss = (loss * M).sum() / M.sum()
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+    with torch.no_grad():
+        pred = model(X)
+
+        bce = -(Y * torch.log(pred + 1e-8) +
+                (1 - Y) * torch.log(1 - pred + 1e-8))
+        bce = (bce * M).sum() / M.sum()
+
+        mse = ((pred - Y) ** 2 * M).sum() / M.sum()
+
+    return {"bce": bce.item(), "mse": mse.item()}
+
+
+# Main Experiment
+
+def main():
+    patients_raw = generate_synthetic_dataset(50)
+    patients = convert_to_mock_patients(patients_raw)
+    dataset = MockDataset(patients)
+
+    windows = [12, 24, 48]
+    results = {}
+
+    print("\n=== Ablation Results ===")
+
+    for w in windows:
+        task = DynamicSurvivalTask(
+            dataset=dataset,
+            observation_window=w,
+            horizon=24,
+        )
+
+        samples = dataset.set_task(task)
+
+        if len(samples) == 0:
+            print(f"Skipping window={w}, no samples")
+            continue
+
+        score = train_and_eval(samples)
+        results[w] = score
+
+        print(f"Window={w} | BCE={score['bce']:.4f} | MSE={score['mse']:.4f}")
+
+
+if __name__ == "__main__":
+    main()