Add WolfeSearch: a strong Wolfe line search

docs/api/searches/searches.md

@@ -21,6 +21,13 @@
 
 ---
 
+::: optimistix.WolfeSearch
+    options:
+        members:
+            - __init__
+
+---
+
 ::: optimistix.ClassicalTrustRegion
     options:
         members:

optimistix/__init__.py

@@ -77,6 +77,7 @@
     OptaxMinimiser as OptaxMinimiser,
     polak_ribiere as polak_ribiere,
     SteepestDescent as SteepestDescent,
+    WolfeSearch as WolfeSearch,
 )
 
 

optimistix/_search.py

@@ -320,6 +320,15 @@ class AbstractSearch(eqx.Module, Generic[Y, _FnInfo, _FnEvalInfo, SearchState]):
     See [this documentation](./introduction.md) for more information.
     """
 
+    requires_grad_eval: ClassVar[bool] = False
+    """Whether [`optimistix.AbstractSearch.step`][] needs gradient information at the
+    trial point `y_eval`, i.e. whether `f_eval_info` must carry a gradient. Searches
+    that test a curvature condition (e.g. [`optimistix.WolfeSearch`][]) set this to
+    `True`, and the calling solver then evaluates the trial-point gradient eagerly.
+    Defaults to `False`, so that searches needing only function values (e.g.
+    [`optimistix.BacktrackingArmijo`][]) incur no extra work on rejected steps.
+    """
+
     @abc.abstractmethod
     def init(self, y: Y, f_info_struct: _FnInfo) -> SearchState:
         """Is called just once, at the very start of the entire optimisation problem.

optimistix/_solver/__init__.py

@@ -49,3 +49,4 @@
     ClassicalTrustRegion as ClassicalTrustRegion,
     LinearTrustRegion as LinearTrustRegion,
 )
+from .wolfe import WolfeSearch as WolfeSearch

optimistix/_solver/gradient_methods.py

@@ -171,17 +171,29 @@ def step(
         f_eval, lin_fn, aux_eval = jax.linearize(
             lambda _y: fn(_y, args), state.y_eval, has_aux=True
         )
+        # Some searches (e.g. `WolfeSearch`) need the gradient at the trial point to
+        # check a curvature condition. Compute it eagerly only when required, and reuse
+        # it on acceptance; otherwise defer it to the `accepted` branch as usual.
+        if self.search.requires_grad_eval:
+            grad_eval = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype)
+            f_eval_info_for_search = FunctionInfo.EvalGrad(f_eval, grad_eval)
+        else:
+            grad_eval = None
+            f_eval_info_for_search = FunctionInfo.Eval(f_eval)
         step_size, accept, search_result, search_state = self.search.step(
             state.first_step,
             y,
             state.y_eval,
             state.f_info,
-            FunctionInfo.Eval(f_eval),
+            f_eval_info_for_search,
             state.search_state,
         )
 
         def accepted(descent_state):
-            grad = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype)
+            if grad_eval is None:
+                grad = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype)
+            else:
+                grad = grad_eval
 
             f_eval_info = FunctionInfo.EvalGrad(f_eval, grad)
             descent_state = self.descent.query(state.y_eval, f_eval_info, descent_state)

optimistix/_solver/limited_memory_bfgs.py

@@ -18,6 +18,7 @@
     tree_dot,
 )
 from .._search import (
+    AbstractSearch,
     FunctionInfo,
 )
 from .backtracking import BacktrackingArmijo
@@ -288,8 +289,9 @@ def _lbfgs_hessian_operator_fn(
             lambda x: jnp.einsum("h,h...->...", descent[: state.history_length], x)
         )
         - ω(state.y_diff_history).call(
-            lambda x: gamma_k
-            * jnp.einsum("h,h...->...", descent[state.history_length :], x),
+            lambda x: (
+                gamma_k * jnp.einsum("h,h...->...", descent[state.history_length :], x)
+            ),
         )
     ).ω
     return descent_step
@@ -568,7 +570,7 @@ class LBFGS(AbstractLBFGS[Y, Aux, _Hessian, _LBFGSUpdateState]):
     norm: Callable[[PyTree], Scalar]
     use_inverse: bool
     descent: NewtonDescent
-    search: BacktrackingArmijo
+    search: AbstractSearch
     history_length: int
     verbose: Callable[..., None]
 
@@ -580,13 +582,14 @@ def __init__(
         use_inverse: bool = True,
         history_length: int = 10,
         verbose: bool | Callable[..., None] = False,
+        search: AbstractSearch | None = None,
     ):
         self.rtol = rtol
         self.atol = atol
         self.norm = norm
         self.use_inverse = use_inverse
         self.descent = NewtonDescent()
-        self.search = BacktrackingArmijo()
+        self.search = BacktrackingArmijo() if search is None else search
         self.history_length = history_length
         self.verbose = default_verbose(verbose)
 
