60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from export_death_burden import (
|
||
|
|
death_probability,
|
||
|
|
find_death_column,
|
||
|
|
validate_horizons,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class DeathBurdenTests(unittest.TestCase):
|
||
|
|
def test_death_column_is_selected_by_code_and_outcome_type(self) -> None:
|
||
|
|
source = {
|
||
|
|
"tokens/column": np.asarray([0, 1, 2], dtype=np.int64),
|
||
|
|
"tokens/token_id": np.asarray([3, 4, 5], dtype=np.int64),
|
||
|
|
"tokens/label_code": np.asarray([b"A00", b"I10", b"Death"]),
|
||
|
|
"tokens/label_text": np.asarray(
|
||
|
|
[b"A00 cholera", b"I10 hypertension", b"Death"]
|
||
|
|
),
|
||
|
|
"tokens/outcome_type": np.asarray(
|
||
|
|
[b"disease", b"disease", b"death"]
|
||
|
|
),
|
||
|
|
}
|
||
|
|
|
||
|
|
self.assertEqual(find_death_column(source), 2)
|
||
|
|
|
||
|
|
def test_death_probability_combines_shape_and_scale(self) -> None:
|
||
|
|
result = death_probability(
|
||
|
|
shape=np.asarray([2.0, 1.0], dtype=np.float32),
|
||
|
|
scale=np.asarray([10.0, 4.0], dtype=np.float32),
|
||
|
|
horizons=np.asarray([5.0, 10.0], dtype=np.float64),
|
||
|
|
)
|
||
|
|
expected = np.asarray(
|
||
|
|
[
|
||
|
|
[1.0 - np.exp(-0.25), 1.0 - np.exp(-1.0)],
|
||
|
|
[1.0 - np.exp(-1.25), 1.0 - np.exp(-2.5)],
|
||
|
|
],
|
||
|
|
dtype=np.float32,
|
||
|
|
)
|
||
|
|
np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-7)
|
||
|
|
|
||
|
|
def test_invalid_parameters_produce_nan(self) -> None:
|
||
|
|
result = death_probability(
|
||
|
|
shape=np.asarray([1.0, -1.0], dtype=np.float32),
|
||
|
|
scale=np.asarray([np.nan, 2.0], dtype=np.float32),
|
||
|
|
horizons=np.asarray([5.0], dtype=np.float64),
|
||
|
|
)
|
||
|
|
self.assertTrue(np.isnan(result).all())
|
||
|
|
|
||
|
|
def test_horizons_must_be_positive_and_unique(self) -> None:
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
validate_horizons([0.0, 5.0])
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
validate_horizons([5.0, 5.0])
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|