torchzero 0.3.15__py3-none-any.whl → 0.4.1__py3-none-any.whl

torchzero/optim/wrappers/pybobyqa.py

@@ -0,0 +1,124 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import torch
+import pybobyqa
+
+from ...utils import TensorList
+from .wrapper import WrapperBase
+
+Closure = Callable[[bool], Any]
+
+
+class PyBobyqaWrapper(WrapperBase):
+    """Use Py-BOBYQA is PyTorch optimizer.
+
+    Note that this performs full minimization on each step,
+    so usually you would want to perform a single step, although performing multiple steps will refine the
+    solution.
+
+    See https://numericalalgorithmsgroup.github.io/pybobyqa/build/html/userguide.html for detailed descriptions of arguments.
+
+    Args:
+        params (Iterable): iterable of parameters to optimize or dicts defining parameter groups.
+        lb (float | None, optional): optional lower bounds. Defaults to None.
+        ub (float | None, optional): optional upper bounds. Defaults to None.
+        projections (list[Callable] | None, optional):
+            a list of functions defining the Euclidean projections for each general convex constraint C_i.
+            Each element of the list projections is a function that takes an input vector x (numpy array)
+            and returns the closest point to that is in C_i. Defaults to None.
+        npt (int | None, optional): the number of interpolation points to use. Defaults to None.
+        rhobeg (float | None, optional):
+            the initial value of the trust region radius. Defaults to None.
+        rhoend (float | None, optional):
+            minimum allowed value of trust region radius, which determines when a successful
+            termination occurs. Defaults to 1e-8.
+        maxfun (int | None, optional):
+            the maximum number of objective evaluations the algorithm may request,
+            default is min(100(n+1), 1000). Defaults to None.
+        nsamples (Callable | None, optional):
+            a Python function nsamples(delta, rho, iter, nrestarts)
+            which returns the number of times to evaluate objfun at a given point.
+            This is only applicable for objectives with stochastic noise,
+            when averaging multiple evaluations at the same point produces a more accurate value.
+            The input parameters are the trust region radius (delta),
+            the lower bound on the trust region radius (rho),
+            how many iterations the algorithm has been running for (iter),
+            and how many restarts have been performed (nrestarts).
+            Default is no averaging (i.e. nsamples(delta, rho, iter, nrestarts)=1).
+            Defaults to None.
+        user_params (dict | None, optional):
+            dictionary of advanced parameters,
+            see https://numericalalgorithmsgroup.github.io/pybobyqa/build/html/advanced.html).
+            Defaults to None.
+        objfun_has_noise (bool, optional):
+            a flag to indicate whether or not objfun has stochastic noise;
+            i.e. will calling objfun(x) multiple times at the same value of x give different results?
+            This is used to set some sensible default parameters (including using multiple restarts),
+            all of which can be overridden by the values provided in user_params. Defaults to False.
+        seek_global_minimum (bool, optional):
+            a flag to indicate whether to search for a global minimum, rather than a local minimum.
+            This is used to set some sensible default parameters,
+            all of which can be overridden by the values provided in user_params.
+            If True, both upper and lower bounds must be set.
+            Note that Py-BOBYQA only implements a heuristic method,
+            so there are no guarantees it will find a global minimum.
+            However, by using this flag, it is more likely to escape local minima
+            if there are better values nearby. The method used is a multiple restart mechanism,
+            where we repeatedly re-initialize Py-BOBYQA from the best point found so far,
+            but where we use a larger trust reigon radius each time
+            (note: this is different to more common multi-start approach to global optimization).
+            Defaults to False.
+        scaling_within_bounds (bool, optional):
+            a flag to indicate whether the algorithm should internally shift and scale the entries of x
+            so that the bounds become 0 <= x <= 1. This is useful is you are setting bounds and the
+            bounds have different orders of magnitude. If scaling_within_bounds=True,
+            the values of rhobeg and rhoend apply to the shifted variables. Defaults to False.
+        do_logging (bool, optional):
+            a flag to indicate whether logging output should be produced.
+            This is not automatically visible unless you use the Python logging module. Defaults to True.
+        print_progress (bool, optional):
+            a flag to indicate whether to print a per-iteration progress log to terminal. Defaults to False.
+    """
+    def __init__(
+        self,
+        params,
+        lb: float | None = None,
+        ub: float | None = None,
+        projections = None,
+        npt: int | None = None,
+        rhobeg: float | None = None,
+        rhoend: float = 1e-8,
+        maxfun: int | None = None,
+        nsamples: Callable | None | None = None,
+        user_params: dict[str, Any] | None = None,
+        objfun_has_noise: bool = False,
+        seek_global_minimum: bool = False,
+        scaling_within_bounds: bool = False,
+        do_logging: bool = True,
+        print_progress: bool = False,
+    ):
+        super().__init__(params, dict(lb=lb, ub=ub))
+        kwargs = locals().copy()
+        for k in ["self", "__class__", "params", "lb", "ub"]:
+            del kwargs[k]
+        self._kwargs = kwargs
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
+        bounds = self._get_bounds()
+
+        soln: pybobyqa.solver.OptimResults = pybobyqa.solve(
+            objfun=partial(self._f, closure=closure, params=params),
+            x0=x0,
+            bounds=bounds,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.as_tensor(soln.x, device = params[0].device, dtype=params[0].dtype,))
+        return soln.f
+

