Implement scaling update as maximum hessian diagonal encountered so far

Author: michael-0brien

optimistix/_solver/levenberg_marquardt.py

@@ -3,6 +3,7 @@
 
 import equinox as eqx
 import jax
+import jax.flatten_util as jfu
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.tree_util as jtu
@@ -21,8 +22,7 @@
 
 
 UpdateScalingFn: TypeAlias = Callable[
-    [FunctionInfo.EvalGradHessian | FunctionInfo.ResidualJac, PyTree],
-    lx.AbstractLinearOperator,
+    [lx.AbstractLinearOperator, lx.DiagonalLinearOperator], lx.DiagonalLinearOperator
 ]
 
 
@@ -38,26 +38,20 @@ def __call__(self, y: PyTree[Array]):
 
 
 def max_diagonal_scaling_update(
-    f_info: FunctionInfo.EvalGradHessian | FunctionInfo.ResidualJac,
-    scaling_operator: lx.AbstractLinearOperator,
-) -> lx.AbstractLinearOperator:
+    hessian: lx.AbstractLinearOperator, scaling_operator: lx.DiagonalLinearOperator
+) -> lx.DiagonalLinearOperator:
     """Update the matrix `D` that controls the relative scaling of each
-    parameter based on the procedure described by More (1977). This takes `D`
-    on each iteration as the maximum diagonal of the hessian
-    so far encountered.
+    parameter based on the procedure described by More (1977).
+    This takes `D` on each iteration as the maximum diagonal of the
+    hessian so far encountered.
     """  # noqa: E501
