Have _find_roots return multiple roots -> allow inverse tests to return union of intervals. (#107)

* Have _find_roots return multiple roots -> allow inverse tests to return union of intervals.

* Changelog.

Author: mlondschien

CHANGELOG.rst

@@ -8,6 +8,11 @@ Changelog
 
 - New test :func:`~ivmodels.tests.j.j_test` of the overidentifying restrictions.
 
+- The tests :func:`~ivmodels.tests.lagrange_multiplier.inverse_lagrange_multiplier_test`
+  and
+  :func:`~ivmodels.tests.conditional_likelihood_ratio.inverse_conditional_likelihood_ratio_test`
+  now possibly return unions of intervals, instead of one conservative large interval.
+
 **Bug fixes:**
 
 - Fixed bug in :func:`~ivmodels.models.kclass.KClass.fit` when `C` is not `None` and

benchmarks/benchmarks/tests.py

@@ -27,7 +27,7 @@ class Tests:
     def setup(self, n, data):
 
         if data == "guggenberger12 (k=10)":
-            Z, X, y, C, W, beta = simulate_guggenberger12(
+            Z, X, y, C, W, _, beta = simulate_guggenberger12(
                 n=n, k=10, seed=0, return_beta=True
             )
         else:

ivmodels/tests/conditional_likelihood_ratio.py

@@ -533,7 +533,7 @@ def f(x):
             )
         )
 
-    boundaries = []
+    roots = []
     for left_upper, right_upper in cs_upper.boundaries:
         left_lower_, right_lower_ = None, None
         for left_lower, right_lower in cs_lower.boundaries:
@@ -550,24 +550,26 @@ def f(x):
             else:
                 continue
 
