174 lines
6.1 KiB
Python
174 lines
6.1 KiB
Python
import unittest
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch.utils.data import Subset
|
|
|
|
from dataset import AllFutureHealthDataset
|
|
from losses import ExponentialLoss, WeibullLoss
|
|
from models import DeepHealth
|
|
from train_util import fit_all_future_baseline
|
|
|
|
|
|
def _inverse_softplus(value: torch.Tensor) -> torch.Tensor:
|
|
return torch.log(torch.expm1(value))
|
|
|
|
|
|
class AllFutureLikelihoodTests(unittest.TestCase):
|
|
def test_exponential_uses_event_specific_first_onset_exposure(self):
|
|
desired_rate = torch.tensor(
|
|
[[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]],
|
|
dtype=torch.float64,
|
|
)
|
|
logits = _inverse_softplus(desired_rate)
|
|
criterion = ExponentialLoss(ignored_idx={0, 1, 2}, eps=1e-12)
|
|
|
|
loss = criterion(
|
|
logits=logits,
|
|
targets=torch.tensor([[4, 0]]),
|
|
exposure=torch.tensor([5.0], dtype=torch.float64),
|
|
dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64),
|
|
history=torch.tensor([[3, 0]]),
|
|
)
|
|
|
|
rate = F.softplus(logits) + criterion.eps
|
|
expected = -(rate[0, 4].log()) + rate[0, 4] * 2.0 + rate[0, 5] * 5.0
|
|
torch.testing.assert_close(loss, expected)
|
|
|
|
def test_weibull_uses_event_time_and_excludes_prevalent_outcome(self):
|
|
desired_rate = torch.tensor(
|
|
[[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]],
|
|
dtype=torch.float64,
|
|
)
|
|
logits = _inverse_softplus(desired_rate)
|
|
rho = torch.tensor(
|
|
[[1.0, 1.0, 1.0, 1.2, 1.5, 0.8]],
|
|
dtype=torch.float64,
|
|
)
|
|
criterion = WeibullLoss(ignored_idx={0, 1, 2}, eps=1e-12)
|
|
|
|
loss = criterion(
|
|
logits=logits,
|
|
weibull_rho=rho,
|
|
targets=torch.tensor([[4, 0]]),
|
|
dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64),
|
|
exposure=torch.tensor([5.0], dtype=torch.float64),
|
|
history=torch.tensor([[3, 0]]),
|
|
)
|
|
|
|
rate = F.softplus(logits) + criterion.eps
|
|
log_hazard = (
|
|
rate[0, 4].log()
|
|
+ rho[0, 4].log()
|
|
+ (rho[0, 4] - 1.0) * torch.tensor(2.0).log()
|
|
)
|
|
expected = (
|
|
-log_hazard
|
|
+ rate[0, 4] * torch.pow(torch.tensor(2.0), rho[0, 4])
|
|
+ rate[0, 5] * torch.pow(torch.tensor(5.0), rho[0, 5])
|
|
)
|
|
torch.testing.assert_close(loss, expected)
|
|
|
|
|
|
class AllFutureQueryDistributionTests(unittest.TestCase):
|
|
@staticmethod
|
|
def _dataset() -> AllFutureHealthDataset:
|
|
dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
|
|
dataset.min_history_events = 1
|
|
dataset.min_future_events = 1
|
|
return dataset
|
|
|
|
@staticmethod
|
|
def _patient():
|
|
return {
|
|
"times": np.asarray([1.0, 2.0, 3.0, 4.0], dtype=np.float32),
|
|
"labels": np.asarray([3, 4, 5, 6], dtype=np.int64),
|
|
"t_obs": 4.0,
|
|
}
|
|
|
|
def test_train_validation_and_test_share_one_query_sampler(self):
|
|
dataset = self._dataset()
|
|
patient = self._patient()
|
|
patient["query_intervals"] = dataset._eligible_query_intervals(patient)
|
|
self.assertEqual(len(patient["query_intervals"]), 3)
|
|
|
|
rng = np.random.RandomState(123)
|
|
fixed = dataset._sample_fixed_validation_queries(patient, rng)
|
|
self.assertEqual(len(fixed), 1)
|
|
self.assertTrue(dataset._is_valid_query(patient, fixed[0]))
|
|
|
|
seen_intervals = set()
|
|
rng = np.random.RandomState(456)
|
|
for _ in range(200):
|
|
query = dataset.sample_query(patient, rng)
|
|
self.assertTrue(dataset._is_valid_query(patient, query))
|
|
seen_intervals.add(int(np.floor(query)))
|
|
self.assertEqual(seen_intervals, {1, 2, 3})
|
|
|
|
|
|
class AllFutureBaselineTests(unittest.TestCase):
|
|
def test_training_baseline_matches_first_onset_exposure(self):
|
|
dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset)
|
|
dataset.vocab_size = 7
|
|
dataset.min_history_events = 1
|
|
dataset.min_future_events = 2
|
|
patient = {
|
|
"times": np.asarray([1.0, 3.0, 5.0], dtype=np.float32),
|
|
"labels": np.asarray([3, 4, 5], dtype=np.int64),
|
|
"t_obs": 5.0,
|
|
"query_intervals": [(1.0, 3.0)],
|
|
}
|
|
dataset.patients = [patient]
|
|
subset = Subset(dataset, [0])
|
|
|
|
stats = fit_all_future_baseline(dataset, subset, seed=17)
|
|
query = dataset.sample_query(patient, np.random.RandomState(17))
|
|
|
|
self.assertEqual(stats.event_count[3], 0)
|
|
self.assertEqual(stats.at_risk_exposure[3], 0.0)
|
|
self.assertEqual(stats.event_count[4], 1)
|
|
self.assertEqual(stats.event_count[5], 1)
|
|
self.assertAlmostEqual(stats.at_risk_exposure[4], 3.0 - query, places=5)
|
|
self.assertAlmostEqual(stats.at_risk_exposure[5], 5.0 - query, places=5)
|
|
self.assertAlmostEqual(
|
|
float(F.softplus(torch.tensor(stats.bias[4]))),
|
|
float(stats.rate[4]),
|
|
places=6,
|
|
)
|
|
|
|
def test_model_starts_exactly_at_fitted_output_baseline(self):
|
|
baseline_rate = torch.tensor([0.0, 0.0, 0.0, 0.02, 0.05, 0.10])
|
|
baseline_bias = torch.zeros_like(baseline_rate)
|
|
baseline_bias[3:] = _inverse_softplus(baseline_rate[3:])
|
|
model = DeepHealth(
|
|
vocab_size=6,
|
|
n_embd=4,
|
|
n_head=1,
|
|
n_layer=1,
|
|
n_types=1,
|
|
n_cont_types=0,
|
|
n_categories=1,
|
|
cont_type_ids=[],
|
|
target_mode="all_future",
|
|
time_mode="absolute",
|
|
dist_mode="weibull",
|
|
model_architecture="transformer_ffn_v1",
|
|
risk_head_bias=True,
|
|
risk_head_bias_init=baseline_bias,
|
|
)
|
|
|
|
torch.testing.assert_close(model.risk_head.weight, torch.zeros_like(model.risk_head.weight))
|
|
output_rate = F.softplus(model.risk_head(torch.randn(3, 4)))
|
|
torch.testing.assert_close(output_rate[:, 3:], baseline_rate[None, 3:].expand(3, -1))
|
|
torch.testing.assert_close(
|
|
F.softplus(model.rho_head.bias),
|
|
torch.ones_like(model.rho_head.bias),
|
|
atol=2e-5,
|
|
rtol=0.0,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|