Optionally pass more info at y_eval to searches

optimistix/_search.py

@@ -29,6 +29,7 @@
 import jax.numpy as jnp
 import jax.tree_util as jtu
 import lineax as lx
+from equinox import AbstractClassVar
 from jaxtyping import Array, Bool, Scalar
 
 from ._custom_types import (
@@ -79,56 +80,39 @@ def as_min(self):
 
 
 # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
-class EvalGrad(FunctionInfo, Generic[Y]):
+class EvalGrad(Eval, Generic[Y]):
     """Has a `.f` attribute as with [`optimistix.FunctionInfo.Eval`][]. Also has a
     `.grad` attribute describing `d(fn)/dy`. Used with first-order solvers for
     minimisation problems. (E.g. gradient descent; nonlinear CG.)
     """
 
-    f: Scalar
     grad: Y
 
-    def as_min(self):
-        return self.f
+    def compute_grad(self):
+        return self.grad
 
     def compute_grad_dot(self, y: Y):
         return tree_dot(self.grad, y)
 
 
 # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
-class EvalGradHessian(FunctionInfo, Generic[Y]):
+class EvalGradHessian(EvalGrad[Y], Generic[Y]):
     """Has `.f` and `.grad` attributes as with [`optimistix.FunctionInfo.EvalGrad`][].
     Also has a `.hessian` attribute describing (an approximation to) the Hessian of
     `fn` at `y`. Used with quasi-Newton minimisation algorithms, like BFGS.
     """
 
-    f: Scalar
-    grad: Y
     hessian: lx.AbstractLinearOperator
 
-    def as_min(self):
-        return self.f
-
-    def compute_grad_dot(self, y: Y):
-        return tree_dot(self.grad, y)
-
 
 # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
-class EvalGradHessianInv(FunctionInfo, Generic[Y]):
+class EvalGradHessianInv(EvalGrad[Y], Generic[Y]):
     """As [`optimistix.FunctionInfo.EvalGradHessian`][], but records the (approximate)
     inverse-Hessian instead. Has `.f` and `.grad` and `.hessian_inv` attributes.
     """
 
-    f: Scalar
-    grad: Y
     hessian_inv: lx.AbstractLinearOperator
 
-    def as_min(self):
-        return self.f
-
-    def compute_grad_dot(self, y: Y):
-        return tree_dot(self.grad, y)
-
 
 # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
 class Residual(FunctionInfo, Generic[Out]):
@@ -143,18 +127,14 @@ def as_min(self):
 
 
 # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
-class ResidualJac(FunctionInfo, Generic[Y, Out]):
+class ResidualJac(Residual[Out], Generic[Y, Out]):
     """Records the Jacobian `d(fn)/dy` as a linear operator. Used for least squares
     problems, for which `fn` returns residuals. Has `.residual` and `.jac` attributes,
     where `residual = fn(y)`, `jac = d(fn)/dy`.
     """
 
-    residual: Out
     jac: lx.AbstractLinearOperator
 
-    def as_min(self):
-        return 0.5 * sum_squares(self.residual)
-
     def compute_grad(self):
         conj_residual = jtu.tree_map(jnp.conj, self.residual)
         return self.jac.transpose().mv(conj_residual)
@@ -320,6 +300,8 @@ class AbstractSearch(eqx.Module, Generic[Y, _FnInfo, _FnEvalInfo, SearchState]):
     See [this documentation](./introduction.md) for more information.
     """
 
+    info_needed_at_y_eval: AbstractClassVar[type[FunctionInfo]]
+
     @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.
@@ -363,7 +345,7 @@ def step(
         - `f_info`: An [`optimistix.FunctionInfo`][] describing information about `f`
             evaluated at `y`, the gradient of `f` at `y`, etc.
         - `f_eval_info`: An [`optimistix.FunctionInfo`][] describing information about
-            `f` evaluated at `y`, the gradient of `f` at `y`, etc.
+            `f` evaluated at `y_eval`, the gradient of `f` at `y_eval`, etc.
         - `state`: the evolving state of the repeated searches.
 
         **Returns:**

optimistix/_solver/_f_eval_info_helpers.py

@@ -0,0 +1,42 @@
+from collections.abc import Callable
+
+from jaxtyping import Scalar
+
+from .._custom_types import Y
+from .._misc import lin_to_grad
+from .._search import FunctionInfo
+
+
+def make_eval_info_for_search(
+    info_needed: type[FunctionInfo],
+    y_eval: Y,
+    f_eval: Scalar,
+    lin_fn: Callable[[Y], Scalar],
+    autodiff_mode: str,
+) -> FunctionInfo.Eval | FunctionInfo.EvalGrad:
+    """Build the cheapest scalar `FunctionInfo` satisfying `info_needed`."""
+    if issubclass(FunctionInfo.Eval, info_needed):
+        return FunctionInfo.Eval(f_eval)
+    elif issubclass(FunctionInfo.EvalGrad, info_needed):
+        grad = lin_to_grad(lin_fn, y_eval, autodiff_mode, f_eval.dtype)
+        return FunctionInfo.EvalGrad(f_eval, grad)
+    else:
+        raise ValueError(
+            f"Cannot provide requested function information at `y_eval`: {info_needed}."
+            " Only FunctionInfo.Eval and FunctionInfo.EvalGrad are supported."
+        )
+
+
+def promote_to_eval_grad(
+    f_eval_info: FunctionInfo.Eval | FunctionInfo.EvalGrad,
+    y_eval: Y,
+    f_eval: Scalar,
+    lin_fn: Callable[[Y], Scalar],
+    autodiff_mode: str,
+) -> FunctionInfo.EvalGrad:
+    """Reuse an existing scalar gradient evaluation, or compute one if needed."""
+    if isinstance(f_eval_info, FunctionInfo.EvalGrad):
+        return f_eval_info
+    else:
+        grad = lin_to_grad(lin_fn, y_eval, autodiff_mode, f_eval.dtype)
+        return FunctionInfo.EvalGrad(f_eval, grad)

optimistix/_solver/backtracking.py

@@ -1,4 +1,4 @@
-from typing import cast, TypeAlias
+from typing import cast, ClassVar, TypeAlias
 
 import equinox as eqx
 import jax.numpy as jnp
@@ -29,6 +29,7 @@ class BacktrackingArmijo(AbstractSearch[Y, _FnInfo, _FnEvalInfo, _BacktrackingSt
     decrease_factor: ScalarLike = 0.5
     slope: ScalarLike = 0.1
     step_init: ScalarLike = 1.0
+    info_needed_at_y_eval: ClassVar[type[FunctionInfo]] = FunctionInfo
 
     def __post_init__(self):
         self.decrease_factor = eqx.error_if(

optimistix/_solver/gradient_methods.py

@@ -13,7 +13,6 @@
 from .._misc import (
     cauchy_termination,
     filter_cond,
-    lin_to_grad,
     max_norm,
     tree_full_like,
     tree_where,
@@ -24,6 +23,7 @@
     FunctionInfo,
 )
 from .._solution import RESULTS
+from ._f_eval_info_helpers import make_eval_info_for_search, promote_to_eval_grad
 from .learning_rate import LearningRate
 
 
@@ -132,7 +132,9 @@ class AbstractGradientDescent(AbstractMinimiser[Y, Aux, _GradientDescentState]):
     norm: AbstractVar[Callable[[PyTree], Scalar]]
     descent: AbstractVar[AbstractDescent[Y, FunctionInfo.EvalGrad, Any]]
     search: AbstractVar[
-        AbstractSearch[Y, FunctionInfo.EvalGrad, FunctionInfo.Eval, Any]
+        AbstractSearch[
+            Y, FunctionInfo.EvalGrad, FunctionInfo.Eval | FunctionInfo.EvalGrad, Any
+        ]
     ]
 
     def init(
@@ -171,20 +173,30 @@ def step(
         f_eval, lin_fn, aux_eval = jax.linearize(
             lambda _y: fn(_y, args), state.y_eval, has_aux=True
         )
+        f_eval_info = make_eval_info_for_search(
+            self.search.info_needed_at_y_eval,
+            state.y_eval,
+            f_eval,
+            lin_fn,
+            autodiff_mode,
+        )
+
         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,
             state.search_state,
         )
 
         def accepted(descent_state):
-            grad = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype)
-
-            f_eval_info = FunctionInfo.EvalGrad(f_eval, grad)
-            descent_state = self.descent.query(state.y_eval, f_eval_info, descent_state)
+            accepted_f_eval_info = promote_to_eval_grad(
+                f_eval_info, state.y_eval, f_eval, lin_fn, autodiff_mode
+            )
+            descent_state = self.descent.query(
+                state.y_eval, accepted_f_eval_info, descent_state
+            )
             y_diff = (state.y_eval**ω - y**ω).ω
             f_diff = (f_eval**ω - state.f_info.f**ω).ω
             terminate = cauchy_termination(
@@ -193,7 +205,13 @@ def accepted(descent_state):
             terminate = jnp.where(
                 state.first_step, jnp.array(False), terminate
             )  # Skip termination on first step
-            return state.y_eval, f_eval_info, aux_eval, descent_state, terminate
+            return (
+                state.y_eval,
+                accepted_f_eval_info,
+                aux_eval,
+                descent_state,
+                terminate,
+            )
 
         def rejected(descent_state):
             return y, state.f_info, state.aux, descent_state, jnp.array(False)

optimistix/_solver/learning_rate.py

@@ -1,4 +1,4 @@
-from typing import cast
+from typing import cast, ClassVar
 
 import equinox as eqx
 import jax.numpy as jnp
@@ -16,6 +16,7 @@ def _typed_asarray(x: ScalarLike) -> Array:
 class LearningRate(AbstractSearch[Y, FunctionInfo, FunctionInfo, None]):
     """Move downhill by taking a step of the fixed size `learning_rate`."""
 
+    info_needed_at_y_eval: ClassVar[type[FunctionInfo]] = FunctionInfo
     learning_rate: ScalarLike = eqx.field(converter=_typed_asarray)
 
     def init(self, y: Y, f_info_struct: FunctionInfo) -> None:

optimistix/_solver/nonlinear_cg.py

@@ -183,7 +183,9 @@ class NonlinearCG(AbstractGradientDescent[Y, Aux]):
     atol: float
     norm: Callable[[PyTree], Scalar]
     descent: NonlinearCGDescent[Y]
-    search: AbstractSearch[Y, FunctionInfo.EvalGrad, FunctionInfo.Eval, Any]
+    search: AbstractSearch[
+        Y, FunctionInfo.EvalGrad, FunctionInfo.Eval | FunctionInfo.EvalGrad, Any
+    ]
 
     def __init__(
         self,
@@ -193,7 +195,10 @@ def __init__(
         method: Callable[[Y, Y, Y], Scalar] = polak_ribiere,
         # TODO(raderj): replace the default line search with something better.
         search: AbstractSearch[
-            Y, FunctionInfo.EvalGrad, FunctionInfo.Eval, Any
+            Y,
+            FunctionInfo.EvalGrad,
+            FunctionInfo.Eval | FunctionInfo.EvalGrad,
+            Any,
         ] = BacktrackingArmijo(decrease_factor=0.5, slope=0.1),
     ):
         """**Arguments:**

optimistix/_solver/quasi_newton.py

@@ -17,7 +17,6 @@
     cauchy_termination,
     default_verbose,
     filter_cond,
-    lin_to_grad,
     max_norm,
     tree_dot,
     tree_full_like,
@@ -29,6 +28,7 @@
     FunctionInfo,
 )
 from .._solution import RESULTS