-        boundaries.append(
-            (
-                _find_roots(
-                    f,
-                    left_lower_,
-                    left_upper,
-                    tol=tol,
-                    max_value=max_value,
-                    max_eval=max_eval,
-                ),
-                _find_roots(
-                    f,
-                    right_lower_,
-                    right_upper,
-                    tol=tol,
-                    max_value=max_value,
-                    max_eval=max_eval,
-                ),
-            )
+        roots += _find_roots(
+            f,
+            left_lower_,
+            left_upper,
+            tol=tol,
+            max_value=max_value,
+            max_eval=max_eval,
+            max_depth=5,
+        )
+        roots += _find_roots(
+            f,
+            right_lower_,
+            right_upper,
+            tol=tol,
+            max_value=max_value,
+            max_eval=max_eval,
         )
+
+    roots = sorted(roots)
+
+    assert len(roots) % 2 == 0
+    boundaries = [(left, right) for left, right in zip(roots[::2], roots[1::2])]
     return ConfidenceSet(boundaries=boundaries)

ivmodels/tests/lagrange_multiplier.py

@@ -490,20 +490,24 @@ def inverse_lagrange_multiplier_test(
     else:
         liml = KClass(kappa="liml", fit_intercept=False).fit(W, y, Z=Z, C=D).coef_[-1]
 
-    left = _find_roots(
+    roots = _find_roots(
         lambda x: lm.lm(x) - critical_value,
         a=liml,
         b=-np.inf,
         tol=tol,
         max_value=max_value,
         max_eval=max_eval,
     )
-    right = _find_roots(
+    roots += _find_roots(
         lambda x: lm.lm(x) - critical_value,
         a=liml,
         b=np.inf,
         tol=tol,
         max_value=max_value,
         max_eval=max_eval,
     )
-    return ConfidenceSet(boundaries=[(left, right)])
+
+    roots = sorted(roots)
+    assert len(roots) % 2 == 0
+    boundaries = [(left, right) for left, right in zip(roots[::2], roots[1::2])]
+    return ConfidenceSet(boundaries=boundaries)

ivmodels/utils.py

@@ -215,13 +215,13 @@ def _check_inputs(Z, X, y, W=None, C=None, D=None, beta=None):
     return Z, X, y, W, C, D, beta
 
 
-def _find_roots(f, a, b, tol, max_value, max_eval, n_points=50):
+def _find_roots(f, a, b, tol, max_value, max_eval, n_points=50, max_depth=5):
     """
-    Find the root of function ``f`` between ``a`` and ``b`` closest to ``b``.
+    Find roots of function ``f`` between ``a`` and ``b``.
 
     Assumes ``f(a) < 0`` and ``f(b) > 0``. Finds root by building a grid between ``a``
-    and ``b`` with ``n_points``, evaluating ``f`` at each point, and finding the last
-    point where ``f`` is negative. If ``b`` is infinite, uses a logarithmic grid between
+    and ``b`` with ``n_points``, evaluating ``f`` at each point, and finding indices
+    where ``f`` switches sign. If ``b`` is infinite, uses a logarithmic grid between
     ``a`` and ``a + sign(b - a) * max_value``. The function is then called recursively
     on the new interval until the size of the interval is less than ``tol`` or the
     maximum number of evaluations ``max_eval`` of ``f`` is reached.
@@ -230,40 +230,57 @@ def _find_roots(f, a, b, tol, max_value, max_eval, n_points=50):
     closest to ``b``. Note that this is also not strictly ensured by this function.
     """
     if np.abs(b - a) < tol or max_eval < 0:
-        return b  # conservative
+        return [b]  # conservative, resulting in a larger interval
+
     if np.isinf(a):
-        return a
+        return [a]
+
+    roots = []
+
     sgn = np.sign(b - a)
     if np.isinf(b):
         grid = np.ones(n_points) * a
-        grid[1:] += sgn * np.logspace(0, np.log10(max_value), n_points - 1)
+        grid[1:] += sgn * np.logspace(tol, np.log10(max_value), n_points - 1)
     else:
         grid = np.linspace(a, b, n_points)
 
     y = np.zeros(n_points)
-    y[-1] = f(grid[-1])
-    if y[-1] < 0:
-        return sgn * np.inf
 
     y[0] = f(grid[0])
     if y[0] >= 0:
         raise ValueError("f(a) must be negative.")
 
-    for i, x in enumerate(grid[:-1]):
-        y[i] = f(x)
+    for i, x in enumerate(grid[1:]):
+        y[i + 1] = f(x)
 
-    last_positive = np.where(y < 0)[0][-1]
+    if y[-1] <= 0:
+        roots = [b]
 
-    # f(a_new) < 0 < f(b_new) -> repeat
-    return _find_roots(
-        f,
-        grid[last_positive],
-        grid[last_positive + 1],
-        tol=tol,
-        n_points=n_points,
-        max_value=None,
-        max_eval=max_eval - n_points,
-    )
+    y[y == 0] = np.finfo(y.dtype).eps
+    where = np.where(np.sign(y[:-1]) != np.sign(y[1:]))[0]
+
+    # Conservative. Focus on change closest to b.
+    if max_depth == 0:
+        where = where[-1:]
+
+    for idx, w in enumerate(where):
+        if idx % 2 == 0:
+            a, b = grid[w], grid[w + 1]
+        else:
+            a, b = grid[w + 1], grid[w]
+
+        roots += _find_roots(
+            f,
+            a,
+            b,
+            tol=tol,
+            n_points=n_points,
+            max_value=max_value,
+            max_eval=max_eval - n_points,
+            max_depth=max_depth - len(where) > 1,
+        )
+
+    return roots
 
 
 def _characteristic_roots(a, b, subset_by_index=None):

tests/test_utils.py

@@ -3,7 +3,7 @@
 import pytest
 import scipy
 
-from ivmodels.utils import _characteristic_roots, oproj, proj, to_numpy
+from ivmodels.utils import _characteristic_roots, _find_roots, oproj, proj, to_numpy
 
 
 def test_proj():
@@ -132,3 +132,20 @@ def test_characteristic_roots(dim, rank):
     assert np.allclose(
         _characteristic_roots(A, B, subset_by_index=[0, 0]), np.min(finite_roots)
     )
+
+
+@pytest.mark.parametrize(
+    "f, a, b, expected",
+    [
+        (np.sin, -1, 8, [0, np.pi, 2 * np.pi]),
+        (lambda x: -np.sin(x), 8, -1, [0, np.pi, 2 * np.pi]),
+        (lambda x: x**2 - 1, 0, 2, [1]),
+        (lambda x: x**2 - 1, 0, -np.inf, [-1]),
+        (lambda x: x**3 - x, -2, 2, [-1, 0, 1]),
+        (lambda x: x**3 - x, 0.5, -np.inf, [-np.inf, -1, 0]),
+    ],
+)
+@pytest.mark.parametrize("tol", [1e-3, 1e-6])
+def test_find_roots(f, a, b, expected, tol):
+    roots = _find_roots(f, a, b, max_value=1e6, max_eval=1e4, tol=tol)
+    assert np.allclose(sorted(roots), expected, atol=tol)