Add Newton-bisection hybrid root-finding method

docs/api/root_find.md

@@ -41,6 +41,13 @@
 
 ---
 
+::: optimistix.NewtonBisection
+    options:
+        members:
+            - __init__
+
+---
+
 ::: optimistix.BestSoFarRootFinder
     options:
         members:

docs/how-to-choose.md

@@ -26,7 +26,7 @@ If your problem is particularly messy, as discussed above, then use a first-orde
 
 ## Root-finding problems
 
-For one-dimensional problems, use [`optimistix.Bisection`][].
+For one-dimensional problems, use [`optimistix.Bisection`][] or [`optimistix.NewtonBisection`][].
 
 For relatively "well-behaved" problems, then either [`optimistix.Newton`][] or [`optimistix.Chord`][] are recommended.
 

optimistix/__init__.py

@@ -71,6 +71,7 @@
     LinearTrustRegion as LinearTrustRegion,
     NelderMead as NelderMead,
     Newton as Newton,
+    NewtonBisection as NewtonBisection,
     NewtonDescent as NewtonDescent,
     NonlinearCG as NonlinearCG,
     NonlinearCGDescent as NonlinearCGDescent,

optimistix/_solver/__init__.py

@@ -5,7 +5,10 @@
     BestSoFarMinimiser as BestSoFarMinimiser,
     BestSoFarRootFinder as BestSoFarRootFinder,
 )
