better-optimize 0.0.0__tar.gz → 0.0.2__tar.gz

{better_optimize-0.0.0 → better_optimize-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: better_optimize
-Version: 0.0.0
+Version: 0.0.2
 Summary: A drop-in replacement for scipy optimize functions with quality of life improvements
 License: MIT License
 Project-URL: Homepage, https://github.com/jessegrabowski/better_optimize
@@ -14,6 +14,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/markdown
 Provides-Extra: tests
 

better_optimize-0.0.2/better_optimize/__init__.py

@@ -0,0 +1,17 @@
+import logging
+
+from better_optimize._version import get_versions
+from better_optimize.minimize import minimize
+from better_optimize.root import root
+
+_log = logging.getLogger(__name__)
+
+if not logging.root.handlers:
+    _log.setLevel(logging.INFO)
+    if len(_log.handlers) == 0:
+        handler = logging.StreamHandler()
+        _log.addHandler(handler)
+
+
+__version__ = get_versions()["version"]
+__all__ = ["minimize", "root"]

better_optimize-0.0.2/better_optimize/_version.py

@@ -0,0 +1,21 @@
+
+# This file was generated by 'versioneer.py' (0.29) from
+# revision-control system data, or from the parent directory name of an
+# unpacked source archive. Distribution tarballs contain a pre-generated copy
+# of this file.
+
+import json
+
+version_json = '''
+{
+ "date": "2024-08-19T18:20:27+0800",
+ "dirty": false,
+ "error": null,
+ "full-revisionid": "3595536b9c1ead4ca6ea916046e4379d86e36ecc",
+ "version": "0.0.2"
+}
+'''  # END VERSION_JSON
+
+
+def get_versions():
+    return json.loads(version_json)

better_optimize-0.0.2/better_optimize/constants.py

@@ -0,0 +1,359 @@
+from typing import Literal
+
+minimize_method = Literal[
+    "nelder-mead",
+    "powell",
+    "CG",
+    "BFGS",
+    "Newton-CG",
+    "L-BFGS-B",
+    "TNC",
+    "COBYLA",
+    "SLSQP",
+    "trust-constr",
+    "dogleg",
+    "trust-ncg",
+    "trust-exact",
+    "trust-krylov",
+]
+
+root_method = Literal[
+    "hybr",
+    "lm",
+    "broyden1",
+    "broyden2",
+    "anderson",
+    "linearmixing",
+    "diagbroyden",
+    "excitingmixing",
+    "krylov",
+    "df-sane",
+]
+
+TOLERANCES = ["xtol", "ftol", "gtol", "fatol", "xatol"]
+
+MINIMIZE_MODE_KWARGS = {
+    "nelder-mead": {
+        "uses_grad": False,
+        "uses_hess": False,
+        "uses_hessp": False,
+        "valid_options": [
+            "disp",
+            "maxiter",
+            "maxfev",
+            "return_all",
+            "initial_simplex",
+            "xatol",
+            "fatol",
+            "adaptive",
+            "bounds",
+        ],
+    },
+    "powell": {
+        "uses_grad": False,
+        "uses_hess": False,
+        "uses_hessp": False,
+        "valid_options": ["disp", "xtol", "ftol", "maxiter", "maxfev", "direc", "return_all"],
+    },
+    "CG": {
+        "uses_grad": True,
+        "uses_hess": False,
+        "uses_hessp": False,
+        "valid_options": [
+            "disp",
+            "maxiter",
+            "gtol",
+            "norm",
+            "eps",
+            "return_all",
+            "finite_diff_rel_step",
+            "c1",
+            "c2",
+        ],
+    },
+    "BFGS": {
+        "uses_grad": True,
+        "uses_hess": False,
+        "uses_hessp": False,
+        "valid_options": [
+            "disp",
+            "maxiter",
+            "gtol",
+            "norm",
+            "eps",
+            "return_all",
+            "finite_diff_rel_step",
+            "xrtol",
+            "c1",
+            "c2",
+            "hess_inv0",
+        ],
+    },
+    "Newton-CG": {
+        "uses_grad": True,
+        "uses_hess": True,
+        "uses_hessp": True,
+        "valid_options": ["disp", "xtol", "maxiter", "eps", "return_all", "c1", "c2"],
+    },
+    "L-BFGS-B": {
+        "uses_grad": True,
+        "uses_hess": False,
+        "uses_hessp": False,
+        "valid_options": [
+            "disp",
+            "maxcor",
+            "ftol",
+            "gtol",
+            "eps",
+            "maxfun",
+            "maxiter",
+            "iprint",
+            "maxls",
+            "finite_diff_rel_step",
+        ],
+    },
+    "TNC": {
+        "uses_grad": True,
+        "uses_hess": False,
+        "uses_hessp": False,
+        "valid_options": [
+            "eps",
+            "scale",
+            "offset",
+            "disp",
+            "maxCGit",
+            "eta",
+            "stepmx",
+            "accuracy",
+            "minfev",
+            "ftol",
+            "xtol",
+            "gtol",
+            "rescale",
+            "finite_diff_rel_step",
+            "maxfun",
+        ],
+    },
+    "COBYLA": {
+        "uses_grad": False,
+        "uses_hess": False,
+        "uses_hessp": False,
+        "valid_options": ["rhobeg", "tol", "disp", "maxiter", "catol"],
+    },
+    "SLSQP": {
+        "uses_grad": True,
+        "uses_hess": False,
+        "uses_hessp": False,
+        "valid_options": [
+            "ftol",
+            "eps",
+            "disp",
+            "maxiter",
+            "finite_diff_rel_step",
+        ],
+    },
+    "trust-constr": {
+        "uses_grad": True,
+        "uses_hess": True,
+        "uses_hessp": True,
+        "valid_options": [
+            "gtol",
+            "xtol",
+            "barrier_tol",
+            "sparse_jacobian",
+            "initial_tr_radius",
+            "initial_constr_penalty",
+            "initial_barrier_parameter",
+            "initial_barrier_tolerance",
+            "factorization_method",
+            "finite_diff_rel_step",
+            "maxiter",
+            "verbose",
+            "disp",
+        ],
+    },
+    "dogleg": {
+        "uses_grad": True,
+        "uses_hess": True,
+        "uses_hessp": False,
+        "valid_options": ["initial_trust_radius", "max_trust_radius", "eta", "gtol"],
+    },
+    "trust-ncg": {
+        "uses_grad": True,
+        "uses_hess": True,
+        "uses_hessp": True,
+        "valid_options": ["initial_trust_radius", "max_trust_radius", "eta", "gtol"],
+    },
+    "trust-exact": {
+        "uses_grad": True,
+        "uses_hess": True,
+        "uses_hessp": False,
+        "valid_options": ["initial_trust_radius", "max_trust_radius", "eta", "gtol"],
+    },
+    "trust-krylov": {
+        "uses_grad": True,
+        "uses_hess": True,
+        "uses_hessp": True,
+        "valid_options": ["inexact"],
+    },
+}
+
+root_method = Literal[
+    "hybr",
+    "lm",
+    "broyden1",
+    "broyden2",
+    "anderson",
+    "linearmixing",
+    "diagbroyden",
+    "excitingmixing",
+    "krylov",
+    "df-sane",
+]
+
+ROOT_MODE_KWARGS = {
+    "hybr": {
+        "uses_jac": True,
+        "valid_options": ["col_deriv", "xtol", "maxfev", "band", "eps", "factor", "diag"],
+    },
+    "lm": {
+        "uses_jac": True,
+        "valid_options": ["col_deriv", "ftol", "xtol", "gtol", "maxiter", "eps", "factor", "diag"],
+    },
+    "broyden1": {
+        "uses_jac": False,
+        "valid_options": [
+            "nit",
+            "disp",
+            "maxiter",
+            "ftol",
+            "fatol",
+            "xtol",
+            "xatol",
+            "tol_norm",
+            "line_search",
+            "max_rank",
+            "jac_options",
+        ],
+        "jac_options": ["alpha", "reduction_method", "to_retain"],
+    },
+    "broyden2": {
+        "uses_jac": False,
+        "valid_options": [
+            "nit",
+            "disp",
+            "maxiter",
+            "ftol",
+            "fatol",
+            "xtol",
+            "xatol",
+            "tol_norm",
+            "line_search",
+            "max_rank",
+            "jac_options",
+        ],
+        "jac_options": ["alpha", "reduction_method", "to_retain"],
+    },
+    "anderson": {
+        "uses_jac": False,
+        "valid_options": [
+            "nit",
+            "disp",
+            "maxiter",
+            "ftol",
+            "fatol",
+            "xtol",
+            "xatol",
+            "tol_norm",
+            "line_search",
+            "jac_options",
+        ],
+        "jac_options": ["alpha", "M", "w0"],
+    },
+    "linearmixing": {
+        "uses_jac": False,
+        "valid_options": [
+            "nit",
+            "disp",
+            "maxiter",
+            "ftol",
+            "fatol",
+            "xtol",
+            "xatol",
+            "tol_norm",
+            "line_search",
+            "jac_options",
+        ],
+        "jac_options": ["alpha"],
+    },
+    "diagbroyden": {
+        "uses_jac": False,
+        "valid_options": [
+            "nit",
+            "disp",
+            "maxiter",
+            "ftol",
+            "fatol",
+            "xtol",
+            "xatol",
+            "tol_norm",
+            "line_search",
+            "jac_options",
+        ],
+        "jac_options": ["alpha"],
+    },
+    "excitingmixing": {
+        "uses_jac": False,
+        "valid_options": [
+            "nit",
+            "disp",
+            "maxiter",
+            "ftol",
+            "fatol",
+            "xtol",
+            "xatol",
+            "tol_norm",
+            "line_search",
+            "jac_options",
+        ],
+        "jac_options": ["alpha", "alphamax"],
+    },
+    "krylov": {
+        "uses_jac": False,
+        "valid_options": [
+            "nit",
+            "disp",
+            "maxiter",
+            "ftol",
+            "fatol",
+            "xtol",
+            "xatol",
+            "tol_norm",
+            "line_search",
+            "jac_options",
+        ],
+        "jac_options": [
+            "rdiff",
+            "method",
+            "inner_M",
+            "inner_tol",
+            "inner_maxiter",
+            "outer_k",
+        ],
+    },
+    "df-sane": {
+        "uses_jac": False,
+        "valid_options": [
+            "ftol",
+            "fatol",
+            "fnorm",
+            "maxfev",
+            "disp",
+            "eta_strategy",
+            "sigma_eps",
+            "sigma_0",
+            "M",
+            "line_search",
+        ],
+    },
+}

better_optimize-0.0.2/better_optimize/minimize.py

@@ -0,0 +1,105 @@
+from collections.abc import Callable
+from functools import partial
+
+import numpy as np
+
+from scipy.optimize import OptimizeResult
+from scipy.optimize import minimize as sp_minimize
+from scipy.sparse.linalg import LinearOperator
+
+from better_optimize.constants import minimize_method
+from better_optimize.utilities import (
+    check_f_is_fused,
+    determine_maxiter,
+    determine_tolerance,
+    kwargs_to_options,
+    validate_provided_functions_minimize,
+)
+from better_optimize.wrapper import ObjectiveWrapper, optimzer_early_stopping_wrapper
+
+
+def minimize(
+    f: Callable[..., float | tuple[float, np.ndarray]],
+    x0: np.ndarray,
+    method: minimize_method,
+    jac: Callable[..., np.ndarray] | None = None,
+    hess: Callable[..., np.ndarray | LinearOperator] | None = None,
+    hessp: Callable[..., np.ndarray] | None = None,
+    progressbar: bool = True,
+    verbose: bool = True,
+    args: tuple | None = None,
+    **optimizer_kwargs,
+) -> OptimizeResult:
+    """
+    Solve a minimization problem using the scipy.optimize.minimize function.
+
+    Parameters
+    ----------
+    x0: np.ndarray
+        The initial values of the parameters to optimize
+    f: Callable
+        The objective function to minimize
+    args: tuple, optional
+        Additional arguments to pass to the objective function. Additional arguments are also passed to the gradient
+        and Hessian functions, if provided
+    jac: Callable, optional
+        The gradient of the objective function
+    hess: Callable, optional
+        The Hessian of the objective function
+    hessp: Callable, optional
+        The Hessian-vector product of the objective function
+    method: str
+        The optimization method to use
+    progressbar: bool
+        Whether to display a progress bar
+    verbose: bool
+        Whether to display verbose output
+    optimizer_kwargs
+        Additional keyword arguments to pass to the optimizer
+
+    Returns
+    -------
+    optimizer_result: OptimizeResult
+        Optimization result
+
+    """
+    has_fused_f_and_grad = check_f_is_fused(f, x0, args)
+    validate_provided_functions_minimize(
+        method, jac, hess, hessp, has_fused_f_and_grad, verbose=verbose
+    )
+
+    options = optimizer_kwargs.pop("options", {})
+    optimizer_kwargs["options"] = options
+
+    optimizer_kwargs = kwargs_to_options(optimizer_kwargs, method)
+    maxiter, optimizer_kwargs = determine_maxiter(optimizer_kwargs, method)
+    optimizer_kwargs = determine_tolerance(optimizer_kwargs, method)
+
+    # Test hessian function -- if it returns a LinearOperator, it can't be used inside the wrapper
+    args = () if args is None else args
+    use_hess = hess is not None and not isinstance(hess(x0, *args), LinearOperator)
+
+    objective = ObjectiveWrapper(
+        maxeval=maxiter,
+        f=f,
+        jac=jac,
+        hess=hess if use_hess else None,
+        args=args,
+        progressbar=progressbar,
+        has_fused_f_and_grad=has_fused_f_and_grad,
+    )
+
+    f_optim = partial(
+        sp_minimize,
+        fun=objective,
+        x0=x0,
+        method=method,
+        jac=True if has_fused_f_and_grad or jac is not None else None,
+        hess=hess,
+        hessp=hessp,
+        callback=objective.callback,
+        **optimizer_kwargs,
+    )
+
+    optimizer_result = optimzer_early_stopping_wrapper(f_optim)
+    return optimizer_result

better_optimize-0.0.2/better_optimize/root.py

@@ -0,0 +1,96 @@
+from collections.abc import Callable
+from functools import partial
+
+import numpy as np
+
+from scipy.optimize import OptimizeResult
+from scipy.optimize import root as sp_root
+
+from better_optimize.constants import root_method
+from better_optimize.utilities import (
+    check_f_is_fused,
+    determine_maxiter,
+    determine_tolerance,
+    kwargs_to_jac_options,
+    kwargs_to_options,
+    validate_provided_functions_root,
+)
+from better_optimize.wrapper import ObjectiveWrapper, optimzer_early_stopping_wrapper
+
+
+def root(
+    f: Callable[..., np.ndarray | tuple[np.ndarray, np.ndarray]],
+    x0: np.ndarray,
+    method: root_method,
+    jac: Callable[..., np.ndarray] | None = None,
+    progressbar: bool = True,
+    verbose: bool = True,
+    args: tuple | None = None,
+    **optimizer_kwargs,
+) -> OptimizeResult:
+    """
+    Solve a minimization problem using the scipy.optimize.minimize function.
+
+    Parameters
+    ----------
+    x0: np.ndarray
+        The initial values of the parameters to optimize
+    f: Callable
+        The objective function to minimize
+    args: tuple, optional
+        Additional arguments to pass to the objective function. Additional arguments are also passed to the gradient
+        and Hessian functions, if provided
+    jac: Callable, optional
+        The gradient of the objective function
+    hess: Callable, optional
+        The Hessian of the objective function
+    hessp: Callable, optional
+        The Hessian-vector product of the objective function
+    method: str
+        The optimization method to use
+    progressbar: bool
+        Whether to display a progress bar
+    verbose: bool
+        Whether to display verbose output
+    optimizer_kwargs
+        Additional keyword arguments to pass to the optimizer
+
+    Returns
+    -------
+    optimizer_result: OptimizeResult
+        Optimization result
+
+    """
+    has_fused_f_and_grad = check_f_is_fused(f, x0, args)
+    validate_provided_functions_root(
+        method, jac, has_fused_f_and_grad, has_fused_f_and_grad, verbose=verbose
+    )
+
+    options = optimizer_kwargs.pop("options", {})
+    optimizer_kwargs["options"] = options
+    optimizer_kwargs = kwargs_to_options(optimizer_kwargs, method)
+    optimizer_kwargs = kwargs_to_jac_options(optimizer_kwargs, method)
+
+    maxiter, optimizer_kwargs = determine_maxiter(optimizer_kwargs, method)
+    optimizer_kwargs = determine_tolerance(optimizer_kwargs, method)
+
+    objective = ObjectiveWrapper(
+        maxeval=maxiter,
+        f=f,
+        jac=jac,
+        args=args,
+        progressbar=progressbar,
+        has_fused_f_and_grad=has_fused_f_and_grad,
+    )
+
+    f_optim = partial(
+        sp_root,
+        fun=objective,
+        x0=x0,
+        method=method,
+        jac=True if has_fused_f_and_grad or jac is not None else None,
+        **optimizer_kwargs,
+    )
+
+    optimizer_result = optimzer_early_stopping_wrapper(f_optim)
+    return optimizer_result

better_optimize-0.0.2/better_optimize/utilities.py

@@ -0,0 +1,188 @@
+import logging
+
+from collections.abc import Callable
+from copy import deepcopy
+
+import numpy as np
+
+from better_optimize.constants import (
+    MINIMIZE_MODE_KWARGS,
+    ROOT_MODE_KWARGS,
+    TOLERANCES,
+    minimize_method,
+    root_method,
+)
+
+_log = logging.getLogger(__name__)
+
+
+def get_option_kwargs(method: minimize_method | root_method):
+    if method in MINIMIZE_MODE_KWARGS.keys():
+        options_kwargs = MINIMIZE_MODE_KWARGS[method]
+    elif method in ROOT_MODE_KWARGS.keys():
+        options_kwargs = ROOT_MODE_KWARGS[method]
+    else:
+        raise ValueError(f"Unknown method: {method}")
+
+    return options_kwargs
+
+
+def validate_provided_functions_minimize(
+    method: minimize_method,
+    f_grad: Callable[[np.ndarray], np.ndarray] | None,
+    f_hess: Callable[[np.ndarray], np.ndarray] | None,
+    f_hessp: Callable[[np.ndarray], np.ndarray] | None,
+    has_fused_f_and_grad: bool,
+    verbose=True,
+) -> None:
+    has_grad, has_hess, has_hessp = map(lambda f: f is not None, [f_grad, f_hess, f_hessp])
+    uses_grad, uses_hess, uses_hessp, _ = MINIMIZE_MODE_KWARGS[method].values()
+
+    if has_fused_f_and_grad and has_grad:
+        _log.warning(
+            "Objective function returns a tuple (interpreted as (loss, gradient), but a gradient function was "
+            "also provided. The gradient function will be ignored."
+        )
+
+    elif has_fused_f_and_grad:
+        has_grad = True
+
+    if method not in MINIMIZE_MODE_KWARGS:
+        raise ValueError(
+            f"Method {method} not recognized. Must be one of {list(MINIMIZE_MODE_KWARGS.keys())}"
+        )
+
+    if has_hess and has_hessp:
+        raise ValueError(
+            "Cannot ask for Hessian and Hessian-vector product at the same time. For all algorithms "
+            "except trust-exact and dogleg, use_hessp is preferred."
+        )
+
+    if has_grad and not uses_grad and verbose:
+        _log.warning(
+            f"Gradient provided but not used by method {method}. Gradients will still be evaluated at each "
+            f"optimzer step and the norm will be reported as a diagnositc. For large problems, this may be "
+            f"computationally intensive."
+        )
+
+    if (has_hess and not uses_hess) or (has_hessp and not uses_hessp) and verbose:
+        _log.warning(
+            f"Hessian provided but not used by method {method}. The full hessian will still be evaluated at "
+            f"each optimzer step and the norm will be reported as a diagnositc. For large problems, this may "
+            f"be computationally intensive."
+        )
+
+    if has_hess and not has_hessp and uses_hessp and uses_hess:
+        _log.warning(
+            f"You provided a function to compute the full hessian, but method {method} allows the use of a "
+            f"hessian-vector product instead. Consider passing hessp instead -- this may be significantly "
+            f"more efficient."
+        )
+
+
+def validate_provided_functions_root(
+    method: root_method, f, jac, has_fused_f_and_grad: bool, verbose: bool = True
+):
+    has_jac = jac is not None
+    info_dict = get_option_kwargs(method)
+    uses_jac = info_dict["uses_jac"]
+
+    if has_fused_f_and_grad and has_jac:
+        _log.warning(
+            "Objective function returns a tuple (interpreted as (loss, gradient), but a gradient function was "
+            "also provided. The gradient function will be ignored."
+        )
+
+    elif has_fused_f_and_grad:
+        has_jac = True
+
+    if has_jac and not uses_jac and verbose:
+        _log.warning(
+            f"Gradient provided but not used by method {method}. Gradients will still be evaluated at each "
+            f"optimzer step and the norm will be reported as a diagnositc. For large problems, this may be "
+            f"computationally intensive."
+        )
+
+
+def check_f_is_fused(f, x0, args):
+    args = () if args is None else args
+    output = f(x0, *args)
+    if not isinstance(output, tuple | list):
+        return False
+
+    # If the output is a tuple, it should be length 2 (returning the value and the grad).If not, something is wrong
+    if len(output) != 2:
+        raise ValueError(
+            "Objective function should return either a scalar loss or a two-tuple of (loss, gradient)"
+        )
+    return True
+
+
+def determine_maxiter(
+    optimizer_kwargs: dict, method: minimize_method | root_method
+) -> tuple[int, dict]:
+    MAXITER_KWARGS = ["maxiter", "maxfun", "maxfev"]
+    maxiter = optimizer_kwargs.pop("maxiter", 5000)
+    optimizer_kwargs["options"].update({"maxiter": maxiter})
+    maxiter_kwargs = [x for x in get_option_kwargs(method)["valid_options"] if x in MAXITER_KWARGS]
+
+    for kwarg in maxiter_kwargs:
+        if kwarg not in optimizer_kwargs["options"]:
+            optimizer_kwargs["options"][kwarg] = maxiter
+
+    return maxiter, optimizer_kwargs
+
+
+def kwargs_to_options(optimizer_kwargs: dict, method: minimize_method | root_method) -> dict:
+    optimizer_kwargs = deepcopy(optimizer_kwargs)
+
+    NEVER_AUTO_PROMOTE = ["bounds", "tol", "jac_options"]
+    option_kwargs = get_option_kwargs(method)["valid_options"]
+
+    provided_kwargs = list(optimizer_kwargs.keys())
+    options = optimizer_kwargs.get("options", {})
+
+    for kwarg in option_kwargs:
+        if kwarg in provided_kwargs and kwarg not in NEVER_AUTO_PROMOTE:
+            options[kwarg] = optimizer_kwargs.pop(kwarg)
+
+    optimizer_kwargs["options"] = options
+    return optimizer_kwargs
+
+
+def kwargs_to_jac_options(optimizer_kwargs: dict, method: root_method) -> dict:
+    provided_kwargs = list(optimizer_kwargs.keys())
+    method_jac_kwargs = get_option_kwargs(method).get("jac_options", None)
+
+    if method_jac_kwargs is None:
+        return optimizer_kwargs
+
+    optimizer_kwargs = deepcopy(optimizer_kwargs)
+
+    if "options" not in optimizer_kwargs:
+        optimizer_kwargs["options"] = {}
+
+    if (
+        any(kwarg in method_jac_kwargs for kwarg in provided_kwargs)
+        and "jac_options" not in optimizer_kwargs["options"]
+    ):
+        optimizer_kwargs["options"]["jac_options"] = {}
+
+    for kwarg in provided_kwargs:
+        if kwarg in method_jac_kwargs:
+            optimizer_kwargs["options"]["jac_options"][kwarg] = optimizer_kwargs.pop(kwarg)
+
+    return optimizer_kwargs
+
+
+def determine_tolerance(optimizer_kwargs: dict, method: minimize_method | root_method) -> dict:
+    tol = optimizer_kwargs.pop("tol", 1e-8)
+    optimizer_kwargs["tol"] = tol
+    method_options = get_option_kwargs(method)["valid_options"]
+    method_tolerances = [tol_type for tol_type in method_options if tol_type in TOLERANCES]
+
+    for tol_name in method_tolerances:
+        if tol_name not in optimizer_kwargs["options"]:
+            optimizer_kwargs["options"][tol_name] = tol
+
+    return optimizer_kwargs

better_optimize-0.0.2/better_optimize/wrapper.py

@@ -0,0 +1,180 @@
+import logging
+
+from collections.abc import Callable
+
+import numpy as np
+
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+from rich.table import Column
+
+_log = logging.getLogger(__name__)
+
+
+class ObjectiveWrapper:
+    def __init__(
+        self,
+        f: Callable[[np.ndarray | float, ...], float | np.ndarray | tuple[float, np.ndarray]],
+        jac: Callable[[np.ndarray | float, ...], np.ndarray] | None = None,
+        hess: Callable[[np.ndarray | float, ...], np.ndarray] | None = None,
+        has_fused_f_and_grad: bool = False,
+        args: tuple | None = None,
+        maxeval: int = 5000,
+        progressbar: bool = True,
+        update_every: int = 10,
+    ):
+        self.n_eval = 0
+        self.maxeval = maxeval
+        self.args = args if args is not None else ()
+        self.f = lambda x: f(x, *self.args)
+        self.use_jac = False
+        self.use_hess = False
+        self.has_fused_f_and_grad = has_fused_f_and_grad
+
+        self.progress = None
+        self.task = None
+
+        self.update_every = update_every
+        self.interrupted = False
+        self.desc = "f = {task.fields[f_value]:,.5g}"
+
+        if jac is not None or has_fused_f_and_grad:
+            self.desc += ", ||grad|| = {task.fields[grad_norm]:,.5g}"
+            self.use_jac = True
+            self.f_jac = lambda x: jac(x, *self.args)
+
+        if hess is not None:
+            self.desc += ", ||hess|| = {task.fields[hess_norm]:,.5g}"
+            self.use_hess = True
+            self.f_hess = lambda x: hess(x, *self.args)
+
+        self.previous_x = None
+        self.progressbar = progressbar
+        if progressbar:
+            self.progress = self.initialize_progress_bar()
+            self.progress.start()
+
+    def step(self, x):
+        grad = None
+        hess = None
+        if self.has_fused_f_and_grad:
+            value, grad = self.f(x)
+        else:
+            value = self.f(x)
+
+        if self.use_jac and not self.has_fused_f_and_grad:
+            grad = self.f_jac(x)
+        if self.use_hess:
+            hess = self.f_hess(x)
+
+        if np.all(np.isfinite(x)):
+            self.previous_x = x
+
+        if self.n_eval % self.update_every == 0:
+            self.update_progressbar(value, grad, hess)
+
+        if self.n_eval > self.maxeval:
+            self.update_progressbar(value, grad, hess)
+            self.interrupted = True
+
+        self.n_eval += 1
+
+        if self.use_hess:
+            return value, grad, hess
+        elif self.use_jac:
+            return value, grad
+        else:
+            return value
+
+    def __call__(self, x):
+        try:
+            return self.step(x)
+        except (KeyboardInterrupt, StopIteration):
+            self.interrupted = True
+            return self.step(x)
+
+    def callback(self, *args):
+        if self.interrupted:
+            raise StopIteration
+
+    def update_progressbar(
+        self, value: float, grad: np.float64 = None, hess: np.float64 = None, completed=False
+    ) -> None:
+        if not self.progressbar:
+            return
+
+        if isinstance(value, np.ndarray):
+            value = (value**2).sum()
+        elif isinstance(value, list | tuple):
+            value = sum([x**2 for x in value])
+
+        value_dict = {"f_value": value}
+        if grad is not None:
+            grad_norm = np.linalg.norm(grad)
+            value_dict["grad_norm"] = grad_norm
+        if hess is not None:
+            hess_norm = np.linalg.norm(hess)
+            value_dict["hess_norm"] = hess_norm
+
+        if self.n_eval == 0:
+            self.task = self.progress.add_task("Optimizing", **value_dict)
+
+        if not completed:
+            self.progress.update(self.task, advance=self.update_every, **value_dict)
+        else:
+            self.progress.update(self.task, total=self.n_eval, completed=self.n_eval, **value_dict)
+
+    def initialize_progress_bar(self):
+        text_column = TextColumn("{task.description}", table_column=Column(ratio=1))
+        bar_column = BarColumn(bar_width=None, table_column=Column(ratio=2))
+        time_column = TimeElapsedColumn()
+        m_of_n = MofNCompleteColumn()
+        spinner = SpinnerColumn()
+
+        stat_column = TextColumn(self.desc, table_column=Column(ratio=1))
+
+        return Progress(
+            text_column, spinner, bar_column, time_column, m_of_n, stat_column, expand=False
+        )
+
+
+def optimzer_early_stopping_wrapper(f_optim):
+    objective = f_optim.keywords["fun"]
+    progressbar = objective.progressbar
+
+    try:
+        # Do the optimization. This calls either optimize.root or optimize.minimize;  all arguments are pre-configured
+        # at this point.
+        res = f_optim()
+    except Exception as e:
+        # Teardown the progress bar if necessary, then forward the error
+        if progressbar:
+            objective.progress.stop()
+        raise e
+
+    x_final = res.x
+
+    if progressbar:
+        # Evaluate the objective function one last time to get the final value, gradient, and hessian
+        # and update the progress bar to show the true final state
+        outputs = objective.step(x_final)
+        if not objective.use_jac and not objective.use_hess:
+            value = outputs
+            grad = None
+            hess = None
+        elif objective.use_jac and not objective.use_hess:
+            value, grad = outputs
+            hess = None
+        else:
+            value, grad, hess = outputs
+
+        objective.update_progressbar(value, grad, hess, completed=True)
+        objective.progress.stop()
+
+    return res

{better_optimize-0.0.0 → better_optimize-0.0.2}/better_optimize.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: better-optimize
-Version: 0.0.0
+Version: 0.0.2
 Summary: A drop-in replacement for scipy optimize functions with quality of life improvements
 License: MIT License
 Project-URL: Homepage, https://github.com/jessegrabowski/better_optimize
@@ -14,6 +14,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/markdown
 Provides-Extra: tests
 

better_optimize-0.0.2/better_optimize.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+README.md
+pyproject.toml
+setup.py
+better_optimize/__init__.py
+better_optimize/_version.py
+better_optimize/constants.py
+better_optimize/minimize.py
+better_optimize/root.py
+better_optimize/utilities.py
+better_optimize/wrapper.py
+better_optimize.egg-info/PKG-INFO
+better_optimize.egg-info/SOURCES.txt
+better_optimize.egg-info/dependency_links.txt
+better_optimize.egg-info/requires.txt
+better_optimize.egg-info/top_level.txt

{better_optimize-0.0.0 → better_optimize-0.0.2}/pyproject.toml

@@ -1,6 +1,15 @@
+[build-system]
+requires = [
+    "numpy",
+    "setuptools",
+    "versioneer[toml]",
+]
+build-backend = "setuptools.build_meta"
+
 [project]
 name = 'better_optimize'
 dynamic = ['version']
+requires-python = ">=3.10,<3.13"
 description = "A drop-in replacement for scipy optimize functions with quality of life improvements"
 readme = "README.md"
 license = {text = "MIT License"}
@@ -58,7 +67,7 @@ exclude_lines = [
 [tool.ruff]
 line-length = 100
 target-version = "py310"
-extend-exclude = ["better_optimize/_version.py"]
+extend-exclude = ["better_optimize/_version.py", "setup.py"]
 
 [tool.ruff.lint]
 select = ["D", "E", "F", "I", "UP", "W", "RUF"]

better_optimize-0.0.2/setup.py

@@ -0,0 +1,19 @@
+from setuptools import find_packages, setup
+from setuptools.dist import Distribution
+
+import versioneer
+
+dist = Distribution()
+dist.parse_config_files()
+
+
+NAME: str = dist.get_name()  # type: ignore
+
+
+if __name__ == "__main__":
+    setup(
+        name=NAME,
+        version=versioneer.get_version(),
+        cmdclass=versioneer.get_cmdclass(),
+        packages=find_packages(exclude=["tests*"]),
+    )

better_optimize-0.0.0/better_optimize.egg-info/SOURCES.txt

@@ -1,8 +0,0 @@
-README.md
-pyproject.toml
-setup.py
-better_optimize.egg-info/PKG-INFO
-better_optimize.egg-info/SOURCES.txt
-better_optimize.egg-info/dependency_links.txt
-better_optimize.egg-info/requires.txt
-better_optimize.egg-info/top_level.txt

better_optimize-0.0.0/setup.py

@@ -1,3 +0,0 @@
-from setuptools import setup
-
-setup()