Add conditional death burden export

tests/test_export_death_burden.py

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