JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove legacy event and mixed distribution paths

Browse Source
This commit is contained in:
Jiarui Li
2026-08-01 14:23:18 +08:00
parent dfb22adf2d
commit de6f9b75b9
22 changed files with 370 additions and 463 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
tests/test_calibration_metrics.py
View File

@@ -207,19 +207,15 @@ class IPCWCalibrationMetricTests(unittest.TestCase):
rho = np.asarray([0.8, 1.0, 1.2, 1.5], dtype=np.float32)
horizons = np.asarray([0.1, 1.0, 5.0], dtype=np.float32)
for dist_mode, token, death_idx, selected_rho in (
("exponential", 4, 9, None),
("weibull", 4, 9, rho),
("mixed", 9, 9, rho),
("mixed", 4, 9, None),
for dist_mode, selected_rho in (
("exponential", None),
("weibull", rho),
):
actual = _risk_probability_matrix(
logits=logits,
rho=selected_rho,
horizons=horizons,
dist_mode=dist_mode,
token=token,
death_idx=death_idx,
)
expected = np.vstack(
[
@@ -229,8 +225,6 @@ class IPCWCalibrationMetricTests(unittest.TestCase):
score_mode="risk",
horizon=float(horizon),
dist_mode=dist_mode,
token=token,
death_idx=death_idx,
)
for horizon in horizons
]
@@ -271,7 +265,6 @@ class IPCWCalibrationMetricTests(unittest.TestCase):
"label_id_to_code": {4: "D4", 5: "D5"},
"dist_mode": "exponential",
"horizons": np.asarray([1.0, 5.0], dtype=np.float32),
"death_index": 9,
"min_cases": 1,
"min_controls": 1,
"max_ipcw_weight": 0.0,

115
tests/test_continuous_value_scaling.py
View File

@@ -5,7 +5,8 @@ import torch
import torch.nn as nn
from torch.utils.data import Subset
from models import OtherInfoTokenizer
from eval_data import build_model_from_dataset
from models import DeepHealth, OtherInfoTokenizer
from train_util import fit_continuous_robust_scaler
@@ -59,6 +60,17 @@ class ContinuousValueScalingTests(unittest.TestCase):
np.testing.assert_allclose(stats.center, np.asarray([2.0, 10.0]))
np.testing.assert_allclose(stats.scale, np.asarray([2.0, 1.0]))
def test_fit_supports_next_step_sample_storage(self):
dataset = _ToyAllFutureDataset()
dataset.samples = dataset.patients
del dataset.patients
train_subset = Subset(dataset, np.asarray([0, 1, 2, 3, 4]))
stats = fit_continuous_robust_scaler(dataset, train_subset)
np.testing.assert_allclose(stats.center, np.asarray([2.0, 10.0]))
np.testing.assert_allclose(stats.scale, np.asarray([2.0, 1.0]))
def test_tokenizer_standardizes_only_continuous_values(self):
tokenizer = OtherInfoTokenizer(
n_embd=4,
@@ -66,7 +78,6 @@ class ContinuousValueScalingTests(unittest.TestCase):
n_cont_types=2,
n_categories=3,
cont_type_ids=[1, 3],
continuous_value_scaling="robust",
continuous_value_center=[10.0, 100.0],
continuous_value_scale=[2.0, 20.0],
)
@@ -89,7 +100,6 @@ class ContinuousValueScalingTests(unittest.TestCase):
n_cont_types=2,
n_categories=2,
cont_type_ids=[1, 3],
continuous_value_scaling="robust",
continuous_value_center=[2.0, 10.0],
continuous_value_scale=[1.5, 4.0],
)
@@ -103,7 +113,8 @@ class ContinuousValueScalingTests(unittest.TestCase):
n_cont_types=2,
n_categories=2,
cont_type_ids=[1, 3],
continuous_value_scaling="robust",
continuous_value_center=[0.0, 0.0],
continuous_value_scale=[1.0, 1.0],
)
restored.load_state_dict(state, strict=True)
@@ -116,27 +127,99 @@ class ContinuousValueScalingTests(unittest.TestCase):
torch.tensor([1.5, 4.0]),
)
def test_legacy_none_mode_keeps_old_state_dict_schema(self):
tokenizer = OtherInfoTokenizer(
n_embd=4,
n_types=4,
n_cont_types=2,
n_categories=2,
cont_type_ids=[1, 3],
)
state = tokenizer.state_dict()
def test_continuous_tokenizer_rejects_missing_scaler_statistics(self):
with self.assertRaisesRegex(ValueError, "require train-split RobustScale"):
OtherInfoTokenizer(
n_embd=4,
n_types=4,
n_cont_types=2,
n_categories=2,
cont_type_ids=[1, 3],
)
self.assertNotIn("continuous_value_center", state)
self.assertNotIn("continuous_value_scale", state)
restored = OtherInfoTokenizer(
def test_evaluation_rejects_unscaled_continuous_checkpoint(self):
dataset = type(
"DatasetMetadata",
(),
{
"vocab_size": 8,
"n_types": 4,
"n_cont_types": 2,
"n_categories": 2,
"cont_type_ids": [1, 3],
},
)()
cfg = {
"model_target_mode": "all_future",
"target_mode": "all_future",
"model_architecture": "transformer_ffn_v1",
"n_layer": 1,
"time_mode": "absolute",
"dist_mode": "exponential",
}
with self.assertRaisesRegex(RuntimeError, "unscaled checkpoints are not supported"):
build_model_from_dataset(
None,
cfg,
dataset,
state_dict={"blocks.0.mlp.w1.weight": torch.zeros(1)},
)
def test_evaluation_restores_required_scaler_buffers(self):
dataset = type(
"DatasetMetadata",
(),
{
"vocab_size": 8,
"n_types": 4,
"n_cont_types": 2,
"n_categories": 2,
"cont_type_ids": [1, 3],
},
)()
source = DeepHealth(
vocab_size=8,
n_embd=4,
n_head=1,
n_layer=1,
n_types=4,
n_cont_types=2,
n_categories=2,
cont_type_ids=[1, 3],
continuous_value_center=[2.0, 10.0],
continuous_value_scale=[1.5, 4.0],
target_mode="all_future",
time_mode="absolute",
dist_mode="exponential",
model_architecture="transformer_ffn_v1",
)
state = source.state_dict()
cfg = {
"model_target_mode": "all_future",
"target_mode": "all_future",
"model_architecture": "transformer_ffn_v1",
"n_embd": 4,
"n_head": 1,
"n_layer": 1,
"n_bins": 16,
"time_mode": "absolute",
"dist_mode": "exponential",
"continuous_value_scaling": "robust",
}
restored = build_model_from_dataset(None, cfg, dataset, state_dict=state)
restored.load_state_dict(state, strict=True)
torch.testing.assert_close(
restored.tokenizer.continuous_value_center,
torch.tensor([2.0, 10.0]),
)
torch.testing.assert_close(
restored.tokenizer.continuous_value_scale,
torch.tensor([1.5, 4.0]),
)
if __name__ == "__main__":
unittest.main()

57
tests/test_dataset_checkup.py
View File

@@ -1,57 +0,0 @@
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()

44
tests/test_dataset_reserved_event.py Normal file
View File

@@ -0,0 +1,44 @@
import unittest
import numpy as np
from dataset import _ExpoBaseDataset
from targets import RESERVED_IDX
class ReservedEventFilteringTests(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, RESERVED_IDX],
[101, 20, 2],
[101, 30, 3],
],
dtype=np.float64,
)
return dataset
def _assert_reserved_event_removed(self, extra_info_types):
dataset = self._base(extra_info_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))
np.testing.assert_array_equal(labels, np.asarray([3, 4], dtype=np.int64))
self.assertNotIn(RESERVED_IDX, labels.tolist())
def test_empty_extra_info_removes_legacy_reserved_event(self):
self._assert_reserved_event_removed([])
def test_selected_extra_info_removes_legacy_reserved_event(self):
self._assert_reserved_event_removed([11])
if __name__ == "__main__":
unittest.main()