-from .bisection import Bisection as Bisection
+from .bisection import (
+    Bisection as Bisection,
+    NewtonBisection as NewtonBisection,
+)
 from .dogleg import Dogleg as Dogleg, DoglegDescent as DoglegDescent
 from .fixed_point import FixedPointIteration as FixedPointIteration
 from .gauss_newton import (

optimistix/_solver/bisection.py

@@ -94,8 +94,8 @@ class Bisection(AbstractRootFinder[Scalar, Scalar, Aux, _BisectionState]):
     the interval is sufficiently small.
 
     If `expand_if_necessary` and `detect` are true, the initial interval will be
-    expanded if it doesn't contain the the root.  This expansion assumes that the
-    function is monotonic.
+    expanded if it doesn't contain the root.  This expansion assumes that the function
+    is monotonic.
     """
 
     rtol: float
@@ -226,3 +226,164 @@ def postprocess(
     `[lower, upper]`. To locate the root correctly then this assumes that the function
     is monotonic.
 """
+
+
+class NewtonBisection(AbstractRootFinder[Scalar, Scalar, Aux, _BisectionState]):
+    """Newton-bisection hybrid. This may only be used with functions
+    `R->R`, i.e. functions with scalar input and scalar output.
+
+    This requires the following `options`:
+
+    - `lower`: The lower bound on the interval which contains the root.
+    - `upper`: The upper bound on the interval which contains the root.
+
+    Which are passed as, for example,
+    `optimistix.root_find(..., options=dict(lower=0, upper=1))`
+
+    This algorithm works by considering the interval `[lower, upper]`, attempting a
+    Newton step first, and falling back to bisection if the Newton step would leave the
+    interval. This is then repeated. The iteration stops once the interval is
+    sufficiently small.
+
+    If `expand_if_necessary` and `detect` are true, the initial interval will be
+    expanded if it doesn't contain the root.  This expansion assumes that the function
+    is monotonic.
+    """
+
+    rtol: float
+    atol: float
+    flip: bool | Literal["detect"] = "detect"
+    expand_if_necessary: bool = False
+    # All norms are the same for scalars.
+    norm: ClassVar[Callable[[PyTree], Scalar]] = jnp.abs
+
+    def init(
+        self,
+        fn: Fn[Scalar, Scalar, Aux],
+        y: Scalar,
+        args: PyTree,
+        options: dict[str, Any],
+        f_struct: jax.ShapeDtypeStruct,
+        aux_struct: PyTree[jax.ShapeDtypeStruct],
+        tags: frozenset[object],
+    ) -> _BisectionState:
+        lower = jnp.asarray(options["lower"], f_struct.dtype)
+        upper = jnp.asarray(options["upper"], f_struct.dtype)
+        del options, aux_struct
+        if jnp.shape(y) != () or jnp.shape(lower) != () or jnp.shape(upper) != ():
+            raise ValueError(
+                "NewtonBisection can only be used to find the roots of a function "
+                "taking a scalar input."
+            )
+        if not isinstance(f_struct, jax.ShapeDtypeStruct) or f_struct.shape != ():
+            raise ValueError(
+                "NewtonBisection can only be used to find the roots of a function "
+                "producing a scalar output."
+            )
+        if isinstance(self.flip, bool):
+            # Make it possible to avoid the extra two function compilations.
+            flip = jnp.array(self.flip)
+        elif self.flip == "detect":
+            lower_val, _ = fn(lower, args)
+            upper_val, _ = fn(upper, args)
+            flip = lower_val > upper_val
+            if self.expand_if_necessary:
+                lower, upper = _expand_interval_repeatedly(
+                    lower,
+                    upper,
+                    upper_val=upper_val,
+                    lower_val=lower_val,
+                    need_positive=lower_val < 0.0,
+                    fn=fn,
+                    args=args,
+                )
+            else:
+                lower_neg = lower_val < 0
+                upper_neg = upper_val < 0
+                root_not_contained = lower_neg == upper_neg
+                flip = eqx.error_if(
+                    flip,
+                    root_not_contained,
+                    msg="The root is not contained in [lower, upper]",
+                )
+        else:
+            raise ValueError("`flip` may only be True, False, or 'detect'.")
+        error, _ = fn(y, args)
+        return _BisectionState(
+            lower=lower,
+            upper=upper,
+            flip=flip,
+            error=error,
+        )
+
+    def step(
+        self,
+        fn: Fn[Scalar, Scalar, Aux],
+        y: Scalar,
+        args: PyTree,
+        options: dict[str, Any],
+        state: _BisectionState,
+        tags: frozenset[object],
+    ) -> tuple[Scalar, _BisectionState, Aux]:
+        del options
+        error = state.error
+        deriv, _ = jax.grad(fn, has_aux=True)(y, args)
+        negative = state.flip ^ (error < 0)
+        new_lower = jnp.where(negative, y, state.lower)
+        new_upper = jnp.where(negative, state.upper, y)
+        new_y_newton = y - error / deriv
+        new_y_bisection = new_lower + 0.5 * (new_upper - new_lower)
+        new_y = jnp.where(
+            (new_y_newton >= new_lower) & (new_y_newton <= new_upper),
+            new_y_newton,
+            new_y_bisection,
+        )
+        new_error, aux = fn(new_y, args)
+        new_state = _BisectionState(
+            lower=new_lower, upper=new_upper, flip=state.flip, error=new_error
+        )
+        return new_y, new_state, aux
+
+    def terminate(
+        self,
+        fn: Fn[Scalar, Scalar, Aux],
+        y: Scalar,
+        args: PyTree,
+        options: dict[str, Any],
+        state: _BisectionState,
+        tags: frozenset[object],
+    ) -> tuple[Bool[Array, ""], RESULTS]:
+        del fn, args, options
+        scale = self.atol + self.rtol * jnp.abs(y)
+        y_small = jnp.abs(state.lower - state.upper) < scale
+        f_small = jnp.abs(state.error) < self.atol
+        return y_small & f_small, RESULTS.successful
+
+    def postprocess(
+        self,
+        fn: Fn[Scalar, Scalar, Aux],
+        y: Scalar,
+        aux: Aux,
+        args: PyTree,
+        options: dict[str, Any],
+        state: _BisectionState,
+        tags: frozenset[object],
+        result: RESULTS,
+    ) -> tuple[Scalar, Aux, dict[str, Any]]:
+        return y, aux, {}
+
+
+NewtonBisection.__init__.__doc__ = """**Arguments:**
+
+- `rtol`: Relative tolerance for terminating solve.
+- `atol`: Absolute tolerance for terminating solve.
+- `flip`: Can be set to any of:
+    - `False`: specify that `fn(lower, args) < 0 < fn(upper, args)`.
+    - `True`: specify that `fn(lower, args) > 0 > fn(upper, args)`.
+    - `"detect"`: automatically check `fn(lower, args)` and `fn(upper, args)`. Note that
+        this option may increase both runtime and compilation time.
+- `expand_if_necessary`: If `True` then the `lower` and `upper` passed as options will
+    be made larger or smaller if the root is not found within the interval
+    `[lower, upper]`. To locate the root correctly then this assumes that the function
+    is monotonic.
+"""

tests/helpers.py

@@ -779,6 +779,14 @@ def trivial(y: Array, args: PyTree):
     return y - 1
 
 
+def _quadratic(y: Array, args: PyTree):
+    """Quadratic function with zero derivative at y=1, roots at y=0 and y=2.
+
+    Used to test that NewtonBisection handles critical points properly.
+    """
+    return (y - 1) ** 2 - 1
+
+
 # ARGS FOR BISECTION AND FIXED POINT
 ones_pytree = ({"a": 0.5 * jnp.ones((3, 2, 4))}, 0.5 * jnp.ones(4))
 flat_ones, _ = jfu.ravel_pytree(ones_pytree)
@@ -835,6 +843,12 @@ def trivial(y: Array, args: PyTree):
         {"upper": jnp.array(1.0), "lower": jnp.array(0.0)},
         jnp.array([-0.137]),
     ),
