58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
|
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()
|