Add disease history ablation modes

tests/test_dataset_checkup.py

@@ -0,0 +1,57 @@
+import unittest
+from pathlib import Path
+
+import numpy as np
+
+from dataset import _ExpoBaseDataset
+from targets import CHECKUP_IDX
+from train_util import load_extra_info_types_file
+
+
+class CheckupSelectionTests(unittest.TestCase):
+    @staticmethod
+    def _base(extra_info_types):
+        dataset = _ExpoBaseDataset.__new__(_ExpoBaseDataset)
+        dataset.extra_info_types = list(extra_info_types)
+        dataset.event_data = np.asarray(
+            [
+                [101, 10, CHECKUP_IDX],
+                [101, 20, 2],
+                [101, 30, 3],
+            ],
+            dtype=np.float64,
+        )
+        return dataset
+
+    def test_explicit_empty_extra_info_removes_checkup(self):
+        project_root = Path(__file__).resolve().parents[1]
+        selected_types = load_extra_info_types_file(
+            str(project_root / "extra_info_types_none.txt")
+        )
+        self.assertEqual(selected_types, [])
+        dataset = self._base(selected_types)
+
+        rows = list(dataset._iter_patient_events(impute_no_event_gaps=False))
+
+        self.assertEqual(len(rows), 1)
+        eid, times, labels = rows[0]
+        self.assertEqual(eid, 101)
+        np.testing.assert_array_equal(times, np.asarray([20, 30], dtype=np.float32))
+        self.assertNotIn(CHECKUP_IDX, labels.tolist())
+
+    def test_selected_extra_info_keeps_checkup(self):
+        dataset = self._base([11])
+
+        rows = list(dataset._iter_patient_events(impute_no_event_gaps=False))
+
+        self.assertEqual(len(rows), 1)
+        _, times, labels = rows[0]
+        np.testing.assert_array_equal(
+            times,
+            np.asarray([10, 20, 30], dtype=np.float32),
+        )
+        self.assertEqual(labels[0], CHECKUP_IDX)
+
+
+if __name__ == "__main__":
+    unittest.main()