torchzero/optim/wrappers/scipy/__init__.py

@@ -0,0 +1,7 @@
+from .basin_hopping import ScipyBasinHopping
+from .brute import ScipyBrute
+from .differential_evolution import ScipyDE
+from .direct import ScipyDIRECT
+from .dual_annealing import ScipyDualAnnealing
+from .minimize import ScipyMinimize
+from .sgho import ScipySHGO

torchzero/optim/wrappers/scipy/basin_hopping.py

@@ -0,0 +1,117 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import scipy.optimize
+import torch
+
+from ....utils import TensorList
+from ..wrapper import WrapperBase
+from .minimize import _use_jac_hess_hessp
+
+Closure = Callable[[bool], Any]
+
+
+class ScipyBasinHopping(WrapperBase):
+    def __init__(
+        self,
+        params,
+        niter: int = 100,
+        T: float = 1,
+        stepsize: float = 0.5,
+        minimizer_kwargs: dict | None = None,
+        take_step: Callable | None = None,
+        accept_test: Callable | None = None,
+        callback: Callable | None = None,
+        interval: int = 50,
+        disp: bool = False,
+        niter_success: int | None = None,
+        rng: int | np.random.Generator | None = None,
+        lb:float | None = None,
+        ub:float | None = None,
+        method: Literal['nelder-mead', 'powell', 'cg', 'bfgs', 'newton-cg',
+                    'l-bfgs-b', 'tnc', 'cobyla', 'cobyqa', 'slsqp',
+                    'trust-constr', 'dogleg', 'trust-ncg', 'trust-exact',
+                    'trust-krylov'] | str | None = None,
+        jac: Literal['2-point', '3-point', 'cs', 'autograd'] = 'autograd',
+        hess: Literal['2-point', '3-point', 'cs', 'autograd'] | scipy.optimize.HessianUpdateStrategy = 'autograd',
+        use_hessp: bool = True,
+
+        *,
+        target_accept_rate: float = 0.5,
+        stepwise_factor: float = 0.9
+    ):
+        super().__init__(params, dict(lb=lb, ub=ub))
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['__class__'], kwargs["minimizer_kwargs"]
+        del kwargs['method'], kwargs["jac"], kwargs['hess'], kwargs['use_hessp']
+        del kwargs["lb"], kwargs["ub"]
+        self._kwargs = kwargs
+
+        self._minimizer_kwargs = minimizer_kwargs
+        self.method = method
+        self.hess = hess
+        self.jac, self.use_jac_autograd, self.use_hess_autograd, self.use_hessp = _use_jac_hess_hessp(method, jac, hess, use_hessp)
+
+    def _jac(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        f,g = self._f_g(x, params, closure)
+        return g
+
+    def _objective(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        if self.use_jac_autograd:
+            f, g = self._f_g(x, params, closure)
+            if self.method is not None and self.method.lower() == 'slsqp': g = g.astype(np.float64) #  slsqp requires float64
+            return f, g
+
+        return self._f(x, params, closure)
+
+    def _hess(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        f,g,H = self._f_g_H(x, params, closure)
+        return H
+
+    def _hessp(self, x: np.ndarray, p:np.ndarray, params: list[torch.Tensor], closure):
+        f,g,Hvp = self._f_g_Hvp(x, p, params, closure)
+        return Hvp
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
+        bounds = self._get_bounds()
+
+        # determine hess argument
+        hess = self.hess
+        hessp = None
+        if hess == 'autograd':
+            if self.use_hess_autograd:
+                if self.use_hessp:
+                    hessp = partial(self._hessp, params=params, closure=closure)
+                    hess = None
+                else:
+                    hess = partial(self._hess, params=params, closure=closure)
+            # hess = 'autograd' but method doesn't use hess
+            else:
+                hess = None
+
+
+        if self.method is not None and (self.method.lower() == 'tnc' or self.method.lower() == 'slsqp'):
+            x0 = x0.astype(np.float64) # those methods error without this
+
+        minimizer_kwargs = self._minimizer_kwargs.copy() if self._minimizer_kwargs is not None else {}
+        minimizer_kwargs.setdefault("method", self.method)
+        minimizer_kwargs.setdefault("jac", self.jac)
+        minimizer_kwargs.setdefault("hess", hess)
+        minimizer_kwargs.setdefault("hessp", hessp)
+        minimizer_kwargs.setdefault("bounds", bounds)
+
+        res = scipy.optimize.basinhopping(
+            partial(self._objective, params = params, closure = closure),
+            x0 = params.to_vec().numpy(force=True),
+            minimizer_kwargs=minimizer_kwargs,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
+        return res.fun

torchzero/optim/wrappers/scipy/brute.py

@@ -0,0 +1,48 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import scipy.optimize
+import torch
+
+from ....utils import TensorList
+from ..wrapper import WrapperBase
+
+Closure = Callable[[bool], Any]
+
+
+
+class ScipyBrute(WrapperBase):
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        Ns: int = 20,
+        finish = scipy.optimize.fmin,
+        disp: bool = False,
+        workers: int = 1
+    ):
+        super().__init__(params, dict(lb=lb,  ub=ub))
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = TensorList(self._get_params())
+        bounds = self._get_bounds()
+        assert bounds is not None
+
+        res,fval,grid,Jout = scipy.optimize.brute(
+            partial(self._f, params = params, closure = closure),
+            ranges=bounds,
+            full_output=True,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.as_tensor(res, device = params[0].device, dtype=params[0].dtype))
+
+        return fval

torchzero/optim/wrappers/scipy/differential_evolution.py

@@ -0,0 +1,80 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import scipy.optimize
+import torch
+
+from ....utils import TensorList
+from ..wrapper import WrapperBase
+
+Closure = Callable[[bool], Any]
+
+
+
+
+
+class ScipyDE(WrapperBase):
+    """Use scipy.minimize.differential_evolution as pytorch optimizer. Note that this performs full minimization on each step,
+    so usually you would want to perform a single step. This also requires bounds to be specified.
+
+    Please refer to https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.differential_evolution.html
+    for all other args.
+
+    Args:
+        params: iterable of parameters to optimize or dicts defining parameter groups.
+        bounds (tuple[float,float], optional): tuple with lower and upper bounds.
+            DE requires bounds to be specified. Defaults to None.
+
+        other args:
+            refer to https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.differential_evolution.html
+    """
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        strategy: Literal['best1bin', 'best1exp', 'rand1bin', 'rand1exp', 'rand2bin', 'rand2exp',
+            'randtobest1bin', 'randtobest1exp', 'currenttobest1bin', 'currenttobest1exp',
+            'best2exp', 'best2bin'] = 'best1bin',
+        maxiter: int = 1000,
+        popsize: int = 15,
+        tol: float = 0.01,
+        mutation = (0.5, 1),
+        recombination: float = 0.7,
+        seed = None,
+        callback = None,
+        disp: bool = False,
+        polish: bool = True,
+        init: str = 'latinhypercube',
+        atol: int = 0,
+        updating: str = 'immediate',
+        workers: int = 1,
+        constraints = (),
+        *,
+        integrality = None,
+
+    ):
+        super().__init__(params, dict(lb=lb, ub=ub))
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    @torch.no_grad
+    def step(self, closure: Closure): # pylint:disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
+        bounds = self._get_bounds()
+        assert bounds is not None
+
+        res = scipy.optimize.differential_evolution(
+            partial(self._f, params = params, closure = closure),
+            x0 = x0,
+            bounds=bounds,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
+        return res.fun

torchzero/optim/wrappers/scipy/direct.py

@@ -0,0 +1,69 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import scipy.optimize
+import torch
+
+from ....utils import TensorList
+from ..wrapper import WrapperBase
+
+Closure = Callable[[bool], Any]
+
+
+
+
+class ScipyDIRECT(WrapperBase):
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        maxfun: int | None = 1000,
+        maxiter: int = 1000,
+        eps: float = 0.0001,
+        locally_biased: bool = True,
+        f_min: float = -np.inf,
+        f_min_rtol: float = 0.0001,
+        vol_tol: float = 1e-16,
+        len_tol: float = 0.000001,
+        callback = None,
+    ):
+        super().__init__(params, dict(lb=lb, ub=ub))
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    def _objective(self, x: np.ndarray, params: TensorList, closure) -> float:
+        if self.raised: return np.inf
+        try:
+            return self._f(x, params, closure)
+
+        except Exception as e:
+            # this makes exceptions work in fcmaes and scipy direct
+            self.e = e
+            self.raised = True
+            return np.inf
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        self.raised = False
+        self.e = None
+
+        params = TensorList(self._get_params())
+        bounds = self._get_bounds()
+        assert bounds is not None
+
+        res = scipy.optimize.direct(
+            partial(self._objective, params=params, closure=closure),
+            bounds=bounds,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
+
+        if self.e is not None: raise self.e from None
+        return res.fun
+

torchzero/optim/wrappers/scipy/dual_annealing.py

@@ -0,0 +1,115 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import scipy.optimize
+import torch
+
+from ....utils import TensorList
+from ..wrapper import WrapperBase
+from .minimize import _use_jac_hess_hessp
+
+Closure = Callable[[bool], Any]
+
+
+
+
+class ScipyDualAnnealing(WrapperBase):
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        maxiter=1000,
+        minimizer_kwargs=None,
+        initial_temp=5230.0,
+        restart_temp_ratio=2.0e-5,
+        visit=2.62,
+        accept=-5.0,
+        maxfun=1e7,
+        rng=None,
+        no_local_search=False,
+        method: Literal['nelder-mead', 'powell', 'cg', 'bfgs', 'newton-cg',
+                    'l-bfgs-b', 'tnc', 'cobyla', 'cobyqa', 'slsqp',
+                    'trust-constr', 'dogleg', 'trust-ncg', 'trust-exact',
+                    'trust-krylov'] | str = 'l-bfgs-b',
+        jac: Literal['2-point', '3-point', 'cs', 'autograd'] = 'autograd',
+        hess: Literal['2-point', '3-point', 'cs', 'autograd'] | scipy.optimize.HessianUpdateStrategy = 'autograd',
+        use_hessp: bool = True,
+    ):
+        super().__init__(params, dict(lb=lb, ub=ub))
+
+        kwargs = locals().copy()
+        for k in ["self", "params", "lb", "ub", "__class__", "method", "jac", "hess", "use_hessp", "minimizer_kwargs"]:
+            del kwargs[k]
+        self._kwargs = kwargs
+
+        self._minimizer_kwargs = minimizer_kwargs
+        self.method = method
+        self.hess = hess
+
+        self.jac, self.use_jac_autograd, self.use_hess_autograd, self.use_hessp = _use_jac_hess_hessp(method, jac, hess, use_hessp)
+
+    def _jac(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        f,g = self._f_g(x, params, closure)
+        return g
+
+    def _objective(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        # dual annealing doesn't support this
+        # if self.use_jac_autograd:
+        #     f, g = self._f_g(x, params, closure)
+        #     if self.method.lower() == 'slsqp': g = g.astype(np.float64) #  slsqp requires float64
+        #     return f, g
+
+        return self._f(x, params, closure)
+
+    def _hess(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        f,g,H = self._f_g_H(x, params, closure)
+        return H
+
+    def _hessp(self, x: np.ndarray, p:np.ndarray, params: list[torch.Tensor], closure):
+        f,g,Hvp = self._f_g_Hvp(x, p, params, closure)
+        return Hvp
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
+        bounds = self._get_bounds()
+        assert bounds is not None
+
+        # determine hess argument
+        hess = self.hess
+        hessp = None
+        if hess == 'autograd':
+            if self.use_hess_autograd:
+                if self.use_hessp:
+                    hessp = partial(self._hessp, params=params, closure=closure)
+                    hess = None
+                else:
+                    hess = partial(self._hess, params=params, closure=closure)
+            # hess = 'autograd' but method doesn't use hess
+            else:
+                hess = None
+
+        if self.method.lower() in ('tnc', 'slsqp'):
+            x0 = x0.astype(np.float64) # those methods error without this
+
+        minimizer_kwargs = self._minimizer_kwargs.copy() if self._minimizer_kwargs is not None else {}
+        minimizer_kwargs.setdefault("method", self.method)
+        minimizer_kwargs.setdefault("jac", partial(self._jac, params = params, closure = closure))
+        minimizer_kwargs.setdefault("hess", hess)
+        minimizer_kwargs.setdefault("hessp", hessp)
+        minimizer_kwargs.setdefault("bounds", bounds)
+
+        res = scipy.optimize.dual_annealing(
+            partial(self._f, params = params, closure = closure),
+            x0 = x0,
+            bounds=bounds,
+            minimizer_kwargs=minimizer_kwargs,
+            **self._kwargs,
+        )
+
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
+        return res.fun