@@ -613,4 +616,7 @@ def __init__(
     or (for customisation) a callable `**kwargs -> None`. If provided as a callable then
     each value will be a 2-tuple of `(str, jax.Array)` providing a human-readable name
     and its corresponding value.
+- `search`: The line search to use. Defaults to [`optimistix.BacktrackingArmijo`][].
+    Pass [`optimistix.WolfeSearch`][] for a strong-Wolfe line search, which often
+    converges in fewer iterations.
 """

optimistix/_solver/quasi_newton.py

@@ -214,17 +214,29 @@ def step(
         f_eval, lin_fn, aux_eval = jax.linearize(
             lambda _y: fn(_y, args), state.y_eval, has_aux=True
         )
+        # Some searches (e.g. `WolfeSearch`) need the gradient at the trial point to
+        # check a curvature condition. Compute it eagerly only when required, and reuse
+        # it on acceptance; otherwise defer it to the `accepted` branch as usual.
+        if self.search.requires_grad_eval:
+            grad_eval = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype)
+            f_eval_info_for_search = FunctionInfo.EvalGrad(f_eval, grad_eval)
+        else:
+            grad_eval = None
+            f_eval_info_for_search = FunctionInfo.Eval(f_eval)
         step_size, accept, search_result, search_state = self.search.step(
             state.first_step,
             y,
             state.y_eval,
             state.f_info,
-            FunctionInfo.Eval(f_eval),
+            f_eval_info_for_search,
             state.search_state,
         )
 
         def accepted(descent_state):
-            grad = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype)
+            if grad_eval is None:
+                grad = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype)
+            else:
+                grad = grad_eval
 
             f_eval_info, hessian_update_state = self.update_hessian(
                 y,
@@ -437,7 +449,7 @@ class BFGS(AbstractBFGS[Y, Aux, _Hessian]):
     norm: Callable[[PyTree], Scalar]
     use_inverse: bool
     descent: NewtonDescent
-    search: BacktrackingArmijo
+    search: AbstractSearch
     verbose: Callable[..., None]
 
     def __init__(
@@ -447,14 +459,14 @@ def __init__(
         norm: Callable[[PyTree], Scalar] = max_norm,
         use_inverse: bool = True,
         verbose: bool | Callable[..., None] = False,
+        search: AbstractSearch | None = None,
     ):
         self.rtol = rtol
         self.atol = atol
         self.norm = norm
         self.use_inverse = use_inverse
         self.descent = NewtonDescent(linear_solver=lx.Cholesky())
-        # TODO(raderj): switch out `BacktrackingArmijo` with a better line search.
-        self.search = BacktrackingArmijo()
+        self.search = BacktrackingArmijo() if search is None else search
         self.verbose = default_verbose(verbose)
 
 
@@ -481,6 +493,11 @@ def __init__(
     or (for customisation) a callable `**kwargs -> None`. If provided as a callable then
     each value will be a 2-tuple of `(str, jax.Array)` providing a human-readable name
     and its corresponding value.
+- `search`: The line search to use. Defaults to [`optimistix.BacktrackingArmijo`][].
+    Pass [`optimistix.WolfeSearch`][] for a strong-Wolfe line search, which often
+    converges in fewer iterations for quasi-Newton methods. Note that searches using
+    the Hessian approximation `B` (e.g. [`optimistix.ClassicalTrustRegion`][]) are
+    incompatible with `use_inverse=True`.
 """
 
 
@@ -598,7 +615,7 @@ class DFP(AbstractDFP[Y, Aux, _Hessian]):
     norm: Callable[[PyTree], Scalar]
     use_inverse: bool
     descent: NewtonDescent
-    search: BacktrackingArmijo
+    search: AbstractSearch
     verbose: Callable[..., None]
 
     def __init__(
@@ -608,14 +625,14 @@ def __init__(
         norm: Callable[[PyTree], Scalar] = max_norm,
         use_inverse: bool = True,
         verbose: bool | Callable[..., None] = False,
+        search: AbstractSearch | None = None,
     ):
         self.rtol = rtol
         self.atol = atol
         self.norm = norm
         self.use_inverse = use_inverse
         self.descent = NewtonDescent(linear_solver=lx.Cholesky())
-        # TODO(raderj): switch out `BacktrackingArmijo` with a better line search.
-        self.search = BacktrackingArmijo()
+        self.search = BacktrackingArmijo() if search is None else search
         self.verbose = default_verbose(verbose)
 
 
@@ -642,4 +659,9 @@ def __init__(
     or (for customisation) a callable `**kwargs -> None`. If provided as a callable then
     each value will be a 2-tuple of `(str, jax.Array)` providing a human-readable name
     and its corresponding value.
+- `search`: The line search to use. Defaults to [`optimistix.BacktrackingArmijo`][].
+    Pass [`optimistix.WolfeSearch`][] for a strong-Wolfe line search, which often
+    converges in fewer iterations for quasi-Newton methods. Note that searches using
+    the Hessian approximation `B` (e.g. [`optimistix.ClassicalTrustRegion`][]) are
+    incompatible with `use_inverse=True`.
 """