From 0bf5278e298a750485e08096e0a3de7fdb38eb02 Mon Sep 17 00:00:00 2001
From: Diane Napolitano <diane.napolitano@washpost.com>
Date: Tue, 30 Jan 2024 14:46:01 -0500
Subject: [PATCH] Adding unit tests for other values that could be specified to
 credible interval and also getting credible interval transitions

---
 tests/test_ei_transition_solver.py | 99 +++++++++++++++++++++++++++++-
 1 file changed, 98 insertions(+), 1 deletion(-)

diff --git a/tests/test_ei_transition_solver.py b/tests/test_ei_transition_solver.py
index 1ea0649a..8d354bcb 100644
--- a/tests/test_ei_transition_solver.py
+++ b/tests/test_ei_transition_solver.py
@@ -3,7 +3,8 @@
 
 from elexsolver.EITransitionSolver import EITransitionSolver
 
-# high tolerance to match PyMC's unit tests
+# high tolerance due to random sampling
+# (which can produce different outcomes on different architectures, despite setting seeds)
 RTOL = 1e-02
 ATOL = 1e-02
 
@@ -140,3 +141,99 @@ def test_ei_credible_interval_percentages():
     (current_lower, current_upper) = ei.get_credible_interval(99, transitions=False)
     np.testing.assert_allclose(expected_lower, current_lower, rtol=RTOL, atol=ATOL)
     np.testing.assert_allclose(expected_upper, current_upper, rtol=RTOL, atol=ATOL)
+
+
+def test_ei_credible_interval_percentages_float_interval():
+    X = np.array(
+        [
+            [1, 2],
+            [3, 4],
+            [5, 6],
+            [7, 8],
+            [9, 10],
+            [11, 12],
+        ]
+    )
+
+    Y = np.array(
+        [
+            [2, 3],
+            [4, 5],
+            [6, 7],
+            [8, 9],
+            [10, 11],
+            [12, 13],
+        ]
+    )
+
+    expected_lower = np.array([[0.037212, 0.356174], [0.424652, 0.117605]])
+    expected_upper = np.array([[0.643826, 0.962788], [0.882395, 0.575348]])
+
+    ei = EITransitionSolver(random_seed=1024, n_samples=10, sampling_chains=1)
+    _ = ei.fit_predict(X, Y)
+    (current_lower, current_upper) = ei.get_credible_interval(0.99, transitions=False)
+    np.testing.assert_allclose(expected_lower, current_lower, rtol=RTOL, atol=ATOL)
+    np.testing.assert_allclose(expected_upper, current_upper, rtol=RTOL, atol=ATOL)
+
+
+def test_ei_credible_interval_invalid():
+    X = np.array(
+        [
+            [1, 2],
+            [3, 4],
+            [5, 6],
+            [7, 8],
+            [9, 10],
+            [11, 12],
+        ]
+    )
+
+    Y = np.array(
+        [
+            [2, 3],
+            [4, 5],
+            [6, 7],
+            [8, 9],
+            [10, 11],
+            [12, 13],
+        ]
+    )
+
+    ei = EITransitionSolver(random_seed=1024, n_samples=10, sampling_chains=1)
+    _ = ei.fit_predict(X, Y)
+
+    with pytest.raises(ValueError):
+        ei.get_credible_interval(3467838976, transitions=False)
+
+
+def test_ei_credible_interval_transitions():
+    X = np.array(
+        [
+            [1, 2],
+            [3, 4],
+            [5, 6],
+            [7, 8],
+            [9, 10],
+            [11, 12],
+        ]
+    )
+
+    Y = np.array(
+        [
+            [2, 3],
+            [4, 5],
+            [6, 7],
+            [8, 9],
+            [10, 11],
+            [12, 13],
+        ]
+    )
+
+    expected_lower = np.array([[0.017175, 0.164388], [0.228659, 0.063326]])
+    expected_upper = np.array([[0.29715, 0.444364], [0.475136, 0.309803]])
+
+    ei = EITransitionSolver(random_seed=1024, n_samples=10, sampling_chains=1)
+    _ = ei.fit_predict(X, Y)
+    (current_lower, current_upper) = ei.get_credible_interval(99, transitions=True)
+    np.testing.assert_allclose(expected_lower, current_lower, rtol=RTOL, atol=ATOL)
+    np.testing.assert_allclose(expected_upper, current_upper, rtol=RTOL, atol=ATOL)