+from ._f_eval_info_helpers import make_eval_info_for_search, promote_to_eval_grad
 from .backtracking import BacktrackingArmijo
 from .gauss_newton import NewtonDescent
 
@@ -139,7 +139,9 @@ class AbstractQuasiNewton(
     norm: AbstractVar[Callable[[PyTree], Scalar]]
     use_inverse: AbstractVar[bool]
     descent: AbstractVar[AbstractDescent[Y, _Hessian, Any]]
-    search: AbstractVar[AbstractSearch[Y, _Hessian, FunctionInfo.Eval, Any]]
+    search: AbstractVar[
+        AbstractSearch[Y, _Hessian, FunctionInfo.Eval | FunctionInfo.EvalGrad, Any]
+    ]
     verbose: AbstractVar[Callable[..., None]]
 
     @abc.abstractmethod
@@ -214,29 +216,39 @@ def step(
         f_eval, lin_fn, aux_eval = jax.linearize(
             lambda _y: fn(_y, args), state.y_eval, has_aux=True
         )
+        f_eval_info = make_eval_info_for_search(
+            self.search.info_needed_at_y_eval,
+            state.y_eval,
+            f_eval,
+            lin_fn,
+            autodiff_mode,
+        )
+
         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,
             state.search_state,
         )
 
         def accepted(descent_state):
-            grad = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype)
+            accepted_f_eval_info = promote_to_eval_grad(
+                f_eval_info, state.y_eval, f_eval, lin_fn, autodiff_mode
+            )
 
