Fix the MIMIC4 LoS and DKA task unit tests (#813)

* Fix MIMIC4 LoS prediction unit tests

* Fix MIMIC4 DKA prediction unit tests

* Remove icustays.csv to avoid merge conflict

Author: EricSchrock

test-resources/core/mimic4demo/hosp/admissions.csv

@@ -1,4 +1,6 @@
 subject_id,hadm_id,admittime,dischtime,admission_type,admission_location,insurance,language,marital_status,race,discharge_location,hospital_expire_flag
+10001,19999,2150-02-15 08:00:00,2150-02-18 14:00:00,ELECTIVE,PHYSICIAN REFERRAL,Medicare,ENGLISH,MARRIED,WHITE,HOME,0
+10002,20000,2151-01-01 07:00:00,2151-01-02 11:00:00,URGENT,TRANSFER FROM HOSPITAL,Medicaid,ENGLISH,SINGLE,BLACK,HOME,0
 10001,20001,2150-03-15 08:00:00,2150-03-18 14:00:00,ELECTIVE,PHYSICIAN REFERRAL,Medicare,ENGLISH,MARRIED,WHITE,HOME,0
 10001,20002,2150-06-20 10:30:00,2150-06-25 16:00:00,EMERGENCY,EMERGENCY ROOM,Medicare,ENGLISH,MARRIED,WHITE,HOME,0
 10002,20003,2151-01-10 07:00:00,2151-01-12 11:00:00,URGENT,TRANSFER FROM HOSPITAL,Medicaid,ENGLISH,SINGLE,BLACK,HOME,0

test-resources/core/mimic4demo/hosp/diagnoses_icd.csv

@@ -1,4 +1,6 @@
 subject_id,hadm_id,seq_num,icd_code,icd_version
+10001,19999,1,E1065,10
+10002,20000,1,E1065,10
 10001,20001,1,E1010,10
 10001,20001,2,E1165,10
 10001,20001,3,I10,10

tests/core/test_mimic4_dka.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 
 from pyhealth.datasets import MIMIC4Dataset
-from pyhealth.tasks.dka import DKAPredictionMIMIC4
+from pyhealth.tasks.dka import DKAPredictionMIMIC4, T1DDKAPredictionMIMIC4
 
 
 class TestMIMIC4DKAPrediction(unittest.TestCase):
@@ -17,7 +17,7 @@ def setUp(self):
     def _setup_dataset_path(self):
         """Get path to local MIMIC-IV demo dataset in test resources."""
         # Get the path to the test-resources/core/mimic4demo directory
-        test_dir = Path(__file__).parent.parent
+        test_dir = Path(__file__).parent.parent.parent
         self.demo_dataset_path = str(
             test_dir / "test-resources" / "core" / "mimic4demo"
         )
@@ -26,12 +26,6 @@ def _setup_dataset_path(self):
         print(f"Setting up MIMIC-IV demo dataset for DKA prediction")
         print(f"Dataset path: {self.demo_dataset_path}")
 
-        # Verify the dataset exists
-        if not os.path.exists(self.demo_dataset_path):
-            raise unittest.SkipTest(
-                f"MIMIC-IV demo dataset not found at {self.demo_dataset_path}"
-            )
-
         # List files in the hosp directory
         hosp_path = os.path.join(self.demo_dataset_path, "hosp")
         if os.path.exists(hosp_path):
@@ -47,9 +41,9 @@ def _load_dataset(self):
         """Load the dataset for testing."""
         tables = ["diagnoses_icd", "procedures_icd", "prescriptions", "labevents"]
         print(f"Loading MIMIC4Dataset with tables: {tables}")
-        self.dataset = MIMIC4Dataset(root=self.demo_dataset_path, tables=tables)
+        self.dataset = MIMIC4Dataset(ehr_root=self.demo_dataset_path, ehr_tables=tables)
         print(f"✓ Dataset loaded successfully")
-        print(f"  Total patients: {len(self.dataset.patients)}")
+        print(f"  Total patients: {len(self.dataset.unique_patient_ids)}")
         print()
 
     def test_dataset_stats(self):
@@ -75,9 +69,33 @@ def test_dka_prediction_task_initialization(self):
         print("Testing default initialization...")
         task = DKAPredictionMIMIC4()
         self.assertEqual(task.task_name, "DKAPredictionMIMIC4")
+        self.assertEqual(task.padding, 0)
+        self.assertIn("icd_codes", task.input_schema)
+        self.assertIn("labs", task.input_schema)
+        self.assertIn("label", task.output_schema)
+        print(f"✓ Default task initialized: {task.task_name}")
+        print(f"  Input schema: {list(task.input_schema.keys())}")
+        print(f"  Output schema: {list(task.output_schema.keys())}")
+
+        # Test custom initialization
+        print("\nTesting custom initialization...")
+        custom_task = DKAPredictionMIMIC4(padding=5)
+        self.assertEqual(custom_task.padding, 5)
+        print(f"✓ Custom task initialized with padding={custom_task.padding}")
+
+    def test_t1ddka_prediction_task_initialization(self):
+        """Test T1DDKAPredictionMIMIC4 task initialization."""
+        print(f"\n{'='*60}")
+        print("TEST: test_t1ddka_prediction_task_initialization()")
+        print(f"{'='*60}")
+
+        # Test default initialization
+        print("Testing default initialization...")
+        task = T1DDKAPredictionMIMIC4()
+        self.assertEqual(task.task_name, "T1DDKAPredictionMIMIC4")
         self.assertEqual(task.dka_window_days, 90)
         self.assertEqual(task.padding, 0)
-        self.assertIn("diagnoses", task.input_schema)
+        self.assertIn("icd_codes", task.input_schema)
         self.assertIn("labs", task.input_schema)
         self.assertIn("label", task.output_schema)
         print(f"✓ Default task initialized: {task.task_name}")
@@ -87,7 +105,7 @@ def test_dka_prediction_task_initialization(self):
 
         # Test custom initialization
         print("\nTesting custom initialization...")
-        custom_task = DKAPredictionMIMIC4(dka_window_days=30, padding=5)
+        custom_task = T1DDKAPredictionMIMIC4(dka_window_days=30, padding=5)
         self.assertEqual(custom_task.dka_window_days, 30)
         self.assertEqual(custom_task.padding, 5)
         print(f"✓ Custom task initialized with window={custom_task.dka_window_days}, padding={custom_task.padding}")
@@ -98,24 +116,37 @@ def test_dka_prediction_class_variables(self):
         print("TEST: test_dka_prediction_class_variables()")
         print(f"{'='*60}")
 
-        # Test T1DM codes
-        self.assertEqual(DKAPredictionMIMIC4.T1DM_ICD10_PREFIX, "E10")
-        self.assertIn("25001", DKAPredictionMIMIC4.T1DM_ICD9_CODES)
-        print(f"✓ T1DM ICD-10 prefix: {DKAPredictionMIMIC4.T1DM_ICD10_PREFIX}")
-        print(f"✓ T1DM ICD-9 codes: {len(DKAPredictionMIMIC4.T1DM_ICD9_CODES)} codes")
-
         # Test DKA codes
-        self.assertEqual(DKAPredictionMIMIC4.DKA_ICD10_PREFIX, "E101")
+        self.assertIn("E101", DKAPredictionMIMIC4.DKA_ICD10_PREFIXES)
         self.assertIn("25011", DKAPredictionMIMIC4.DKA_ICD9_CODES)
-        print(f"✓ DKA ICD-10 prefix: {DKAPredictionMIMIC4.DKA_ICD10_PREFIX}")
+        print(f"✓ DKA ICD-10 prefix: {DKAPredictionMIMIC4.DKA_ICD10_PREFIXES}")
         print(f"✓ DKA ICD-9 codes: {len(DKAPredictionMIMIC4.DKA_ICD9_CODES)} codes")
 
         # Test lab categories
-        self.assertEqual(len(DKAPredictionMIMIC4.LAB_CATEGORY_ORDER), 6)
-        expected_categories = ["glucose", "bicarbonate", "anion_gap", "potassium", "sodium", "chloride"]
+        self.assertEqual(len(DKAPredictionMIMIC4.LAB_CATEGORY_ORDER), 10)
+        expected_categories = ["Sodium", "Potassium", "Chloride", "Bicarbonate", "Glucose", "Calcium", "Magnesium", "Anion Gap", "Osmolality", "Phosphate"]
         self.assertEqual(DKAPredictionMIMIC4.LAB_CATEGORY_ORDER, expected_categories)
         print(f"✓ Lab categories: {DKAPredictionMIMIC4.LAB_CATEGORY_ORDER}")
-        print(f"✓ Total lab item IDs: {len(DKAPredictionMIMIC4.ALL_LAB_ITEMIDS)}")
+        print(f"✓ Total lab item IDs: {len(DKAPredictionMIMIC4.LABITEMS)}")
+
+    def test_t1ddka_prediction_class_variables(self):
+        """Test that class variables are properly defined."""
+        print(f"\n{'='*60}")
+        print("TEST: test_t1ddka_prediction_class_variables()")
+        print(f"{'='*60}")
+
+        # Test T1DM codes
+        self.assertEqual(T1DDKAPredictionMIMIC4.T1DM_ICD10_PREFIX, "E10")
+        self.assertIn("25001", T1DDKAPredictionMIMIC4.T1DM_ICD9_CODES)
+        print(f"✓ T1DM ICD-10 prefix: {T1DDKAPredictionMIMIC4.T1DM_ICD10_PREFIX}")
+        print(f"✓ T1DM ICD-9 codes: {len(T1DDKAPredictionMIMIC4.T1DM_ICD9_CODES)} codes")
+
+        # Test lab categories
+        self.assertEqual(len(T1DDKAPredictionMIMIC4.LAB_CATEGORY_ORDER), 10)
+        expected_categories = ["Sodium", "Potassium", "Chloride", "Bicarbonate", "Glucose", "Calcium", "Magnesium", "Anion Gap", "Osmolality", "Phosphate"]
+        self.assertEqual(T1DDKAPredictionMIMIC4.LAB_CATEGORY_ORDER, expected_categories)
+        print(f"✓ Lab categories: {T1DDKAPredictionMIMIC4.LAB_CATEGORY_ORDER}")
+        print(f"✓ Total lab item IDs: {len(T1DDKAPredictionMIMIC4.LABITEMS)}")
 
     def test_dka_prediction_set_task(self):
         """Test DKAPredictionMIMIC4 task with set_task() method."""
@@ -131,18 +162,15 @@ def test_dka_prediction_set_task(self):
             print("\nCalling dataset.set_task()...")
             sample_dataset = self.dataset.set_task(task)
             self.assertIsNotNone(sample_dataset, "set_task should return a dataset")
-            self.assertTrue(
-                hasattr(sample_dataset, "samples"), "Sample dataset should have samples"
-            )
             print(f"✓ set_task() completed")
 
             # Check sample count
-            num_samples = len(sample_dataset.samples)
+            num_samples = len(sample_dataset)
             print(f"✓ Generated {num_samples} DKA prediction samples")
 
             if num_samples > 0:
-                sample = sample_dataset.samples[0]
-                required_keys = ["patient_id", "record_id", "diagnoses", "labs", "label"]
+                sample = sample_dataset[0]
+                required_keys = ["patient_id", "record_id", "icd_codes", "labs", "label"]
 
                 print(f"\nFirst sample structure:")
                 print(f"  Sample keys: {list(sample.keys())}")
@@ -151,7 +179,7 @@ def test_dka_prediction_set_task(self):
                     self.assertIn(key, sample, f"Sample should contain key: {key}")
 
                 # Verify diagnoses format (tuple of times and sequences)
-                diagnoses = sample["diagnoses"]
+                diagnoses = sample["icd_codes"]
                 self.assertIsInstance(diagnoses, tuple, "diagnoses should be a tuple")
                 self.assertEqual(len(diagnoses), 2, "diagnoses tuple should have 2 elements")
                 print(f"  - diagnoses: {len(diagnoses[1])} admission(s)")
@@ -168,8 +196,8 @@ def test_dka_prediction_set_task(self):
 
                 # Count label distribution
                 label_counts = {0: 0, 1: 0}
-                for s in sample_dataset.samples:
-                    label_counts[s["label"]] += 1
+                for s in sample_dataset:
+                    label_counts[int(s["label"].item())] += 1
 
                 print(f"\nLabel distribution:")
                 print(f"  - No DKA (0): {label_counts[0]} ({label_counts[0]/num_samples*100:.1f}%)")
@@ -183,6 +211,69 @@ def test_dka_prediction_set_task(self):
             traceback.print_exc()
             self.fail(f"Failed to use set_task with DKAPredictionMIMIC4: {e}")
 
+    def test_t1ddka_prediction_set_task(self):
+        """Test T1DDKAPredictionMIMIC4 task with set_task() method."""
+        print(f"\n{'='*60}")
+        print("TEST: test_t1ddka_prediction_set_task()")
+        print(f"{'='*60}")
+
+        print("Initializing T1DDKAPredictionMIMIC4 task...")
+        task = T1DDKAPredictionMIMIC4()
+
+        # Test using set_task method
+        try:
+            print("\nCalling dataset.set_task()...")
+            sample_dataset = self.dataset.set_task(task)
+            self.assertIsNotNone(sample_dataset, "set_task should return a dataset")
+            print(f"✓ set_task() completed")
+
+            # Check sample count
+            num_samples = len(sample_dataset)
+            print(f"✓ Generated {num_samples} DKA prediction samples")
+
+            if num_samples > 0:
+                sample = sample_dataset[0]
+                required_keys = ["patient_id", "record_id", "icd_codes", "labs", "label"]
+
+                print(f"\nFirst sample structure:")
+                print(f"  Sample keys: {list(sample.keys())}")
+
+                for key in required_keys:
+                    self.assertIn(key, sample, f"Sample should contain key: {key}")
+
+                # Verify diagnoses format (tuple of times and sequences)
+                diagnoses = sample["icd_codes"]
+                self.assertIsInstance(diagnoses, tuple, "diagnoses should be a tuple")
+                self.assertEqual(len(diagnoses), 2, "diagnoses tuple should have 2 elements")
+                print(f"  - diagnoses: {len(diagnoses[1])} admission(s)")
+
+                # Verify labs format (tuple of times and sequences)
+                labs = sample["labs"]
+                self.assertIsInstance(labs, tuple, "labs should be a tuple")
+                self.assertEqual(len(labs), 2, "labs tuple should have 2 elements")
+                print(f"  - labs: {len(labs[1])} lab vector(s)")
+
+                # Verify label is binary
+                self.assertIn(sample["label"], [0, 1], "Label should be binary (0 or 1)")
+                print(f"  - label: {sample['label']}")
+
+                # Count label distribution
+                label_counts = {0: 0, 1: 0}
+                for s in sample_dataset:
+                    label_counts[int(s["label"].item())] += 1
+
+                print(f"\nLabel distribution:")
+                print(f"  - No DKA (0): {label_counts[0]} ({label_counts[0]/num_samples*100:.1f}%)")
+                print(f"  - Has DKA (1): {label_counts[1]} ({label_counts[1]/num_samples*100:.1f}%)")
+
+            print(f"\n✓ test_t1ddka_prediction_set_task() passed successfully")
+
+        except Exception as e:
+            print(f"✗ Failed with error: {e}")
+            import traceback
+            traceback.print_exc()
+            self.fail(f"Failed to use set_task with T1DDKAPredictionMIMIC4: {e}")
+
     def test_dka_prediction_helper_methods(self):
         """Test helper methods of DKAPredictionMIMIC4."""
         print(f"\n{'='*60}")
@@ -191,13 +282,25 @@ def test_dka_prediction_helper_methods(self):
 
         task = DKAPredictionMIMIC4()
 
-        # Test _normalize_icd
-        print("Testing _normalize_icd()...")
-        self.assertEqual(task._normalize_icd("E10.10"), "E1010")
-        self.assertEqual(task._normalize_icd("e10.10"), "E1010")
-        self.assertEqual(task._normalize_icd(None), "")
-        self.assertEqual(task._normalize_icd("  25001  "), "25001")
-        print("✓ _normalize_icd() works correctly")
+        # Test _is_dka_code
+        print("\nTesting _is_dka_code()...")
+        self.assertTrue(task._is_dka_code("E10.10", 10))
+        self.assertTrue(task._is_dka_code("E1011", "10"))
+        self.assertTrue(task._is_dka_code("25011", 9))
+        self.assertTrue(task._is_dka_code("25013", "9"))
+        self.assertFalse(task._is_dka_code("E10.65", 10))  # Not DKA
+        self.assertFalse(task._is_dka_code(None, 10))
+        print("✓ _is_dka_code() works correctly")
+
+        print(f"\n✓ test_dka_prediction_helper_methods() passed successfully")
+
+    def test_t1ddka_prediction_helper_methods(self):
+        """Test helper methods of DKAPredictionMIMIC4."""
+        print(f"\n{'='*60}")
+        print("TEST: test_t1ddka_prediction_helper_methods()")
+        print(f"{'='*60}")
+
+        task = T1DDKAPredictionMIMIC4()
 
         # Test _is_t1dm_code
         print("\nTesting _is_t1dm_code()...")
@@ -220,19 +323,7 @@ def test_dka_prediction_helper_methods(self):
         self.assertFalse(task._is_dka_code(None, 10))
         print("✓ _is_dka_code() works correctly")
 
-        # Test _deduplicate_preserve_order
-        print("\nTesting _deduplicate_preserve_order()...")
-        self.assertEqual(
-            task._deduplicate_preserve_order(["A", "B", "A", "C", "B"]),
-            ["A", "B", "C"]
-        )
-        self.assertEqual(
-            task._deduplicate_preserve_order([]),
-            []
-        )
-        print("✓ _deduplicate_preserve_order() works correctly")
-
-        print(f"\n✓ test_dka_prediction_helper_methods() passed successfully")
+        print(f"\n✓ test_t1ddka_prediction_helper_methods() passed successfully")
 
 
 if __name__ == "__main__":

tests/core/test_mimic4_los.py

@@ -17,7 +17,7 @@ def setUp(self):
     def _setup_dataset_path(self):
         """Get path to local MIMIC-IV demo dataset in test resources."""
         # Get the path to the test-resources/core/mimic4demo directory
-        test_dir = Path(__file__).parent.parent
+        test_dir = Path(__file__).parent.parent.parent
         self.demo_dataset_path = str(
             test_dir / "test-resources" / "core" / "mimic4demo"
         )
@@ -26,12 +26,6 @@ def _setup_dataset_path(self):
         print(f"Setting up MIMIC-IV demo dataset for length of stay prediction")
         print(f"Dataset path: {self.demo_dataset_path}")
 
-        # Verify the dataset exists
-        if not os.path.exists(self.demo_dataset_path):
-            raise unittest.SkipTest(
-                f"MIMIC-IV demo dataset not found at {self.demo_dataset_path}"
-            )
-
         # List files in the hosp directory
         hosp_path = os.path.join(self.demo_dataset_path, "hosp")
         if os.path.exists(hosp_path):
@@ -47,9 +41,9 @@ def _load_dataset(self):
         """Load the dataset for testing."""
         tables = ["diagnoses_icd", "procedures_icd", "prescriptions"]
         print(f"Loading MIMIC4Dataset with tables: {tables}")
-        self.dataset = MIMIC4Dataset(root=self.demo_dataset_path, tables=tables)
+        self.dataset = MIMIC4Dataset(ehr_root=self.demo_dataset_path, ehr_tables=tables)
         print(f"✓ Dataset loaded successfully")
-        print(f"  Total patients: {len(self.dataset.patients)}")
+        print(f"  Total patients: {len(self.dataset.unique_patient_ids)}")
         print()
 
     def test_dataset_stats(self):
@@ -89,19 +83,16 @@ def test_length_of_stay_prediction_mimic4_set_task(self):
             print("\nCalling dataset.set_task()...")
             sample_dataset = self.dataset.set_task(task)
             self.assertIsNotNone(sample_dataset, "set_task should return a dataset")
-            self.assertTrue(
-                hasattr(sample_dataset, "samples"), "Sample dataset should have samples"
-            )
             print(f"✓ set_task() completed")
 
             # Verify we got some samples
-            num_samples = len(sample_dataset.samples)
+            num_samples = len(sample_dataset)
             self.assertGreater(num_samples, 0, "Should generate at least one sample")
             print(f"✓ Generated {num_samples} length of stay prediction samples")
 
             # Test sample structure
             if num_samples > 0:
-                sample = sample_dataset.samples[0]
+                sample = sample_dataset[0]
                 required_keys = [
                     "visit_id",
                     "patient_id",
@@ -128,20 +119,10 @@ def test_length_of_stay_prediction_mimic4_set_task(self):
                     "Length of stay category should be 0-9",
                 )
 
-                # Verify conditions include icd_version prefix (MIMIC-4 specific)
-                if sample["conditions"]:
-                    first_condition = sample["conditions"][0]
-                    self.assertIn(
-                        "_",
-                        first_condition,
-                        "MIMIC-4 conditions should include version prefix (e.g., '10_E1010')",
-                    )
-                    print(f"  Sample condition format: {first_condition}")
-
                 # Count LOS distribution
                 los_counts = {i: 0 for i in range(10)}
-                for s in sample_dataset.samples:
-                    los_counts[s["los"]] += 1
+                for s in sample_dataset:
+                    los_counts[int(s["los"].item())] += 1
 
                 print(f"\nLength of stay category distribution:")
                 category_labels = [