Getting proximal to work with augmented Lagrangian optimizers

desc/objectives/utils.py

@@ -80,13 +80,14 @@ def factorize_linear_constraints(objective, constraint, x_scale="auto"):  # noqa
 
     if isinstance(objective, ProximalProjection):
         # remove cols of A corresponding to ["R_lmn", "Z_lmn", "L_lmn", "Ra_n", "Za_n"]
-        # see desc.optimize._constraint_wrappers.ProximalProjection._set_eq_state_vector
+        # see desc.optimize._constraint_wrappers.ProximalState._set_eq_state_vector
         c = 0
         cols = np.array([], dtype=int)
+        eq = objective._state.eq
         for t in objective.things:
-            if t is objective._eq:
-                for arg, dim in objective._eq.dimensions.items():
-                    if arg in objective._args:  # these Equilibrium args are kept
+            if t is eq:
+                for arg, dim in eq.dimensions.items():
+                    if arg in objective._state.args:  # these Equilibrium args are kept
                         cols = np.append(cols, np.arange(c, c + dim))
                     c += dim  # other Equilibrium args are removed
             else:  # non-Equilibrium args are always included

desc/optimize/__init__.py

@@ -1,7 +1,11 @@
 """Functions for minimization and wrappers for scipy methods."""
 
 from . import _desc_wrappers, _scipy_wrappers
-from ._constraint_wrappers import LinearConstraintProjection, ProximalProjection
+from ._constraint_wrappers import (
+    LinearConstraintProjection,
+    ProximalProjection,
+    ProximalState,
+)
 from .aug_lagrangian import fmin_auglag
 from .aug_lagrangian_ls import lsq_auglag
 from .fmin_scalar import fmintr

desc/optimize/aug_lagrangian_ls.py

@@ -25,7 +25,7 @@
     inequality_to_bounds,
     print_header_nonlinear,
     print_iteration_nonlinear,
-    scale_columns,
+    scale_matrix,
     solve_triangular_regularized,
 )
 
@@ -288,7 +288,7 @@ def lagjac(z, y, mu, *args):
     # we don't need unscaled J anymore, so we overwrite it with J_h = J * d to avoid
     # carrying so many J-sized matrices in memory, which can be large. The buffer is
     # donated so the scaling doesn't allocate a second copy of J.