-            f_eval_info, hessian_update_state = self.update_hessian(
+            updated_f_eval_info, hessian_update_state = self.update_hessian(
                 y,
                 state.y_eval,
                 state.f_info,
-                FunctionInfo.EvalGrad(f_eval, grad),
+                accepted_f_eval_info,
                 state.hessian_update_state,
             )
 
             descent_state = self.descent.query(
                 state.y_eval,
-                f_eval_info,  # pyright: ignore
+                updated_f_eval_info,
                 descent_state,
             )
             y_diff = (state.y_eval**ω - y**ω).ω
@@ -249,7 +261,7 @@ def accepted(descent_state):
             )  # Skip termination on first step
             return (
                 state.y_eval,
-                f_eval_info,
+                updated_f_eval_info,
                 aux_eval,
                 descent_state,
                 terminate,

optimistix/_solver/trust_region.py

@@ -1,5 +1,5 @@
 import abc
-from typing import TypeAlias, TypeVar
+from typing import ClassVar, TypeAlias, TypeVar
 
 import equinox as eqx
 import jax.numpy as jnp
@@ -47,6 +47,8 @@ class _AbstractTrustRegion(AbstractSearch[Y, _FnInfo, _FnEvalInfo, _TrustRegionS
     high_constant: AbstractVar[ScalarLike]
     low_constant: AbstractVar[ScalarLike]
 
+    info_needed_at_y_eval: ClassVar[type[FunctionInfo]] = FunctionInfo
+
     def __post_init__(self):
         # You would not expect `self.low_cutoff` or `self.high_cutoff` to
         # be below zero, but this is technically not incorrect so we don't