diff --git a/optimistix/_solver/gauss_newton.py b/optimistix/_solver/gauss_newton.py index 67918412..38a850da 100644 --- a/optimistix/_solver/gauss_newton.py +++ b/optimistix/_solver/gauss_newton.py @@ -17,6 +17,7 @@ max_norm, sum_squares, tree_full_like, + tree_where, verbose_print, ) from .._search import ( @@ -331,6 +332,7 @@ def rejected(descent_state): search_result == RESULTS.successful, descent_result, search_result ) + prev_aux = tree_where(state.first_step, aux, state.aux) state = _GaussNewtonState( first_step=jnp.array(False), y_eval=y_eval, @@ -344,7 +346,7 @@ def rejected(descent_state): num_accepted_steps=num_accepted_steps, num_steps_since_acceptance=num_steps_since_acceptance, ) - return y, state, aux + return y, state, prev_aux def terminate( self, diff --git a/optimistix/_solver/gradient_methods.py b/optimistix/_solver/gradient_methods.py index bb6c31bd..6efd4aef 100644 --- a/optimistix/_solver/gradient_methods.py +++ b/optimistix/_solver/gradient_methods.py @@ -16,6 +16,7 @@ lin_to_grad, max_norm, tree_full_like, + tree_where, ) from .._search import ( AbstractDescent, @@ -207,6 +208,7 @@ def rejected(descent_state): search_result == RESULTS.successful, descent_result, search_result ) + prev_aux = tree_where(state.first_step, aux, state.aux) state = _GradientDescentState( first_step=jnp.array(False), y_eval=y_eval, @@ -217,7 +219,7 @@ def rejected(descent_state): terminate=terminate, result=result, ) - return y, state, aux + return y, state, prev_aux def terminate( self, diff --git a/optimistix/_solver/quasi_newton.py b/optimistix/_solver/quasi_newton.py index d664c0de..c55c3fb2 100644 --- a/optimistix/_solver/quasi_newton.py +++ b/optimistix/_solver/quasi_newton.py @@ -20,6 +20,7 @@ max_norm, tree_dot, tree_full_like, + tree_where, verbose_print, ) from .._search import ( @@ -286,6 +287,7 @@ def rejected(descent_state): search_result == RESULTS.successful, descent_result, search_result ) + prev_aux = tree_where(state.first_step, aux, state.aux) state = _QuasiNewtonState( first_step=jnp.array(False), y_eval=y_eval, @@ -298,7 +300,7 @@ def rejected(descent_state): num_accepted_steps=state.num_accepted_steps + jnp.where(accept, 1, 0), hessian_update_state=hessian_update_state, ) - return y, state, aux + return y, state, prev_aux def terminate( self, diff --git a/tests/test_best_so_far.py b/tests/test_best_so_far.py index 3697884a..a1551621 100644 --- a/tests/test_best_so_far.py +++ b/tests/test_best_so_far.py @@ -1,5 +1,6 @@ import jax.numpy as jnp import optimistix as optx +import pytest def test_fixed_point(): @@ -41,3 +42,25 @@ def fn(y, _): solver = optx.BestSoFarMinimiser(solver) sol = optx.minimise(fn, solver, jnp.array(0.0)) assert jnp.allclose(sol.value, 0.96118069, rtol=1e-5, atol=1e-5) + + +# https://github.com/patrick-kidger/optimistix/issues/33 +@pytest.mark.parametrize( + "solver", (optx.BFGS(rtol=1e-5, atol=1e-5), optx.NonlinearCG(atol=1e-5, rtol=1e-5)) +) +def test_checks_last_point_minimiser(solver): + def fn(y, _): + return (y - 3.0) ** 2 + + solver = optx.BestSoFarMinimiser(solver) + sol = optx.minimise(fn, solver, jnp.array(0.0)) + assert sol.value == 3.0 + + +def test_checks_last_point_least_squares(): + def fn(y, _): + return y - 3.0 + + solver = optx.BestSoFarLeastSquares(optx.GaussNewton(rtol=1e-5, atol=1e-5)) + sol = optx.least_squares(fn, solver, jnp.array(0.0)) + assert sol.value == 3.0