-    J_h = scale_columns(J, d)
+    J_h = scale_matrix(J, d)
     del J
     g_norm = jnp.linalg.norm(
         (g * v * scale if scaled_termination else g * v), ord=jnp.inf
@@ -546,6 +546,7 @@ def lagjac(z, y, mu, *args):
 
             # updating augmented lagrangian params
             if g_norm < gtolk:
+                mu_old = mu
                 y = jnp.where(jnp.abs(c) < ctolk, y - mu * c, y)
                 mu = jnp.where(jnp.abs(c) >= ctolk, tau * mu, mu)
                 if constr_violation < ctolk:
@@ -554,12 +555,14 @@ def lagjac(z, y, mu, *args):
                 else:
                     ctolk = max(eta / (jnp.mean(mu) ** alpha_eta), ctol)
                     gtolk = max(omega / (jnp.mean(mu) ** alpha_omega), gtol)
-                # if we update lagrangian params, need to recompute L and J
+                # if we update lagrangian params, need to recompute L and J.
                 L = lagfun(f, c, y, mu)
                 Lcost = 0.5 * jnp.dot(L, L)
-                del J
-                J = lagjac(z, y, mu, *args)
-                njev += 1
+                row_scale = jnp.concatenate(
+                    [jnp.ones_like(f), jnp.sqrt(mu) / jnp.sqrt(mu_old)]
+                )
+                J = scale_matrix(J, row_scale[:, None])
+
                 g = jnp.dot(L, J)
 
                 if jac_scale:
@@ -577,7 +580,7 @@ def lagjac(z, y, mu, *args):
             d = v**0.5 * scale
             diag_h = g * dv * scale
             g_h = g * d
-            J_h = scale_columns(J, d)
+            J_h = scale_matrix(J, d)
             del J
 
             if g_norm < gtol and constr_violation < ctol:
@@ -611,7 +614,7 @@ def lagjac(z, y, mu, *args):
     active_mask = find_active_constraints(z, zbounds[0], zbounds[1], rtol=xtol)
     # after overwriting J_h with J*d, we have to revert back and store the
     # unscaled version
-    J_h = scale_columns(J_h, 1 / d)
+    J_h = scale_matrix(J_h, 1 / d)
     result = OptimizeResult(
         x=x,
         s=s,

desc/optimize/fmin_scalar.py

@@ -24,7 +24,7 @@
     compute_hess_scale,
     print_header_nonlinear,
     print_iteration_nonlinear,
-    scale_columns,
+    scale_matrix,
 )
 
 
@@ -243,9 +243,9 @@ def fmintr(  # noqa: C901
     # we don't need unscaled H anymore, so we overwrite it with H_h = d[:, None] * H * d
     # to avoid carrying so many H-sized matrices in memory, which can be large. The
     # buffer is donated so the scaling doesn't allocate a second copy of H.
-    H_h = scale_columns(H, d)
+    H_h = scale_matrix(H, d)
     del H
-    H_h = scale_columns(H_h, d[:, None])
+    H_h = scale_matrix(H_h, d[:, None])
 
     g_norm = jnp.linalg.norm(
         (g * v * scale if scaled_termination else g * v), ord=jnp.inf
@@ -440,9 +440,9 @@ def fmintr(  # noqa: C901
 
             g_h = g * d
 
-            H_h = scale_columns(H, d)
+            H_h = scale_matrix(H, d)
             del H
-            H_h = scale_columns(H_h, d[:, None])
+            H_h = scale_matrix(H_h, d[:, None])
 
             x_norm = jnp.linalg.norm(
                 ((x * scale_inv) if scaled_termination else x), ord=2
@@ -474,8 +474,8 @@ def fmintr(  # noqa: C901
     active_mask = find_active_constraints(x, lb, ub, rtol=xtol)
     # after overwriting H_h with the scaled version, we have to revert back and
     # store the unscaled one
-    H_h = scale_columns(H_h, 1 / d)
-    H_h = scale_columns(H_h, 1 / d[:, None])
+    H_h = scale_matrix(H_h, 1 / d)
+    H_h = scale_matrix(H_h, 1 / d[:, None])
     result = OptimizeResult(
         x=x,
         success=success,

desc/optimize/least_squares.py

@@ -24,7 +24,7 @@
     compute_jac_scale,
     print_header_nonlinear,
     print_iteration_nonlinear,
-    scale_columns,
+    scale_matrix,
     solve_triangular_regularized,
 )
 
@@ -209,7 +209,7 @@ def lsqtr(  # noqa: C901
     # we don't need unscaled J anymore, so we overwrite it with J_h = J * d to avoid
     # carrying so many J-sized matrices in memory, which can be large. The buffer is
     # donated so the scaling doesn't allocate a second copy of J.
-    J_h = scale_columns(J, d)
+    J_h = scale_matrix(J, d)
     del J
     g_norm = jnp.linalg.norm(
         (g * v * scale if scaled_termination else g * v), ord=jnp.inf
@@ -427,7 +427,7 @@ def lsqtr(  # noqa: C901
             diag_h = g * dv * scale
 
             g_h = g * d
-            J_h = scale_columns(J, d)
+            J_h = scale_matrix(J, d)
             del J
             x_norm = jnp.linalg.norm(
                 ((x * scale_inv) if scaled_termination else x), ord=2
@@ -458,7 +458,7 @@ def lsqtr(  # noqa: C901
     active_mask = find_active_constraints(x, lb, ub, rtol=xtol)
     # after overwriting J_h with J*d, we have to revert back and store the
     # unscaled version
-    J_h = scale_columns(J_h, 1 / d)
+    J_h = scale_matrix(J_h, 1 / d)
     result = OptimizeResult(
         x=x,
         success=success,

desc/optimize/optimizer.py

@@ -28,7 +28,11 @@
     warnif,
 )
 
-from ._constraint_wrappers import LinearConstraintProjection, ProximalProjection
+from ._constraint_wrappers import (
+    LinearConstraintProjection,
+    ProximalProjection,
+    ProximalState,
+)
 
 
 class Optimizer(IOAble):
@@ -335,7 +339,7 @@
             # reset eq params to initial
             if eq is not None:
                 eq.params_dict = eq_params_init
-            result["history"] = objective.history
+            result["history"] = objective.history(result["allx"])
             objective = objective._objective
         else:
             result["history"] = [
@@ -460,10 +464,11 @@
         # Project equilibrium part: remove excluded parameters
         excluded_params = ["R_lmn", "Z_lmn", "L_lmn", "Ra_n", "Za_n"]
         included_idx = []
-        for arg in prox_obj._eq.optimizable_params:
+        for arg in prox_obj._state.eq.optimizable_params:
             if arg not in excluded_params:
-                included_idx.extend(prox_obj._eq.x_idx[arg])
-        x_scale[prox_obj._eq_idx] = x_scale[prox_obj._eq_idx][jnp.array(included_idx)]
+                included_idx.extend(prox_obj._state.eq.x_idx[arg])
+        eq_idx = prox_obj._eq_idx
+        x_scale[eq_idx] = x_scale[eq_idx][jnp.array(included_idx)]
         x_scale = jnp.concatenate(x_scale)
 
     if isinstance(objective, LinearConstraintProjection):
@@ -530,6 +535,8 @@
         Otherwise returns None.
 
     """
+    if len(constraints) == 1 and isinstance(constraints[0], ObjectiveFunction):
+        return constraints[0]
     if len(constraints):
         objective = ObjectiveFunction(constraints)
     else:
@@ -576,6 +583,32 @@
     return linear_constraints, nonlinear_constraints
 
 
+def _parse_nonlinear_constraints(eq, nonlinear_constraints):
+    """Split nonlinear constraints for Proximal.
+
+    Parameters
+    ----------
+    eq: Equilibrium
+    nonlinear constraints : tuple of Objective
+        Nonlinear constraints to parse.
+
+    Returns
+    -------
+    eq_constraints : tuple of Objective
+        Constraints which can be handled directly by Proximal,
+        i.e. ones involving the equilibrium and without bounds.
+    other_constraints : tuple of Objective
+        Remaining nonlinear constraints. Can be empty.
+    """
+    eq_constraints, other_constraints = [], []
+    for con in nonlinear_constraints:
+        if con._equilibrium and con.bounds is None and con.things == [eq]:
+            eq_constraints.append(con)
+        else:
+            other_constraints.append(con)
+    return eq_constraints, other_constraints
+
+
 def _maybe_wrap_nonlinear_constraints(
     eq, objective, nonlinear_constraints, method, options
 ):
@@ -599,16 +632,38 @@
                 """))
         wrapper = "proximal"
     if wrapper is not None and wrapper.lower() in ["prox", "proximal"]:
+        eq_constraints, other_constraints = _parse_nonlinear_constraints(
+            eq, nonlinear_constraints
+        )
+        errorif(
+            not len(eq_constraints),
+            ValueError,
+            f"Method {wrapper}-{method} requires at least one equilibrium constraint "
+            + "(e.g. ForceBalance), but none were given.",
+        )
+        if len(other_constraints) and not optimizers[method]["equality_constraints"]:
+            raise ValueError(
+                f"Method {wrapper}-{method} cannot accept nonlinear constraints. "
+                + "Consider an Augmented Lagrangian solver such as proximal-lsq-auglag."
+            )
+
         perturb_options = options.pop("perturb_options", {})
         solve_options = options.pop("solve_options", {})
-        objective = ProximalProjection(
-            objective,
-            constraint=_combine_constraints(nonlinear_constraints),
+        state = ProximalState(
+            eq,
+            _combine_constraints(eq_constraints),
             perturb_options=perturb_options,
             solve_options=solve_options,
-            eq=eq,
+            cache_tangents=bool(other_constraints),
         )
-        nonlinear_constraints = ()
+
+        objective = ProximalProjection(objective, state=state)
+        nonlinear_constraints = (
+            (ProximalProjection(_combine_constraints(other_constraints), state=state),)
+            if len(other_constraints)
+            else ()
+        )
+
     return objective, nonlinear_constraints
 
 
@@ -718,7 +773,7 @@
         objective.build(verbose=verbose)
         if nonlinear_constraint is not None:
             nonlinear_constraint = LinearConstraintProjection(
-                nonlinear_constraint, linear_constraint
+                nonlinear_constraint, linear_constraint, **linear_constraint_options
             )
             nonlinear_constraint.build(verbose=verbose)
 

desc/optimize/utils.py

@@ -504,11 +504,11 @@ def compute_jac_scale(A, prev_scale_inv=None):
 
 
 @functools.partial(jit, donate_argnums=0)
-def scale_columns(A, d):
+def scale_matrix(A, d):
     """Compute `A * d` reusing `A`'s buffer instead of allocating a second copy.
 
     `A` is invalid after this call, so callers must rebind, ie
-    `A = scale_columns(A, d)`.
+    `A = scale_matrix(A, d)`.
     """
     return A * d