-    raise NotImplementedError
-    # assert lx.is_diagonal(scaling_operator)
-    # if isinstance(f_info, FunctionInfo.EvalGradHessian):
-    #     hessian = f_info.hessian
-    # else:
-    #     hessian = f_info.jac.transpose() @ f_info.jac
-    # hessian_diagonal = lx.TaggedLinearOperator(
-    #     hessian @ lx.IdentityLinearOperator(hessian.in_structure()), lx.diagonal_tag
-    # )
-
-
-def step_levenberg_damping(
+    diagonal, unflatten_fn = jfu.ravel_pytree(scaling_operator.diagonal)
+    return lx.DiagonalLinearOperator(
+        unflatten_fn(jnp.maximum(lx.diagonal(hessian), diagonal))
+    )
+
+
+def damped_newton_step_levenberg(
     step_size: Scalar,
     f_info: FunctionInfo.EvalGradHessian | FunctionInfo.ResidualJac,
     linear_solver: lx.AbstractLinearSolver,
@@ -104,11 +98,11 @@ def step_levenberg_damping(
     return linear_sol.value, RESULTS.promote(linear_sol.result)
 
 
-def step_scaling_operator_damping(
+def damped_newton_step_scaled(
     step_size: Scalar,
     f_info: FunctionInfo.EvalGradHessian | FunctionInfo.ResidualJac,
-    scaling_operator: lx.AbstractLinearOperator,
     linear_solver: lx.AbstractLinearSolver,
+    scaling_operator: lx.DiagonalLinearOperator,
     update_scaling_fn: UpdateScalingFn,
 ) -> tuple[PyTree[Array], RESULTS]:
     """Compute a damped Newton step that keeps track of an operator
@@ -126,14 +120,14 @@ def step_scaling_operator_damping(
         grad = f_info.grad
     elif isinstance(f_info, FunctionInfo.ResidualJac):
         jac, residual = f_info.jac, f_info.residual
-        hessian = jac.transpose() @ jac
-        grad = jac.transpose() @ residual
+        hessian = lx.TaggedLinearOperator(jac.T @ jac, lx.positive_semidefinite_tag)
+        grad = jac.T.mv(residual)
     else:
         raise ValueError(
             "Damped newton descent cannot be used with a solver that does not "
             "provide (approximate) Hessian information."
         )
-    scaling_operator = update_scaling_fn(f_info, scaling_operator)
+    scaling_operator = update_scaling_fn(hessian, scaling_operator)
     operator = hessian + lm_param * scaling_operator
     if lx.is_positive_semidefinite(hessian):
         operator = lx.TaggedLinearOperator(operator, lx.positive_semidefinite_tag)
@@ -191,7 +185,7 @@ def query(
     def step(
         self, step_size: Scalar, state: _DampedNewtonDescentState
     ) -> tuple[Y, RESULTS]:
-        sol_value, result = step_levenberg_damping(
+        sol_value, result = damped_newton_step_levenberg(
             step_size, state.f_info, self.linear_solver
         )
         y_diff = (-(sol_value**ω)).ω
@@ -281,7 +275,7 @@ def step(
         scaled_step_size = state.newton_norm * step_size
 
         def comparison_fn(lambda_i: Scalar, _):
-            step, _ = step_levenberg_damping(
+            step, _ = damped_newton_step_levenberg(
                 1 / lambda_i, state.f_info, self.linear_solver
             )
             return self.trust_region_norm(step) - scaled_step_size
@@ -296,7 +290,7 @@ def reject_newton():
                 max_steps=32,
                 throw=False,
             ).value
-            y_diff, result = step_levenberg_damping(
+            y_diff, result = damped_newton_step_levenberg(
                 1 / lambda_out, state.f_info, self.linear_solver
             )
             return y_diff, result
@@ -325,7 +319,7 @@ def accept_newton():
 
 class _ScaledDampedNewtonDescentState(eqx.Module):
     f_info: FunctionInfo.EvalGradHessian | FunctionInfo.ResidualJac
-    scaling_operator: lx.AbstractLinearOperator
+    scaling_operator: lx.DiagonalLinearOperator
 
 
 class ScaledDampedNewtonDescent(
@@ -366,12 +360,9 @@ def init(
         y: Y,
         f_info_struct: FunctionInfo.EvalGradHessian | FunctionInfo.ResidualJac,
     ) -> _ScaledDampedNewtonDescentState:
-        if isinstance(f_info_struct, FunctionInfo.ResidualJac):
-            raise NotImplementedError
-        del y
         f_info_init = tree_full_like(f_info_struct, 0, allow_static=True)
-        scaling_operator = -jnp.inf * lx.IdentityLinearOperator(
-            f_info_struct.hessian.in_structure()
+        scaling_operator = lx.DiagonalLinearOperator(
+            tree_full_like(y, -jnp.inf, allow_static=True)
         )
         return _ScaledDampedNewtonDescentState(f_info_init, scaling_operator)
 
@@ -387,11 +378,11 @@ def query(
     def step(
         self, step_size: Scalar, state: _ScaledDampedNewtonDescentState
     ) -> tuple[Y, RESULTS]:
-        sol_value, result = step_scaling_operator_damping(
+        sol_value, result = damped_newton_step_scaled(
             step_size,
             state.f_info,
-            state.scaling_operator,
             self.linear_solver,
+            state.scaling_operator,
             self.update_scaling_fn,
         )
         y_diff = (-(sol_value**ω)).ω
@@ -401,10 +392,11 @@ def step(
 ScaledDampedNewtonDescent.__init__.__doc__ = """**Arguments:**
 
 - `linear_solver`: The linear solver used to compute the Newton step.
-- `update_scaling_fn`: A function with signature `fn(f_info, descent_state)`
-    that returns a `lineax.AbstractLinearOperator` of the same structure as
-    `f_info.hessian.in_structure()`. By default, this is the procedure used
-    in [`scipy.least_squares`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.least_squares.html)
+- `update_scaling_fn`: A function with signature `fn(hessian, scaling)`,
+    where `hessian` is an `AbstractLinearOperator` and `scaling` is a
+    `DiagonalLinearOperator` for the scaling operator on the kth iteration.
+    By default, this is the procedure used in [`scipy.least_squares`]
+    (https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.least_squares.html)
 """  # noqa: E501