From ff1b70a077008f207f430d6bf9711c2b6649e24f Mon Sep 17 00:00:00 2001 From: MashdedP Date: Tue, 9 Jun 2026 16:13:54 +0200 Subject: [PATCH] Add WolfeSearch (strong Wolfe line search) and wire it into (L)BFGS WolfeSearch implements the strong Wolfe conditions via the bracketing/zoom algorithm of Nocedal & Wright (Alg. 3.5-3.6). Since Optimistix calls AbstractSearch.step once per solver step, the textbook nested loops are expressed as a single per-call state machine dispatched by a `mode` flag. The curvature condition needs the gradient at each trial point. To avoid penalising searches that only need function values (e.g. BacktrackingArmijo), AbstractSearch gains an opt-in `requires_grad_eval` class flag; the gradient/ quasi-Newton solvers compute the trial-point gradient eagerly only when the search requests it, reusing it on acceptance. BFGS, DFP and LBFGS gain an optional `search=` argument (default unchanged: BacktrackingArmijo) so users can select WolfeSearch. Co-Authored-By: Claude Opus 4.8 --- docs/api/searches/searches.md | 7 + optimistix/__init__.py | 1 + optimistix/_search.py | 9 + optimistix/_solver/__init__.py | 1 + optimistix/_solver/gradient_methods.py | 16 +- optimistix/_solver/limited_memory_bfgs.py | 14 +- optimistix/_solver/quasi_newton.py | 38 +++- optimistix/_solver/wolfe.py | 229 ++++++++++++++++++++++ tests/test_wolfe_search.py | 108 ++++++++++ 9 files changed, 409 insertions(+), 14 deletions(-) create mode 100644 optimistix/_solver/wolfe.py create mode 100644 tests/test_wolfe_search.py diff --git a/docs/api/searches/searches.md b/docs/api/searches/searches.md index 9dd01d2b..d4bebb5e 100644 --- a/docs/api/searches/searches.md +++ b/docs/api/searches/searches.md @@ -21,6 +21,13 @@ --- +::: optimistix.WolfeSearch + options: + members: + - __init__ + +--- + ::: optimistix.ClassicalTrustRegion options: members: diff --git a/optimistix/__init__.py b/optimistix/__init__.py index 309d429f..05f501f2 100644 --- a/optimistix/__init__.py +++ b/optimistix/__init__.py @@ -77,6 +77,7 @@ OptaxMinimiser as OptaxMinimiser, polak_ribiere as polak_ribiere, SteepestDescent as SteepestDescent, + WolfeSearch as WolfeSearch, ) diff --git a/optimistix/_search.py b/optimistix/_search.py index 2f4ec172..cbd2cbab 100644 --- a/optimistix/_search.py +++ b/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. diff --git a/optimistix/_solver/__init__.py b/optimistix/_solver/__init__.py index 90e2832d..e2f5c035 100644 --- a/optimistix/_solver/__init__.py +++ b/optimistix/_solver/__init__.py @@ -49,3 +49,4 @@ ClassicalTrustRegion as ClassicalTrustRegion, LinearTrustRegion as LinearTrustRegion, ) +from .wolfe import WolfeSearch as WolfeSearch diff --git a/optimistix/_solver/gradient_methods.py b/optimistix/_solver/gradient_methods.py index 67a86498..a17ed396 100644 --- a/optimistix/_solver/gradient_methods.py +++ b/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) diff --git a/optimistix/_solver/limited_memory_bfgs.py b/optimistix/_solver/limited_memory_bfgs.py index f90ee4fb..ca574419 100644 --- a/optimistix/_solver/limited_memory_bfgs.py +++ b/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. """ diff --git a/optimistix/_solver/quasi_newton.py b/optimistix/_solver/quasi_newton.py index 94f38810..9719f8db 100644 --- a/optimistix/_solver/quasi_newton.py +++ b/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`. """ diff --git a/optimistix/_solver/wolfe.py b/optimistix/_solver/wolfe.py new file mode 100644 index 00000000..68337a13 --- /dev/null +++ b/optimistix/_solver/wolfe.py @@ -0,0 +1,229 @@ +from typing import cast, ClassVar, TypeAlias + +import equinox as eqx +import jax.numpy as jnp +from equinox.internal import ω +from jaxtyping import Array, Bool, Int, Scalar, ScalarLike + +from .._custom_types import Y +from .._search import AbstractSearch, FunctionInfo +from .._solution import RESULTS + + +# `mode` flags for the line-search state machine. +_BRACKET = 0 +_ZOOM = 1 + + +# `f_info`, at the last accepted point, must carry a gradient so that we can form the +# directional derivative `phi'(0)`. +_FnInfo: TypeAlias = ( + FunctionInfo.EvalGrad + | FunctionInfo.EvalGradHessian + | FunctionInfo.EvalGradHessianInv +) +# `f_eval_info`, at the trial point, must also carry a gradient so that we can form +# `phi'(alpha)` for the curvature condition. +_FnEvalInfo: TypeAlias = FunctionInfo.EvalGrad + + +class _WolfeState(eqx.Module): + mode: Int[Array, ""] + step_size: Scalar # the step size that produced the current `y_eval` + a_prev: Scalar # previous bracket point (bracketing phase) + phi_prev: Scalar + a_lo: Scalar # zoom interval [a_lo, a_hi]; `a_lo` is the best point so far + phi_lo: Scalar + a_hi: Scalar + first_in_search: Bool[Array, ""] # True before the first bracketing evaluation + num_evals: Int[Array, ""] + + +class WolfeSearch(AbstractSearch[Y, _FnInfo, _FnEvalInfo, _WolfeState]): + """A line search satisfying the strong Wolfe conditions. + + That is, this finds a step size `alpha` such that, writing + `phi(alpha) = f(y + alpha p)` for the descent direction `p`, + + - `phi(alpha) <= phi(0) + c1 alpha phi'(0)` (the Armijo / sufficient-decrease + condition), and + - `|phi'(alpha)| <= c2 |phi'(0)|` (the strong curvature condition). + + This is the line search most commonly paired with quasi-Newton methods such as + [`optimistix.BFGS`][] and [`optimistix.LBFGS`][]: the curvature condition keeps the + step large enough for the Hessian update to remain well-conditioned. + + The search follows the bracketing/zoom algorithm of Nocedal & Wright, "Numerical + Optimization" (2nd ed.), Algorithms 3.5 and 3.6, using bisection for the zoom + interpolation. As Optimistix calls [`optimistix.AbstractSearch.step`][] once per + solver step (the iteration loop belongs to the solver), the two nested loops of the + textbook algorithm are expressed here as a single per-call state machine, dispatched + by a `mode` flag. + + This search requires the gradient at each trial point, so it must be used with a + solver that provides one. (It sets `requires_grad_eval = True`; see + [`optimistix.AbstractSearch`][].) All of Optimistix's gradient-based minimisers do + so. + """ + + requires_grad_eval: ClassVar[bool] = True + + c1: ScalarLike = 1e-4 + c2: ScalarLike = 0.9 + step_init: ScalarLike = 1.0 + step_max: ScalarLike = 1e6 + max_evals: int = 30 + + def __post_init__(self): + self.c1 = eqx.error_if( + self.c1, + (self.c1 <= 0) | (self.c1 >= self.c2), # pyright: ignore + "`WolfeSearch` requires `0 < c1 < c2 < 1`.", + ) + self.c2 = eqx.error_if( + self.c2, + (self.c2 <= self.c1) | (self.c2 >= 1), # pyright: ignore + "`WolfeSearch` requires `0 < c1 < c2 < 1`.", + ) + self.step_init = eqx.error_if( + self.step_init, + self.step_init <= 0, # pyright: ignore + "`WolfeSearch(step_init=...)` must be strictly greater than 0.", + ) + + def init(self, y: Y, f_info_struct: _FnInfo) -> _WolfeState: + del y, f_info_struct + zero = jnp.array(0.0) + return _WolfeState( + mode=jnp.array(_BRACKET), + step_size=jnp.array(self.step_init), + a_prev=zero, + phi_prev=zero, + a_lo=zero, + phi_lo=zero, + a_hi=zero, + first_in_search=jnp.array(True), + num_evals=jnp.array(0), + ) + + def step( + self, + first_step: Bool[Array, ""], + y: Y, + y_eval: Y, + f_info: _FnInfo, + f_eval_info: _FnEvalInfo, + state: _WolfeState, + ) -> tuple[Scalar, Bool[Array, ""], RESULTS, _WolfeState]: + if not isinstance( + f_info, + ( + FunctionInfo.EvalGrad, + FunctionInfo.EvalGradHessian, + FunctionInfo.EvalGradHessianInv, + ), + ): + raise ValueError( + "Cannot use `WolfeSearch` with this solver. This is because " + "`WolfeSearch` requires gradients of the target function, but this " + "solver does not evaluate such gradients." + ) + if not isinstance(f_eval_info, FunctionInfo.EvalGrad): + raise ValueError( + "Cannot use `WolfeSearch` with this solver. This is because " + "`WolfeSearch` requires the gradient at each trial point, which this " + "solver does not provide." + ) + + # Restrict to the 1D line `phi(a) = f(y + a p)`. We never need `p` explicitly: + # with `y_diff = y_eval - y = alpha p`, the directional derivatives are + # `phi'(0) = = / alpha` and likewise at the trial + # point. `compute_grad_dot(y_diff)` returns the relevant ``. + y_diff = (y_eval**ω - y**ω).ω + alpha = state.step_size + safe_alpha = jnp.where(alpha > 0, alpha, 1.0) + phi0 = f_info.as_min() + phi = f_eval_info.as_min() + dphi0 = f_info.compute_grad_dot(y_diff) / safe_alpha + dphi = f_eval_info.compute_grad_dot(y_diff) / safe_alpha + + armijo = phi <= phi0 + self.c1 * alpha * dphi0 + curvature = jnp.abs(dphi) <= -self.c2 * dphi0 + + # Bracketing transition (Nocedal & Wright Alg. 3.5). The minimum is bracketed by + # [a_prev, alpha] if sufficient decrease fails or `phi` stopped decreasing; by + # [alpha, a_prev] if the slope has turned positive; otherwise we grow the step. + b_zoom_prev_alpha = (~armijo) | ( + (phi >= state.phi_prev) & ~state.first_in_search + ) + b_accept = curvature & ~b_zoom_prev_alpha + b_zoom_alpha_prev = (dphi >= 0) & ~b_zoom_prev_alpha & ~b_accept + b_grow = ~b_zoom_prev_alpha & ~b_accept & ~b_zoom_alpha_prev + b_a_lo = jnp.where(b_zoom_prev_alpha, state.a_prev, alpha) + b_phi_lo = jnp.where(b_zoom_prev_alpha, state.phi_prev, phi) + b_a_hi = jnp.where(b_zoom_prev_alpha, alpha, state.a_prev) + b_mode = jnp.where(b_grow, _BRACKET, _ZOOM) + b_next = jnp.where( + b_grow, jnp.minimum(2.0 * alpha, self.step_max), 0.5 * (b_a_lo + b_a_hi) + ) + + # Zoom transition (Nocedal & Wright Alg. 3.6), by bisection of [a_lo, a_hi]. + z_shrink_hi = (~armijo) | (phi >= state.phi_lo) + z_accept = ~z_shrink_hi & curvature + z_flip = ~z_shrink_hi & ~curvature & (dphi * (state.a_hi - state.a_lo) >= 0) + z_a_hi = jnp.where( + z_shrink_hi, alpha, jnp.where(z_flip, state.a_lo, state.a_hi) + ) + z_a_lo = jnp.where(z_shrink_hi, state.a_lo, alpha) + z_phi_lo = jnp.where(z_shrink_hi, state.phi_lo, phi) + z_next = 0.5 * (z_a_lo + z_a_hi) + + # Select the active branch by the current mode. + in_bracket = state.mode == _BRACKET + accept_ls = jnp.where(in_bracket, b_accept, z_accept) + cont_mode = jnp.where(in_bracket, b_mode, _ZOOM) + cont_next = jnp.where(in_bracket, b_next, z_next) + cont_a_prev = jnp.where(in_bracket, alpha, state.a_prev) + cont_phi_prev = jnp.where(in_bracket, phi, state.phi_prev) + cont_a_lo = jnp.where(in_bracket, b_a_lo, z_a_lo) + cont_phi_lo = jnp.where(in_bracket, b_phi_lo, z_phi_lo) + cont_a_hi = jnp.where(in_bracket, b_a_hi, z_a_hi) + + num_evals = state.num_evals + 1 + out_of_budget = num_evals >= self.max_evals + # On the first solver step, `y`/`f_info` are dummy data and we must accept. On + # budget exhaustion we accept the best point seen so far, leaving the solver's + # gradient-based termination to judge convergence (a best-effort fallback, as in + # SciPy). This is benign near an optimum, where `phi'(0) -> 0` makes the strong + # curvature condition impossible to satisfy exactly. + accept = first_step | accept_ls | out_of_budget + + # On acceptance a new line search begins along a new descent direction, so we + # reset to `step_init`. `first_in_search` guards the first bracket comparison. + step_size = cast(Scalar, jnp.where(accept, self.step_init, cont_next)) + new_state = _WolfeState( + mode=jnp.where(accept, _BRACKET, cont_mode), + step_size=step_size, + a_prev=jnp.where(accept, 0.0, cont_a_prev), + phi_prev=jnp.where(accept, phi, cont_phi_prev), + a_lo=jnp.where(accept, 0.0, cont_a_lo), + phi_lo=jnp.where(accept, phi0, cont_phi_lo), + a_hi=jnp.where(accept, 0.0, cont_a_hi), + first_in_search=jnp.where(accept, True, False), + num_evals=jnp.where(accept, 0, num_evals), + ) + return step_size, accept, RESULTS.successful, new_state + + +WolfeSearch.__init__.__doc__ = """**Arguments:** + +- `c1`: the Armijo (sufficient decrease) constant. Must satisfy `0 < c1 < c2 < 1`. +- `c2`: the curvature constant. Must satisfy `0 < c1 < c2 < 1`. Smaller values impose a + stricter curvature condition; `0.9` is the usual choice for quasi-Newton methods and + `0.1` for nonlinear conjugate gradients. +- `step_init`: the step size tried first along each new descent direction. Must be + strictly greater than 0. +- `step_max`: an upper bound on the step size during bracketing. +- `max_evals`: the maximum number of trial points evaluated per line search, after + which the best point seen so far is accepted. +""" diff --git a/tests/test_wolfe_search.py b/tests/test_wolfe_search.py new file mode 100644 index 00000000..5bdce9b6 --- /dev/null +++ b/tests/test_wolfe_search.py @@ -0,0 +1,108 @@ +import jax +import jax.numpy as jnp +import optimistix as optx +import pytest + +from .helpers import tree_allclose + + +# Test functions and a descent direction at the start point. +def _rosenbrock(x, args=None): + return jnp.sum(100.0 * (x[1:] - x[:-1] ** 2) ** 2 + (1.0 - x[:-1]) ** 2) + + +def _quadratic(x, args=None): + matrix = jnp.array([[3.0, 0.5], [0.5, 2.0]]) + offset = jnp.array([1.0, -2.0]) + return 0.5 * x @ matrix @ x - offset @ x + + +_LINE_SEARCH_CASES = [ + (_rosenbrock, jnp.array([-1.2, 1.0])), + (_quadratic, jnp.array([5.0, 5.0])), +] + +_MINIMISE_CASES = [ + (_rosenbrock, jnp.array([-1.2, 1.0]), jnp.array([1.0, 1.0])), + (_rosenbrock, jnp.full((6,), -1.0), jnp.ones((6,))), + (_quadratic, jnp.array([5.0, 5.0]), jnp.array([0.52173913, -1.13043478])), +] + + +def _run_single_line_search(search, fn, y, p): + """Drive `search` along the fixed direction `p`, exactly as a solver would, and + return the accepted step size.""" + value_and_grad = jax.value_and_grad(fn) + f0, g0 = value_and_grad(y) + f_info = optx.FunctionInfo.EvalGrad(f0, g0) + f_info_struct = jax.eval_shape(lambda: f_info) + state = search.init(y, f_info_struct) + + # First solver step: dummy data, accept is forced, returns the initial step size. + step_size, _, _, state = search.step(jnp.array(True), y, y, f_info, f_info, state) + alpha = step_size + for _ in range(40): + y_eval = y + alpha * p + f_eval, g_eval = value_and_grad(y_eval) + f_eval_info = optx.FunctionInfo.EvalGrad(f_eval, g_eval) + step_size, accept, _, state = search.step( + jnp.array(False), y, y_eval, f_info, f_eval_info, state + ) + if bool(accept): + return float(alpha) # the step size that produced the accepted `y_eval` + alpha = step_size + raise AssertionError("line search did not terminate") + + +@pytest.mark.parametrize("c1, c2", [(1e-4, 0.9), (1e-3, 0.4), (0.1, 0.5)]) +@pytest.mark.parametrize("fn, y", _LINE_SEARCH_CASES) +def test_satisfies_strong_wolfe(fn, y, c1, c2): + value_and_grad = jax.value_and_grad(fn) + _, g0 = value_and_grad(y) + p = -g0 # steepest descent direction + search = optx.WolfeSearch(c1=c1, c2=c2) + alpha = _run_single_line_search(search, fn, y, p) + + phi0 = fn(y) + dphi0 = jnp.vdot(g0, p) + assert dphi0 < 0 + f_alpha, g_alpha = value_and_grad(y + alpha * p) + dphi_alpha = jnp.vdot(g_alpha, p) + # Strong Wolfe: sufficient decrease and curvature. + assert f_alpha <= phi0 + c1 * alpha * dphi0 + 1e-10 + assert jnp.abs(dphi_alpha) <= -c2 * dphi0 + 1e-10 + + +def test_requires_grad_eval_flag(): + assert optx.WolfeSearch.requires_grad_eval is True + assert optx.BacktrackingArmijo.requires_grad_eval is False + assert optx.AbstractSearch.requires_grad_eval is False + + +@pytest.mark.parametrize("solver_cls", [optx.BFGS, optx.LBFGS]) +@pytest.mark.parametrize("use_inverse", [True, False]) +@pytest.mark.parametrize("fn, y0, expected", _MINIMISE_CASES) +def test_minimise_with_wolfe(solver_cls, use_inverse, fn, y0, expected): + solver = solver_cls( + rtol=1e-9, atol=1e-9, use_inverse=use_inverse, search=optx.WolfeSearch() + ) + sol = optx.minimise(fn, solver, y0, max_steps=1024, throw=False) + assert sol.result == optx.RESULTS.successful + assert tree_allclose(sol.value, expected, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("fn, y0, expected", _MINIMISE_CASES) +def test_wolfe_jit(fn, y0, expected): + solver = optx.LBFGS(rtol=1e-9, atol=1e-9, search=optx.WolfeSearch()) + run = jax.jit( + lambda y: optx.minimise(fn, solver, y, max_steps=1024, throw=False).value + ) + assert tree_allclose(run(y0), expected, rtol=1e-4, atol=1e-4) + + +def test_backtracking_still_works(): + # The opt-in eager-gradient path must not change the default-search behaviour. + fn, y0, expected = _MINIMISE_CASES[2] # quadratic + armijo = optx.LBFGS(rtol=1e-9, atol=1e-9) + default = optx.minimise(fn, armijo, y0, max_steps=1024, throw=False) + assert tree_allclose(default.value, expected, rtol=1e-4, atol=1e-4)