Add length method to confidence_set.

mlondschien · Aug 22, 2024 · 5f1ffa8 · 5f1ffa8
1 parent bfd62ba
commit 5f1ffa8
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/ivmodels/confidence_set.py b/ivmodels/confidence_set.py
@@ -64,3 +64,7 @@ def from_quadric(quadric):
             return ConfidenceSet([(-np.inf, np.inf)])
         else:
             return ConfidenceSet([(-np.inf, boundary[0]), (boundary[1], np.inf)])
+
+    def length(self):
+        """Return the length of the confidence set."""
+        return sum(right - left for left, right in self.boundaries)
diff --git a/tests/test_confidence_set.py b/tests/test_confidence_set.py
@@ -5,17 +5,18 @@
 
 
 @pytest.mark.parametrize(
-    "boundaries, values, finite, empty",
+    "boundaries, values, finite, empty, length",
     [
-        ([], {0: 1}, True, True),
-        ([(-1, 1)], {-2: 1, 1: -1, 0: -1}, True, False),
-        ([(-np.inf, np.inf)], {-7: -1}, False, False),
-        ([(-np.inf, 0), (1, 20)], {0.5: 1, 12: -1}, False, False),
+        ([], {0: 1}, True, True, 0),
+        ([(-1, 1)], {-2: 1, 1: -1, 0: -1}, True, False, 2),
+        ([(-np.inf, np.inf)], {-7: -1}, False, False, np.inf),
+        ([(-np.inf, 0), (1, 20)], {0.5: 1, 12: -1}, False, False, np.inf),
     ],
 )
-def test_confidence_set_call(boundaries, values, finite, empty):
+def test_confidence_set_call(boundaries, values, finite, empty, length):
     confidence_set = ConfidenceSet(boundaries)
     for value, expected in values.items():
         assert confidence_set(value) == expected
     assert confidence_set.is_finite() == finite
     assert confidence_set.is_empty() == empty
+    assert confidence_set.length() == length