<commit_before>import scipy.optimize
import warnings
from termcolor import colored

from desc.backend import jnp
from desc.utils import Timer
from desc.io import IOAble
from desc.objectives import (
    ObjectiveFunction,
    ForceBalance,
    RadialForceBalance,
    HelicalForceBalance,
    WrappedEquilibriumObjective,
)
from desc.objectives.utils import factorize_linear_constraints
from desc.optimize import fmintr, lsqtr
from .utils import check_termination, print_header_nonlinear, print_iteration_nonlinear


class Optimizer(IOAble):
    """A helper class to wrap several different optimization routines

    Offers all of the ``scipy.optimize.least_squares`` routines  and several of the most
    useful ``scipy.optimize.minimize`` routines.
    Also offers several custom routines specifically designed for DESC, both scalar and
    least squares routines with and without jacobian/hessian information.

    Parameters
    ----------
    method : str
        name of the optimizer to use. Options are:

        * scipy scalar routines: ``'scipy-bfgs'``, ``'scipy-trust-exact'``,
          ``'scipy-trust-ncg'``, ``'scipy-trust-krylov'``
        * scipy least squares routines: ``'scipy-trf'``, ``'scipy-lm'``, ``'scipy-dogbox'``
        * desc scalar routines: ``'dogleg'``, ``'subspace'``, ``'dogleg-bfgs'``,
          ``'subspace-bfgs'``
        * desc least squares routines: ``'lsq-exact'``

    objective : ObjectiveFunction
        objective to be optimized

    """

    _io_attrs_ = ["_method"]

    # TODO: better way to organize these:
    _scipy_least_squares_methods = ["scipy-trf", "scipy-lm", "scipy-dogbox"]
    _scipy_scalar_methods = [
        "scipy-bfgs",
        "scipy-trust-exact",
        "scipy-trust-ncg",
        "scipy-trust-krylov",
    ]
    _desc_scalar_methods = ["dogleg", "subspace", "dogleg-bfgs", "subspace-bfgs"]
    _desc_least_squares_methods = ["lsq-exact"]
    _hessian_free_methods = ["scipy-bfgs", "dogleg-bfgs", "subspace-bfgs"]
    _scipy_constrained_scalar_methods = ["scipy-trust-constr"]
    _scipy_constrained_least_squares_methods = []
    _desc_constrained_scalar_methods = []
    _desc_constrained_least_squares_methods = []
    _scalar_methods = (
        _desc_scalar_methods
        + _scipy_scalar_methods
        + _scipy_constrained_scalar_methods
        + _desc_constrained_scalar_methods
    )
    _least_squares_methods = (
        _scipy_least_squares_methods
        + _desc_least_squares_methods
        + _scipy_constrained_least_squares_methods
        + _desc_constrained_least_squares_methods
    )
    _scipy_methods = (
        _scipy_least_squares_methods
        + _scipy_scalar_methods
        + _scipy_constrained_scalar_methods
        + _scipy_constrained_least_squares_methods
    )
    _desc_methods = (
        _desc_least_squares_methods
        + _desc_scalar_methods
        + _desc_constrained_scalar_methods
        + _desc_constrained_least_squares_methods
    )
    _constrained_methods = (
        _desc_constrained_scalar_methods
        + _desc_constrained_least_squares_methods
        + _scipy_constrained_scalar_methods
        + _scipy_constrained_least_squares_methods
    )
    _all_methods = (
        _scipy_least_squares_methods
        + _scipy_scalar_methods
        + _desc_scalar_methods
        + _desc_least_squares_methods
    )

    def __init__(self, method):

        self.method = method

    def __repr__(self):
        """string form of the object"""
        return (
            type(self).__name__
            + " at "
            + str(hex(id(self)))
            + " (method={})".format(self.method)
        )

    @property
    def method(self):
        """str : name of the optimization method"""
        return self._method

    @method.setter
    def method(self, method):
        if method not in Optimizer._all_methods:
            raise NotImplementedError(
                colored(
                    "method must be one of {}".format(
                        ".".join([Optimizer._all_methods])
                    ),
                    "red",
                )
            )
        self._method = method

    # TODO: add copy argument and return the equilibrium?
    def optimize(
        self,
        eq,
        objective,
        constraints=(),
        ftol=1e-6,
        xtol=1e-6,
        gtol=1e-6,
        x_scale="auto",
        verbose=1,
        maxiter=None,
        options={},
    ):
        """Optimize an objective function.

        Parameters
        ----------
        eq : Equilibrium
            Initial equilibrium.
        objective : ObjectiveFunction
            Objective function to optimize.
        constraints : tuple of Objective, optional
            List of objectives to be used as constraints during optimization.
        ftol : float or None, optional
            Tolerance for termination by the change of the cost function.
            Default is 1e-8. The optimization process is stopped when ``dF < ftol * F``,
            and there was an adequate agreement between a local quadratic model and the
            true model in the last step.
            If None, the termination by this condition is disabled.
        xtol : float or None, optional
            Tolerance for termination by the change of the independent variables.
            Default is 1e-8.
            Optimization is stopped when ``norm(dx) < xtol * (xtol + norm(x))``.
            If None, the termination by this condition is disabled.
        gtol : float or None, optional
            Absolute tolerance for termination by the norm of the gradient.
            Default is 1e-8. Optimizer teriminates when ``norm(g) < gtol``, where
            # FIXME: missing documentation!
            If None, the termination by this condition is disabled.
        x_scale : array_like or ``'auto'``, 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 ``'auto'``, the scale is iteratively updated using the
            inverse norms of the columns of the jacobian or hessian matrix.
        verbose : integer, optional
            * 0  : work silently.
            * 1-2 : display a termination report.
            * 3 : display progress during iterations
        maxiter : int, optional
            Maximum number of iterations. Defaults to size(x)*100.
        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.

        """
        # TODO: document options
        # scipy optimizers expect disp={0,1,2} while we use verbose={0,1,2,3}
        disp = verbose - 1 if verbose > 1 else verbose

        timer = Timer()

        if self.method in Optimizer._desc_methods:
            if not isinstance(x_scale, str) and jnp.allclose(x_scale, 1):
                options.setdefault("initial_trust_radius", 0.5)
                options.setdefault("max_trust_radius", 1.0)

        if not isinstance(constraints, tuple):
            constraints = (constraints,)
        linear_constraints = tuple(
            constraint for constraint in constraints if constraint.linear
        )
        nonlinear_constraints = tuple(
            constraint for constraint in constraints if not constraint.linear
        )

        # wrap nonlinear constraints if necessary
        wrapped = False
        if len(nonlinear_constraints) > 0 and (
            self.method not in Optimizer._constrained_methods
        ):
            wrapped = True
            for constraint in nonlinear_constraints:
                if not isinstance(
                    constraint, (ForceBalance, RadialForceBalance, HelicalForceBalance)
                ):
                    raise ValueError(
                        "optimizer method {} ".format(self.method)
                        + "cannot handle general nonlinear constraint {}.".format(
                            constraint
                        )
                    )
            objective = WrappedEquilibriumObjective(
                objective,
                eq_objective=ObjectiveFunction(nonlinear_constraints),
                perturb_options=options.pop("perturb_options", {}),
                solve_options=options.pop("solve_options", {}),
            )

        if not objective.built:
            objective.build(eq, verbose=verbose)
        if not objective.compiled:
            mode = "scalar" if self.method in Optimizer._scalar_methods else "lsq"
            objective.compile(mode, verbose)
        for constraint in linear_constraints:
            if not constraint.built:
                constraint.build(eq, verbose=verbose)

        if objective.scalar and (self.method in Optimizer._least_squares_methods):
            warnings.warn(
                colored(
                    "method {} is not intended for scalar objective function".format(
                        ".".join([self.method])
                    ),
                    "yellow",
                )
            )

        if verbose > 0:
            print("Factorizing linear constraints")
        timer.start("linear constraint factorize")
        _, _, _, _, Z, unfixed_idx, project, recover = factorize_linear_constraints(
            linear_constraints
        )
        timer.stop("linear constraint factorize")
        if verbose > 1:
            timer.disp("linear constraint factorize")

        if verbose > 0:
            print("Starting optimization")
        timer.start("Solution time")

        def compute_wrapped(x_reduced):
            x = recover(x_reduced)
            f = objective.compute(x)
            if self.method in Optimizer._scalar_methods:
                return f.squeeze()
            else:
                return jnp.atleast_1d(f)

        def compute_scalar_wrapped(x_reduced):
            x = recover(x_reduced)
            return objective.compute_scalar(x)

        def grad_wrapped(x_reduced):
            x = recover(x_reduced)
            df = objective.grad(x)
            return df[unfixed_idx] @ Z

        def hess_wrapped(x_reduced):
            x = recover(x_reduced)
            df = objective.hess(x)
            return Z.T @ df[unfixed_idx, :][:, unfixed_idx] @ Z

        def jac_wrapped(x_reduced):
            x = recover(x_reduced)
            df = objective.jac(x)
            return df[:, unfixed_idx] @ Z

        x0_reduced = project(objective.x(eq))
