Add configurable minimum gap for MIMIC4 readmission

docs/api/tasks/pyhealth.tasks.readmission_prediction.rst

@@ -11,6 +11,27 @@
     :undoc-members:
     :show-inheritance:
 
+Minimum admission gap
+---------------------
+
+``ReadmissionPredictionMIMIC4`` supports an optional ``min_gap`` parameter to
+exclude very short intervals between discharge and the next admission from
+being labeled as readmissions. This can be useful when short gaps may reflect
+internal transfers rather than a new hospital admission.
+
+For example:
+
+.. code-block:: python
+
+    from datetime import timedelta
+    from pyhealth.tasks import ReadmissionPredictionMIMIC4
+
+    task = ReadmissionPredictionMIMIC4(
+        window=timedelta(days=30),
+        min_gap=timedelta(hours=3),
+    )
+
+To maintain backwards compatibility, ``min_gap`` defaults to ``None``, meaning that some 'readmissions' are in fact internal transfers.
 .. autoclass:: pyhealth.tasks.readmission_prediction.ReadmissionPredictionEICU
     :members:
     :undoc-members:

examples/readmission/readmission_mimic4_rnn.py

@@ -1,4 +1,5 @@
 import tempfile
+from datetime import timedelta
 
 from pyhealth.datasets import MIMIC4Dataset
 from pyhealth.datasets import split_by_patient, get_dataloader
@@ -18,7 +19,9 @@
     base_dataset.stats()
 
     # STEP 2: Set task
-    task = ReadmissionPredictionMIMIC4()
+    task = ReadmissionPredictionMIMIC4(
+        min_gap=timedelta(hours=3),
+    )
     sample_dataset = base_dataset.set_task(task)
 
     # STEP 3: Split and create dataloaders

pyhealth/tasks/readmission_prediction.py

@@ -175,7 +175,11 @@ class ReadmissionPredictionMIMIC4(BaseTask):
     output_schema: Dict[str, str] = {"readmission": "binary"}
 
     def __init__(
-        self, window: timedelta = timedelta(days=15), exclude_minors: bool = True, **kwargs
+        self,
+        window: timedelta = timedelta(days=15),
+        exclude_minors: bool = True,
+        min_gap: timedelta | None = None,
+        **kwargs,
     ) -> None:
         """Initializes the task object.
 
@@ -184,11 +188,17 @@ def __init__(
                 considered a readmission. Defaults to 15 days.
             exclude_minors: Whether to exclude patients whose
                 ``anchor_age`` is less than 18. Defaults to True.
+            min_gap: Minimum time between discharge and the next admission
+                for it to be considered a readmission. Admissions occuring within the window of discharge time + ``min_gap`` are considered as a single admission. Defaults to None, which applies no
+                minimum gap and treats some internal admissions as readmissions.
+
+                Consider setting ``min_gap`` to avoid this behaviour.
             **kwargs: Passed to :class:`~pyhealth.tasks.BaseTask`, e.g.
                 ``code_mapping``.
         """
         super().__init__(**kwargs)
         self.window = window
+        self.min_gap = min_gap
         self.exclude_minors = exclude_minors
 
     def __call__(self, patient: Patient) -> List[Dict]:
@@ -260,8 +270,11 @@ def __call__(self, patient: Patient) -> List[Dict]:
             except ValueError:
                 discharge_time = datetime.strptime(admissions[i].dischtime, "%Y-%m-%d")
 
+            gap = admissions[i + 1].timestamp - discharge_time
+
             readmission = int(
-                (admissions[i + 1].timestamp - discharge_time) < self.window
+                gap < self.window
+                and (self.min_gap is None or gap > self.min_gap)
             )
 
             samples.append(

tests/core/test_mimic4_readmission_prediction.py

@@ -20,12 +20,19 @@ def setUpClass(cls):
 
         cls.samples15days = dataset.set_task(ReadmissionPredictionMIMIC4(window=timedelta(days=15)))
         cls.samples5days = dataset.set_task(ReadmissionPredictionMIMIC4(window=timedelta(days=5)))
+        cls.samples5daymingap = dataset.set_task(
+            ReadmissionPredictionMIMIC4(
+                window=timedelta(days=15),
+                min_gap=timedelta(days=5),
+            )
+        )
         cls.sampleswithminors = dataset.set_task(ReadmissionPredictionMIMIC4(exclude_minors=False))
 
     @classmethod
     def tearDownClass(cls):
         cls.samples15days.close()
         cls.samples5days.close()
+        cls.samples5daymingap.close()
         cls.sampleswithminors.close()
 
     def test_task_schema(self):
@@ -44,6 +51,14 @@ def test_task_defaults(self):
 
         self.assertEqual(task.window, timedelta(days=15))
         self.assertTrue(task.exclude_minors)
+        self.assertIsNone(task.min_gap)
+
+    def test_positional_arguments_backward_compatible(self):
+        task = ReadmissionPredictionMIMIC4(timedelta(days=30), False)
+
+        self.assertEqual(task.window, timedelta(days=30))
+        self.assertFalse(task.exclude_minors)
+        self.assertIsNone(task.min_gap)
 
     def test_sample_schema(self):
         for sample in self.samples15days:
@@ -73,6 +88,24 @@ def test_time_window(self):
         self.assertTrue(bool(readmitted5days[0].item()))
         self.assertTrue(bool(readmitted15days[0].item()))
 
+    def test_min_gap(self):
+        readmitted3 = [
+            s["readmission"]
+            for s in self.samples5daymingap
+            if s["visit_id"] == "3"
+        ]
+        readmitted4 = [
+            s["readmission"]
+            for s in self.samples5daymingap
+            if s["visit_id"] == "4"
+        ]
+
+        self.assertEqual(len(readmitted3), 1)
+        self.assertEqual(len(readmitted4), 1)
+
+        self.assertTrue(bool(readmitted3[0].item()))
+        self.assertFalse(bool(readmitted4[0].item()))
+
     def test_last_admission_is_excluded(self):
         self.assertNotIn("2", [ s["visit_id"] for s in self.samples15days ])
         self.assertNotIn("2", [ s["visit_id"] for s in self.samples5days ])