diff --git a/ivmodels/confidence_set.py b/ivmodels/confidence_set.py index b029c77..1c7a8a5 100644 --- 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 index e806bde..7b1db15 100644 --- 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