Merge pull request #503 from PlasmaControl/constrained2

Adding augmented lagrangian for constrained optimization

desc/equilibrium/equilibrium.py

@@ -548,6 +548,7 @@ def optimize(
         ftol=None,
         xtol=None,
         gtol=None,
+        ctol=None,
         maxiter=None,
         x_scale="auto",
         options=None,
@@ -564,7 +565,7 @@ def optimize(
             Objective function to satisfy. Default = fixed-boundary force balance.
         optimizer : str or Optimizer (optional)
             optimizer to use
-        ftol, xtol, gtol : float
+        ftol, xtol, gtol, ctol : float
             stopping tolerances. `None` will use defaults for given optimizer.
         maxiter : int
             Maximum number of solver steps.
@@ -618,6 +619,7 @@ def optimize(
             ftol=ftol,
             xtol=xtol,
             gtol=gtol,
+            ctol=ctol,
             x_scale=x_scale,
             verbose=verbose,
             maxiter=maxiter,

desc/optimize/__init__.py

@@ -2,6 +2,8 @@
 
 from . import _desc_wrappers, _scipy_wrappers
 from ._constraint_wrappers import LinearConstraintProjection, ProximalProjection
+from .aug_lagrangian import fmin_auglag
+from .aug_lagrangian_ls import lsq_auglag
 from .fmin_scalar import fmintr
 from .least_squares import lsqtr
 from .optimizer import Optimizer, optimizers, register_optimizer

desc/optimize/_desc_wrappers.py

@@ -1,11 +1,200 @@
+from scipy.optimize import NonlinearConstraint
+
 from desc.backend import jnp
 
+from .aug_lagrangian import fmin_auglag
+from .aug_lagrangian_ls import lsq_auglag
 from .fmin_scalar import fmintr
 from .least_squares import lsqtr
 from .optimizer import register_optimizer
 from .stochastic import sgd
 
 
+@register_optimizer(
+    name="fmin-auglag",
+    description="Augmented Lagrangian method with trust region subproblem.",
+    scalar=True,
+    equality_constraints=True,
+    inequality_constraints=True,
+    stochastic=False,
+    hessian=True,
+    GPU=True,
+)
+def _optimize_desc_aug_lagrangian(
+    objective, constraint, x0, method, x_scale, verbose, stoptol, options=None
+):
+    """Wrapper for desc.optimize.fmin_lag_ls_stel.
+
+    Parameters
+    ----------
+    objective : ObjectiveFunction
+        Function to minimize.
+    constraint : ObjectiveFunction
+        Constraint to satisfy
+    x0 : ndarray
+        Starting point.
+    method : str
+        Name of the method to use. Any of the options in scipy.optimize.least_squares.
+    x_scale : array_like or ‘jac’, optional
+        Characteristic scale of each variable. Setting x_scale is equivalent to
+        reformulating the problem in scaled variables xs = x / x_scale. An alternative
+        view is that the size of a trust region along jth dimension is proportional to
+        x_scale[j]. Improved convergence may be achieved by setting x_scale such that
+        a step of a given size along any of the scaled variables has a similar effect
+        on the cost function. If set to ‘jac’, the scale is iteratively updated using
+        the inverse norms of the columns of the Jacobian matrix.
+    verbose : int
+        * 0  : work silently.
+        * 1 : display a termination report.
+        * 2 : display progress during iterations
+    stoptol : dict
+        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol", "ctol",
+        "maxiter", "max_nfev", "max_njev", "max_ngev", "max_nhev"}
+    options : dict, optional
+        Dictionary of optional keyword arguments to override default solver
+        settings. See the code for more details.
+
+    Returns
+    -------
+    res : OptimizeResult
+       The optimization result represented as a ``OptimizeResult`` object.
+       Important attributes are: ``x`` the solution array, ``success`` a
+       Boolean flag indicating if the optimizer exited successfully and
+       ``message`` which describes the cause of the termination. See
+       `OptimizeResult` for a description of other attributes.
+
+    """
+    options = {} if options is None else options
+    if not isinstance(x_scale, str) and jnp.allclose(x_scale, 1):
+        options.setdefault("initial_trust_ratio", 1e-3)
+        options.setdefault("max_trust_radius", 1.0)
+    options["max_nfev"] = stoptol["max_nfev"]
+    options["max_ngev"] = stoptol["max_ngev"]
+    options["max_nhev"] = stoptol["max_nhev"]
+
+    if constraint is not None:
+        lb, ub = constraint.bounds_scaled
+        constraint_wrapped = NonlinearConstraint(
+            constraint.compute_scaled,
+            lb,
+            ub,
+            constraint.jac_scaled,
+        )
+    else:
+        constraint_wrapped = None
+
+    result = fmin_auglag(
+        objective.compute_scalar,
+        x0=x0,
+        grad=objective.grad,
+        hess=objective.hess,
+        bounds=(-jnp.inf, jnp.inf),
+        constraint=constraint_wrapped,
+        args=(),
+        x_scale=x_scale,
+        ftol=stoptol["ftol"],
+        xtol=stoptol["xtol"],
+        gtol=stoptol["gtol"],
+        ctol=stoptol["ctol"],
+        verbose=verbose,
+        maxiter=stoptol["maxiter"],
+        options=options,
+    )
+    return result
+
+
+@register_optimizer(
+    name="lsq-auglag",
+    description="Least Squares Augmented Lagrangian approach "
+    + "to constrained optimization",
+    scalar=False,
+    equality_constraints=True,
+    inequality_constraints=True,
+    stochastic=False,
+    hessian=False,
+    GPU=True,
+)
+def _optimize_desc_aug_lagrangian_least_squares(
+    objective, constraint, x0, method, x_scale, verbose, stoptol, options=None
+):
+    """Wrapper for desc.optimize.fmin_lag_ls_stel.
+
+    Parameters
+    ----------
+    objective : ObjectiveFunction
+        Function to minimize.
+    constraint : ObjectiveFunction
+        Constraint to satisfy
+    x0 : ndarray
+        Starting point.
+    method : str
+        Name of the method to use. Any of the options in scipy.optimize.least_squares.
+    x_scale : array_like or ‘jac’, optional
+        Characteristic scale of each variable. Setting x_scale is equivalent to
+        reformulating the problem in scaled variables xs = x / x_scale. An alternative
+        view is that the size of a trust region along jth dimension is proportional to
+        x_scale[j]. Improved convergence may be achieved by setting x_scale such that
+        a step of a given size along any of the scaled variables has a similar effect
+        on the cost function. If set to ‘jac’, the scale is iteratively updated using
+        the inverse norms of the columns of the Jacobian matrix.
+    verbose : int
+        * 0  : work silently.
+        * 1 : display a termination report.
+        * 2 : display progress during iterations
+    stoptol : dict
+        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol", "ctol",
+        "maxiter", "max_nfev", "max_njev", "max_ngev", "max_nhev"}
+    options : dict, optional
+        Dictionary of optional keyword arguments to override default solver
+        settings. See the code for more details.
+
+    Returns
+    -------
+    res : OptimizeResult
+       The optimization result represented as a ``OptimizeResult`` object.
+       Important attributes are: ``x`` the solution array, ``success`` a
+       Boolean flag indicating if the optimizer exited successfully and
+       ``message`` which describes the cause of the termination. See
+       `OptimizeResult` for a description of other attributes.
+
+    """
+    options = {} if options is None else options
+    if not isinstance(x_scale, str) and jnp.allclose(x_scale, 1):
+        options.setdefault("initial_trust_radius", 1e-3)
+        options.setdefault("max_trust_radius", 1.0)
+    options["max_nfev"] = stoptol["max_nfev"]
+    options["max_njev"] = stoptol["max_njev"]
+
+    if constraint is not None:
+        lb, ub = constraint.bounds_scaled
+        constraint_wrapped = NonlinearConstraint(
+            constraint.compute_scaled,
+            lb,
+            ub,
+            constraint.jac_scaled,
+        )
+    else:
+        constraint_wrapped = None
+
+    result = lsq_auglag(
+        objective.compute_scaled_error,
+        x0=x0,
+        jac=objective.jac_scaled,
+        bounds=(-jnp.inf, jnp.inf),
+        constraint=constraint_wrapped,
+        args=(),
+        x_scale=x_scale,
+        ftol=stoptol["ftol"],
+        xtol=stoptol["xtol"],
+        gtol=stoptol["gtol"],
+        ctol=stoptol["ctol"],
+        verbose=verbose,
+        maxiter=stoptol["maxiter"],
+        options=options,
+    )
+    return result
+
+
 @register_optimizer(
     name="lsq-exact",
     description="Trust region least squares method, "
@@ -45,7 +234,7 @@ def _optimize_desc_least_squares(
         * 1 : display a termination report.
         * 2 : display progress during iterations
     stoptol : dict
-        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol",
+        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol", "ctol",
         "maxiter", "max_nfev", "max_njev", "max_ngev", "max_nhev"}
     options : dict, optional
         Dictionary of optional keyword arguments to override default solver
@@ -89,12 +278,23 @@ def _optimize_desc_least_squares(
 
 
 @register_optimizer(
-    name=["dogleg", "subspace", "dogleg-bfgs", "subspace-bfgs"],
+    name=[
+        "fmin-exact",
+        "fmin-dogleg",
+        "fmin-subspace",
+        "fmin-exact-bfgs",
+        "fmin-dogleg-bfgs",
+        "fmin-subspace-bfgs",
+    ],
     description=[
+        "Trust region method using iterative cholesky method to exactly solve the "
+        + "trust region subproblem.",
         "Trust region method using Powell's dogleg method to approximately solve the "
         + "trust region subproblem.",
         "Trust region method solving the subproblem over the 2d subspace spanned by "
         + "the gradient and newton direction.",
+        "Trust region method using iterative cholesky method to exactly solve the "
+        + "trust region subproblem. Uses BFGS to approximate hessian",
         "Trust region method using Powell's dogleg method to approximately solve the "
         + "trust region subproblem. Uses BFGS to approximate hessian",
         "Trust region method solving the subproblem over the 2d subspace spanned by "
@@ -104,7 +304,7 @@ def _optimize_desc_least_squares(
     equality_constraints=False,
     inequality_constraints=False,
     stochastic=False,
-    hessian=[True, True, False, False],
+    hessian=[True, True, True, False, False, False],
     GPU=True,
 )
 def _optimize_desc_fmin_scalar(
@@ -135,7 +335,7 @@ def _optimize_desc_fmin_scalar(
         * 1 : display a termination report.
         * 2 : display progress during iterations
     stoptol : dict
-        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol",
+        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol", "ctol",
         "maxiter", "max_nfev", "max_njev", "max_ngev", "max_nhev"}
     options : dict, optional
         Dictionary of optional keyword arguments to override default solver
@@ -169,7 +369,7 @@ def _optimize_desc_fmin_scalar(
         grad=objective.grad,
         hess=hess,
         args=(),
-        method=method.replace("-bfgs", ""),
+        method=method.replace("-bfgs", "").replace("fmin-", ""),
         x_scale=x_scale,
         ftol=stoptol["ftol"],
         xtol=stoptol["xtol"],
@@ -220,7 +420,7 @@ def _optimize_desc_stochastic(
         * 1 : display a termination report.
         * 2 : display progress during iterations
     stoptol : dict
-        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol",
+        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol", "ctol",
         "maxiter", "max_nfev", "max_njev", "max_ngev", "max_nhev"}
     options : dict, optional
         Dictionary of optional keyword arguments to override default solver

desc/optimize/_scipy_wrappers.py

@@ -82,7 +82,7 @@ def _optimize_scipy_minimize(  # noqa: C901 - FIXME: simplify this
         * 1 : display a termination report.
         * 2 : display progress during iterations
     stoptol : dict
-        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol",
+        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol", "ctol",
         "maxiter", "max_nfev", "max_njev", "max_ngev", "max_nhev"}
     options : dict, optional
         Dictionary of optional keyword arguments to override default solver
@@ -341,7 +341,7 @@ def _optimize_scipy_least_squares(  # noqa: C901 - FIXME: simplify this
         * 1 : display a termination report.
         * 2 : display progress during iterations
     stoptol : dict
-        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol",
+        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol", "ctol",
         "maxiter", "max_nfev", "max_njev", "max_ngev", "max_nhev"}
     options : dict, optional
         Dictionary of optional keyword arguments to override default solver
@@ -544,7 +544,7 @@ def _optimize_scipy_constrained(  # noqa: C901 - FIXME: simplify this
         * 1 : display a termination report.
         * 2 : display progress during iterations
     stoptol : dict
-        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol",
+        Dictionary of stopping tolerances, with keys {"xtol", "ftol", "gtol", "ctol",
         "maxiter", "max_nfev", "max_njev", "max_ngev", "max_nhev"}
     options : dict, optional
         Dictionary of optional keyword arguments to override default solver