DeepHealth/tests/test_continuous_value_scaling.py

import unittest

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Subset

from eval_data import build_model_from_dataset
from models import DeepHealth, OtherInfoTokenizer
from train_util import fit_continuous_robust_scaler


class _ToyAllFutureDataset:
    def __init__(self):
        self.cont_type_ids = [1, 3]
        self.n_types = 4
        self.patients = [
            self._patient([1, 3], [0.0, 10.0]),
            self._patient([1, 3], [1.0, 10.0]),
            self._patient([1, 3], [2.0, 10.0]),
            self._patient([1, 3], [3.0, 10.0]),
            self._patient([1, 3], [4.0, 10.0]),
            self._patient([1, 3], [1000.0, 999.0]),
        ]

    @staticmethod
    def _patient(types, values):
        return {
            "other_type": np.asarray(types, dtype=np.int64),
            "other_value": np.asarray(values, dtype=np.float32),
            "other_value_kind": np.ones(len(types), dtype=np.int64),
        }

    def __len__(self):
        return len(self.patients)


class _CaptureContinuousEncoder(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.n_embd = n_embd
        self.last_type = None
        self.last_value = None

    def forward(self, cont_type_idx, value):
        self.last_type = cont_type_idx.detach().clone()
        self.last_value = value.detach().clone()
        return value[:, None].expand(-1, self.n_embd)


class ContinuousValueScalingTests(unittest.TestCase):
    def test_fit_uses_only_training_subset_and_handles_constant_features(self):
        dataset = _ToyAllFutureDataset()
        train_subset = Subset(dataset, np.asarray([0, 1, 2, 3, 4]))

        stats = fit_continuous_robust_scaler(dataset, train_subset)

        self.assertEqual(stats.cont_type_ids, (1, 3))
        np.testing.assert_array_equal(stats.observation_count, np.asarray([5, 5]))
        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,
            n_types=4,
            n_cont_types=2,
            n_categories=3,
            cont_type_ids=[1, 3],
            continuous_value_center=[10.0, 100.0],
            continuous_value_scale=[2.0, 20.0],
        )
        capture = _CaptureContinuousEncoder(n_embd=4)
        tokenizer.cont_value_encoder = capture

        tokenizer(
            other_type=torch.tensor([[1, 2, 3]], dtype=torch.long),
            other_value=torch.tensor([[14.0, 1.0, 80.0]]),
            other_value_kind=torch.tensor([[1, 2, 1]], dtype=torch.long),
        )

        torch.testing.assert_close(capture.last_type, torch.tensor([0, 1]))
        torch.testing.assert_close(capture.last_value, torch.tensor([2.0, -1.0]))

    def test_scaler_buffers_round_trip_in_new_checkpoint(self):
        tokenizer = OtherInfoTokenizer(
            n_embd=4,
            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],
        )
        state = tokenizer.state_dict()
        self.assertIn("continuous_value_center", state)
        self.assertIn("continuous_value_scale", state)

        restored = OtherInfoTokenizer(
            n_embd=4,
            n_types=4,
            n_cont_types=2,
            n_categories=2,
            cont_type_ids=[1, 3],
            continuous_value_center=[0.0, 0.0],
            continuous_value_scale=[1.0, 1.0],
        )
        restored.load_state_dict(state, strict=True)

        torch.testing.assert_close(
            restored.continuous_value_center,
            torch.tensor([2.0, 10.0]),
        )
        torch.testing.assert_close(
            restored.continuous_value_scale,
            torch.tensor([1.5, 4.0]),
        )

    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],
            )

    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()