+    (
+        _quadratic,
+        jnp.array(1.0),  # Purposely start where derivative is 0
+        {"upper": jnp.array(3.0), "lower": jnp.array(0.5)},
+        None,
+    ),
 )
 
 fixed_point_fn_init_args = (

tests/test_bisection.py

@@ -0,0 +1,98 @@
+import random
+
+import jax.numpy as jnp
+import jax.tree_util as jtu
+import optimistix as optx
+import pytest
+from equinox.internal import ω
+
+from .helpers import (
+    bisection_fn_init_options_args,
+    tree_allclose,
+    trivial,
+)
+
+
+atol = rtol = 1e-6
+_fp_solvers = (optx.FixedPointIteration(rtol, atol),)
+smoke_aux = (jnp.ones((2, 3)), {"smoke_aux": jnp.ones(2)})
+
+
+@pytest.mark.parametrize(
+    "_fn, init, bisection_options, args", bisection_fn_init_options_args
+)
+def test_bisection(_fn, init, bisection_options, args):
+    solver = optx.Bisection(rtol=1e-6, atol=1e-6)
+    atol = rtol = 1e-4
+    has_aux = random.choice([True, False])
+
+    def root_find_problem(y, args):
+        f_val = _fn(y, args)
+        return (f_val**ω - y**ω).ω
+
+    if _fn == trivial:
+        bisection_problem = _fn
+    else:
+        bisection_problem = root_find_problem
+    if has_aux:
+        fn = lambda x, args: (bisection_problem(x, args), smoke_aux)
+    else:
+        fn = bisection_problem
+
+    optx_fp = optx.root_find(
+        fn,
+        solver,
+        init,
+        has_aux=has_aux,
+        args=args,
+        options=bisection_options,
+        max_steps=10_000,
+        throw=False,
+    ).value
+    out = fn(optx_fp, args)
+    if has_aux:
+        f_val, _ = out
+    else:
+        f_val = out
+    zeros = jtu.tree_map(jnp.zeros_like, f_val)
+    assert tree_allclose(f_val, zeros, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize(
+    "_fn, init, bisection_options, args", bisection_fn_init_options_args
+)
+def test_newton_bisection(_fn, init, bisection_options, args):
+    solver = optx.NewtonBisection(rtol=1e-6, atol=1e-6)
+    atol = rtol = 1e-4
+    has_aux = random.choice([True, False])
+
+    def root_find_problem(y, args):
+        f_val = _fn(y, args)
+        return (f_val**ω - y**ω).ω
+
+    if _fn == trivial:
+        bisection_problem = _fn
+    else:
+        bisection_problem = root_find_problem
+    if has_aux:
+        fn = lambda x, args: (bisection_problem(x, args), smoke_aux)
+    else:
+        fn = bisection_problem
+
+    optx_fp = optx.root_find(
+        fn,
+        solver,
+        init,
+        has_aux=has_aux,
+        args=args,
+        options=bisection_options,
+        max_steps=10_000,
+        throw=False,
+    ).value
+    out = fn(optx_fp, args)
+    if has_aux:
+        f_val, _ = out
+    else:
+        f_val = out
+    zeros = jtu.tree_map(jnp.zeros_like, f_val)
+    assert tree_allclose(f_val, zeros, atol=atol, rtol=rtol)