<commit_msg>add CurrentDensity to wrapped objectives<commit_after>import scipy.optimize
import warnings
from termcolor import colored

from desc.backend import jnp
from desc.utils import Timer
from desc.io import IOAble
from desc.objectives import (
    ObjectiveFunction,
    ForceBalance,
    RadialForceBalance,
    HelicalForceBalance,
    CurrentDensity,
    WrappedEquilibriumObjective,
)
from desc.objectives.utils import factorize_linear_constraints
from desc.optimize import fmintr, lsqtr
from .utils import check_termination, print_header_nonlinear, print_iteration_nonlinear


class Optimizer(IOAble):
    """A helper class to wrap several different optimization routines

    Offers all of the ``scipy.optimize.least_squares`` routines  and several of the most
    useful ``scipy.optimize.minimize`` routines.
    Also offers several custom routines specifically designed for DESC, both scalar and
    least squares routines with and without jacobian/hessian information.

    Parameters
    ----------
    method : str
        name of the optimizer to use. Options are:

        * scipy scalar routines: ``'scipy-bfgs'``, ``'scipy-trust-exact'``,
          ``'scipy-trust-ncg'``, ``'scipy-trust-krylov'``
        * scipy least squares routines: ``'scipy-trf'``, ``'scipy-lm'``, ``'scipy-dogbox'``
        * desc scalar routines: ``'dogleg'``, ``'subspace'``, ``'dogleg-bfgs'``,
          ``'subspace-bfgs'``
        * desc least squares routines: ``'lsq-exact'``

    objective : ObjectiveFunction
        objective to be optimized

    """

    _io_attrs_ = ["_method"]

    # TODO: better way to organize these:
    _scipy_least_squares_methods = ["scipy-trf", "scipy-lm", "scipy-dogbox"]
    _scipy_scalar_methods = [
        "scipy-bfgs",
        "scipy-trust-exact",
        "scipy-trust-ncg",
        "scipy-trust-krylov",
    ]
    _desc_scalar_methods = ["dogleg", "subspace", "dogleg-bfgs", "subspace-bfgs"]
    _desc_least_squares_methods = ["lsq-exact"]
    _hessian_free_methods = ["scipy-bfgs", "dogleg-bfgs", "subspace-bfgs"]
    _scipy_constrained_scalar_methods = ["scipy-trust-constr"]
    _scipy_constrained_least_squares_methods = []
    _desc_constrained_scalar_methods = []
    _desc_constrained_least_squares_methods = []
    _scalar_methods = (
        _desc_scalar_methods
        + _scipy_scalar_methods
        + _scipy_constrained_scalar_methods
        + _desc_constrained_scalar_methods
    )
    _least_squares_methods = (
        _scipy_least_squares_methods
        + _desc_least_squares_methods
        + _scipy_constrained_least_squares_methods
        + _desc_constrained_least_squares_methods
    )
    _scipy_methods = (
        _scipy_least_squares_methods
        + _scipy_scalar_methods
        + _scipy_constrained_scalar_methods
        + _scipy_constrained_least_squares_methods
    )
    _desc_methods = (
        _desc_least_squares_methods
        + _desc_scalar_methods
        + _desc_constrained_scalar_methods
        + _desc_constrained_least_squares_methods
    )
    _constrained_methods = (
        _desc_constrained_scalar_methods
        + _desc_constrained_least_squares_methods
        + _scipy_constrained_scalar_methods
        + _scipy_constrained_least_squares_methods
    )
    _all_methods = (
        _scipy_least_squares_methods
        + _scipy_scalar_methods
        + _desc_scalar_methods
        + _desc_least_squares_methods
    )

    def __init__(self, method):

        self.method = method

    def __repr__(self):
        """string form of the object"""
        return (
            type(self).__name__
            + " at "
            + str(hex(id(self)))
            + " (method={})".format(self.method)
        )

    @property
    def method(self):
        """str : name of the optimization method"""
        return self._method

    @method.setter
    def method(self, method):
        if method not in Optimizer._all_methods:
            raise NotImplementedError(
                colored(
                    "method must be one of {}".format(
                        ".".join([Optimizer._all_methods])
                    ),
                    "red",
                )
            )
        self._method = method

    # TODO: add copy argument and return the equilibrium?
    def optimize(
        self,
        eq,
        objective,
        constraints=(),
        ftol=1e-6,
        xtol=1e-6,
        gtol=1e-6,
        x_scale="auto",
        verbose=1,
        maxiter=None,
        options={},
    ):
        """Optimize an objective function.

        Parameters
        ----------
        eq : Equilibrium
            Initial equilibrium.
        objective : ObjectiveFunction
            Objective function to optimize.
        constraints : tuple of Objective, optional
            List of objectives to be used as constraints during optimization.
        ftol : float or None, optional
            Tolerance for termination by the change of the cost function.
            Default is 1e-8. The optimization process is stopped when ``dF < ftol * F``,
            and there was an adequate agreement between a local quadratic model and the
            true model in the last step.
            If None, the termination by this condition is disabled.
        xtol : float or None, optional
            Tolerance for termination by the change of the independent variables.
            Default is 1e-8.
            Optimization is stopped when ``norm(dx) < xtol * (xtol + norm(x))``.
            If None, the termination by this condition is disabled.
        gtol : float or None, optional
            Absolute tolerance for termination by the norm of the gradient.
            Default is 1e-8. Optimizer teriminates when ``norm(g) < gtol``, where
            # FIXME: missing documentation!
            If None, the termination by this condition is disabled.
        x_scale : array_like or ``'auto'``, 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 ``'auto'``, the scale is iteratively updated using the
            inverse norms of the columns of the jacobian or hessian matrix.
        verbose : integer, optional
            * 0  : work silently.
            * 1-2 : display a termination report.
            * 3 : display progress during iterations
        maxiter : int, optional
            Maximum number of iterations. Defaults to size(x)*100.
        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.

        """
        # TODO: document options
        # scipy optimizers expect disp={0,1,2} while we use verbose={0,1,2,3}
        disp = verbose - 1 if verbose > 1 else verbose

        timer = Timer()

        if self.method in Optimizer._desc_methods:
            if not isinstance(x_scale, str) and jnp.allclose(x_scale, 1):
                options.setdefault("initial_trust_radius", 0.5)
                options.setdefault("max_trust_radius", 1.0)

        if not isinstance(constraints, tuple):
            constraints = (constraints,)
        linear_constraints = tuple(
            constraint for constraint in constraints if constraint.linear
        )
        nonlinear_constraints = tuple(
            constraint for constraint in constraints if not constraint.linear
        )

        # wrap nonlinear constraints if necessary
        wrapped = False
        if len(nonlinear_constraints) > 0 and (
            self.method not in Optimizer._constrained_methods
        ):
            wrapped = True
            for constraint in nonlinear_constraints:
                if not isinstance(
                    constraint,
                    (
                        ForceBalance,
                        RadialForceBalance,
                        HelicalForceBalance,
                        CurrentDensity,
                    ),
                ):
                    raise ValueError(
                        "optimizer method {} ".format(self.method)
                        + "cannot handle general nonlinear constraint {}.".format(
                            constraint
                        )
                    )
            objective = WrappedEquilibriumObjective(
                objective,
                eq_objective=ObjectiveFunction(nonlinear_constraints),
                perturb_options=options.pop("perturb_options", {}),
                solve_options=options.pop("solve_options", {}),
            )

        if not objective.built:
            objective.build(eq, verbose=verbose)
        if not objective.compiled:
            mode = "scalar" if self.method in Optimizer._scalar_methods else "lsq"
            objective.compile(mode, verbose)
        for constraint in linear_constraints:
            if not constraint.built:
                constraint.build(eq, verbose=verbose)

        if objective.scalar and (self.method in Optimizer._least_squares_methods):
            warnings.warn(
                colored(
                    "method {} is not intended for scalar objective function".format(
                        ".".join([self.method])
                    ),
                    "yellow",
                )
            )

        if verbose > 0:
            print("Factorizing linear constraints")
        timer.start("linear constraint factorize")
        _, _, _, _, Z, unfixed_idx, project, recover = factorize_linear_constraints(
            linear_constraints, extra_args=objective.args
        )
        timer.stop("linear constraint factorize")
        if verbose > 1:
            timer.disp("linear constraint factorize")

        x0_reduced = project(objective.x(eq))

        if verbose > 0:
            print("Number of parameters: {}".format(x0_reduced.size))
            print("Number of objectives: {}".format(objective.dim_f))

        if verbose > 0:
            print("Starting optimization")
        timer.start("Solution time")

        def compute_wrapped(x_reduced):
            x = recover(x_reduced)
            f = objective.compute(x)
            if self.method in Optimizer._scalar_methods:
                return f.squeeze()
            else:
                return jnp.atleast_1d(f)

        def compute_scalar_wrapped(x_reduced):
            x = recover(x_reduced)
            return objective.compute_scalar(x)

        def grad_wrapped(x_reduced):
            x = recover(x_reduced)
            df = objective.grad(x)
            return df[unfixed_idx] @ Z

        def hess_wrapped(x_reduced):
            x = recover(x_reduced)
            df = objective.hess(x)
            return Z.T @ df[unfixed_idx, :][:, unfixed_idx] @ Z

        def jac_wrapped(x_reduced):
            x = recover(x_reduced)
            df = objective.jac(x)
            return df[:, unfixed_idx] @ Z

{'fraction_non_alphanumeric': 0.07982772871623091, 'fraction_numerical': 0.003567233653804324, 'mean_word_length': 2.0191752035723667, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4708183', 'n_tokens_mistral': 6478, 'n_tokens_neox': 6015, 'n_words': 1702}