torchzero 0.3.10__py3-none-any.whl → 0.3.11__py3-none-any.whl

torchzero/optim/wrappers/directsearch.py

@@ -33,8 +33,45 @@ class DirectSearch(Optimizer):
     solution.
 
     Args:
-        params (_type_): _description_
-        maxevals (_type_, optional): _description_. Defaults to DEFAULT_PARAMS['maxevals'].
+        params: iterable of parameters to optimize or dicts defining parameter groups.
+
+        rho: Choice of the forcing function.
+
+        sketch_dim: Reduced dimension to generate polling directions in.
+
+        sketch_type: Sketching technique to be used.
+
+        maxevals: Maximum number of calls to f performed by the algorithm.
+
+        poll_type: Type of polling directions generated in the reduced spaces.
+
+        alpha0: Initial value for the stepsize parameter.
+
+        alpha_max: Maximum value for the stepsize parameter.
+
+        alpha_min: Minimum value for the stepsize parameter.
+
+        gamma_inc: Increase factor for the stepsize update.
+
+        gamma_dec: Decrease factor for the stepsize update.
+
+        verbose:
+            Boolean indicating whether information should be displayed during an algorithmic run.
+
+        print_freq:
+            Value indicating how frequently information should be displayed.
+
+        use_stochastic_three_points:
+            Boolean indicating whether the specific stochastic three points method should be used.
+
+        poll_scale_prob: Probability of scaling the polling directions.
+
+        poll_scale_factor: Factor used to scale the polling directions.
+
+        rho_uses_normd:
+            Boolean indicating whether the forcing function should account for the norm of the direction.
+
+
     """
     def __init__(
         self,

torchzero/optim/wrappers/fcmaes.py

@@ -27,18 +27,25 @@ class FcmaesWrapper(Optimizer):
     Note that this performs full minimization on each step, so only perform one step with this.
 
     Args:
-        params (_type_): _description_
-        lb (float): _description_
-        ub (float): _description_
-        optimizer (fcmaes.optimizer.Optimizer | None, optional): _description_. Defaults to None.
-        max_evaluations (int | None, optional): _description_. Defaults to 50000.
-        value_limit (float | None, optional): _description_. Defaults to np.inf.
-        num_retries (int | None, optional): _description_. Defaults to 1.
-        workers (int, optional): _description_. Defaults to 1.
-        popsize (int | None, optional): _description_. Defaults to 31.
-        capacity (int | None, optional): _description_. Defaults to 500.
-        stop_fitness (float | None, optional): _description_. Defaults to -np.inf.
-        statistic_num (int | None, optional): _description_. Defaults to 0.
+        params: iterable of parameters to optimize or dicts defining parameter groups.
+        lb (float): lower bounds, this can also be specified in param_groups.
+        ub (float): upper bounds, this can also be specified in param_groups.
+        optimizer (fcmaes.optimizer.Optimizer | None, optional):
+            optimizer to use. Default is a sequence of differential evolution and CMA-ES.
+        max_evaluations (int | None, optional):
+            Forced termination of all optimization runs after `max_evaluations` function evaluations.
+            Only used if optimizer is undefined, otherwise this setting is defined in the optimizer. Defaults to 50000.
+        value_limit (float | None, optional): Upper limit for optimized function values to be stored. Defaults to np.inf.
+        num_retries (int | None, optional): Number of optimization retries. Defaults to 1.
+        popsize (int | None, optional):
+            CMA-ES population size used for all CMA-ES runs.
+            Not used for differential evolution.
+            Ignored if parameter optimizer is defined. Defaults to 31.
+        capacity (int | None, optional): capacity of the evaluation store.. Defaults to 500.
+        stop_fitness (float | None, optional):
+            Limit for fitness value. optimization runs terminate if this value is reached. Defaults to -np.inf.
+        statistic_num (int | None, optional):
+            if > 0 stores the progress of the optimization. Defines the size of this store. Defaults to 0.
     """
     def __init__(
         self,
@@ -49,7 +56,7 @@ class FcmaesWrapper(Optimizer):
         max_evaluations: int | None = 50000,
         value_limit: float | None = np.inf,
         num_retries: int | None = 1,
-        workers: int = 1,
+        # workers: int = 1,
         popsize: int | None = 31,
         capacity: int | None = 500,
         stop_fitness: float | None = -np.inf,
@@ -60,6 +67,7 @@ class FcmaesWrapper(Optimizer):
         kwargs = locals().copy()
         del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__']
         self._kwargs = kwargs
+        self._kwargs['workers'] = 1
 
     def _objective(self, x: np.ndarray, params: TensorList, closure) -> float:
         if self.raised: return np.inf

torchzero/optim/wrappers/mads.py

@@ -31,16 +31,15 @@ class MADS(Optimizer):
     solution.
 
     Args:
-        params (params): params
-        lb (float): lower bounds
-        ub (float): upper bounds
+        params: iterable of parameters to optimize or dicts defining parameter groups.
+        lb (float): lower bounds, this can also be specified in param_groups.
+        ub (float): upper bounds, this can also be specified in param_groups.
         dp (float, optional): Initial poll size as percent of bounds. Defaults to 0.1.
         dm (float, optional): Initial mesh size as percent of bounds. Defaults to 0.01.
-        dp_tol (_type_, optional): Minimum poll size stopping criteria. Defaults to -float('inf').
-        nitermax (_type_, optional): Maximum objective function evaluations. Defaults to float('inf').
+        dp_tol (float, optional): Minimum poll size stopping criteria. Defaults to -float('inf').
+        nitermax (float, optional): Maximum objective function evaluations. Defaults to float('inf').
         displog (bool, optional): whether to show log. Defaults to False.
         savelog (bool, optional): whether to save log. Defaults to False.
-
     """
     def __init__(
         self,

torchzero/optim/wrappers/nevergrad.py

@@ -29,6 +29,12 @@ class NevergradWrapper(Optimizer):
             use certain rule for first 50% of the steps, and then switch to another rule.
             This parameter doesn't actually limit the maximum number of steps!
             But it doesn't have to be exact. Defaults to None.
+        lb (float | None, optional):
+            lower bounds, this can also be specified in param_groups. Bounds are optional, however
+            some nevergrad algorithms will raise an exception of bounds are not specified.
+        ub (float, optional):
+            upper bounds, this can also be specified in param_groups. Bounds are optional, however
+            some nevergrad algorithms will raise an exception of bounds are not specified.
         mutable_sigma (bool, optional):
             nevergrad parameter, sets whether the mutation standard deviation must mutate as well
             (for mutation based algorithms). Defaults to False.
@@ -44,11 +50,20 @@ class NevergradWrapper(Optimizer):
         params,
         opt_cls:"type[ng.optimizers.base.Optimizer] | abc.Callable[..., ng.optimizers.base.Optimizer]",
         budget: int | None = None,
-        mutable_sigma = False,
         lb: float | None = None,
         ub: float | None = None,
+        mutable_sigma = False,
         use_init = True,
     ):
+        """_summary_
+
+        Args:
+            params (_type_): _description_
+            opt_cls (type[ng.optimizers.base.Optimizer] | abc.Callable[..., ng.optimizers.base.Optimizer]): _description_
+            budget (int | None, optional): _description_. Defaults to None.
+            mutable_sigma (bool, optional): _description_. Defaults to False.
+            use_init (bool, optional): _description_. Defaults to True.
+        """
         defaults = dict(lb=lb, ub=ub, use_init=use_init, mutable_sigma=mutable_sigma)
         super().__init__(params, defaults)
         self.opt_cls = opt_cls

torchzero/optim/wrappers/optuna.py

@@ -23,7 +23,7 @@ class OptunaSampler(Optimizer):
     Note - optuna is surprisingly scalable to large number of parameters (up to 10,000), despite literally requiring a for-loop because it only supports scalars. Default TPESampler is good for BBO. Maybe not for NNs...
 
     Args:
-        params (_type_): parameters
+        params: iterable of parameters to optimize or dicts defining parameter groups.
         lb (float): lower bounds.
         ub (float): upper bounds.
         sampler (optuna.samplers.BaseSampler | type[optuna.samplers.BaseSampler] | None, optional): sampler. Defaults to None.

torchzero/optim/wrappers/scipy.py

@@ -139,9 +139,11 @@ class ScipyMinimize(Optimizer):
 
         # make bounds
         lb, ub = self.group_vals('lb', 'ub', cls=list)
-        bounds = []
-        for p, l, u in zip(params, lb, ub):
-            bounds.extend([(l, u)] * p.numel())
+        bounds = None
+        if any(b is not None for b in lb) or any(b is not None for b in ub):
+            bounds = []
+            for p, l, u in zip(params, lb, ub):
+                bounds.extend([(l, u)] * p.numel())
 
         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

torchzero/utils/__init__.py

@@ -18,6 +18,6 @@ from .params import (
     _copy_param_groups,
     _make_param_groups,
 )
-from .python_tools import flatten, generic_eq, reduce_dim, unpack_dicts
-from .tensorlist import TensorList, as_tensorlist, Distributions, generic_clamp, generic_numel, generic_vector_norm, generic_zeros_like, generic_randn_like
+from .python_tools import flatten, generic_eq, generic_ne, reduce_dim, unpack_dicts
+from .tensorlist import TensorList, as_tensorlist, Distributions, generic_clamp, generic_numel, generic_vector_norm, generic_zeros_like, generic_randn_like, generic_finfo_eps
 from .torch_tools import tofloat, tolist, tonumpy, totensor, vec_to_tensors, vec_to_tensors_, set_storage_

torchzero/utils/derivatives.py

@@ -158,7 +158,7 @@ def hessian_mat(
     method="func",
     vectorize=False,
     outer_jacobian_strategy="reverse-mode",
-):
+) -> torch.Tensor:
     """
     returns hessian matrix for parameters (as if they were flattened and concatenated into a vector).
 
@@ -190,7 +190,7 @@ def hessian_mat(
         return loss
 
     if method == 'func':
-        return torch.func.hessian(func)(torch.cat([p.view(-1) for p in params]).detach().requires_grad_(create_graph))
+        return torch.func.hessian(func)(torch.cat([p.view(-1) for p in params]).detach().requires_grad_(create_graph)) # pyright:ignore[reportReturnType]
 
     if method == 'autograd.functional':
         return torch.autograd.functional.hessian(
@@ -199,7 +199,7 @@ def hessian_mat(
             create_graph=create_graph,
             vectorize=vectorize,
             outer_jacobian_strategy=outer_jacobian_strategy,
-        )
+        ) # pyright:ignore[reportReturnType]
     raise ValueError(method)
 
 def jvp(fn, params: Iterable[torch.Tensor], tangent: Iterable[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:

torchzero/utils/linalg/__init__.py

@@ -2,4 +2,4 @@ from .matrix_funcs import inv_sqrt_2x2, eigvals_func, singular_vals_func, matrix
 from .orthogonalize import gram_schmidt
 from .qr import qr_householder
 from .svd import randomized_svd
-from .solve import cg, nystrom_approximation, nystrom_sketch_and_solve
+from .solve import cg, nystrom_approximation, nystrom_sketch_and_solve, steihaug_toint_cg

torchzero/utils/linalg/solve.py

@@ -1,12 +1,41 @@
+# pyright: reportArgumentType=false
 from collections.abc import Callable
-from typing import overload
+from typing import Any, overload
+
 import torch
 
-from .. import TensorList, generic_zeros_like, generic_vector_norm, generic_numel, generic_randn_like, generic_eq
+from .. import (
+    TensorList,
+    generic_eq,
+    generic_finfo_eps,
+    generic_numel,
+    generic_randn_like,
+    generic_vector_norm,
+    generic_zeros_like,
+)
+
+
+def _make_A_mm_reg(A_mm: Callable | torch.Tensor, reg):
+    if callable(A_mm):
+        def A_mm_reg(x): # A_mm with regularization
+            Ax = A_mm(x)
+            if not generic_eq(reg, 0): Ax += x*reg
+            return Ax
+        return A_mm_reg
+
+    if not isinstance(A_mm, torch.Tensor): raise TypeError(type(A_mm))
+
+    def Ax_reg(x): # A_mm with regularization
+        if A_mm.ndim == 1: Ax = A_mm * x
+        else: Ax = A_mm @ x
+        if reg != 0: Ax += x*reg
+        return Ax
+    return Ax_reg
+
 
 @overload
 def cg(
-    A_mm: Callable[[torch.Tensor], torch.Tensor],
+    A_mm: Callable[[torch.Tensor], torch.Tensor] | torch.Tensor,
     b: torch.Tensor,
     x0_: torch.Tensor | None = None,
     tol: float | None = 1e-4,
@@ -24,17 +53,17 @@ def cg(
 ) -> TensorList: ...
 
 def cg(
-    A_mm: Callable,
+    A_mm: Callable | torch.Tensor,
     b: torch.Tensor | TensorList,
     x0_: torch.Tensor | TensorList | None = None,
     tol: float | None = 1e-4,
     maxiter: int | None = None,
     reg: float | list[float] | tuple[float] = 0,
 ):
-    def A_mm_reg(x): # A_mm with regularization
-        Ax = A_mm(x)
-        if not generic_eq(reg, 0): Ax += x*reg
-        return Ax
+    A_mm_reg = _make_A_mm_reg(A_mm, reg)
+    eps = generic_finfo_eps(b)
+
+    if tol is None: tol = eps
 
     if maxiter is None: maxiter = generic_numel(b)
     if x0_ is None: x0_ = generic_zeros_like(b)
@@ -44,9 +73,10 @@ def cg(
     p = residual.clone() # search direction
     r_norm = generic_vector_norm(residual)
     init_norm = r_norm
-    if tol is not None and r_norm < tol: return x
+    if r_norm < tol: return x
     k = 0
 
+
     while True:
         Ap = A_mm_reg(p)
         step_size = (r_norm**2) / p.dot(Ap)
@@ -55,7 +85,7 @@ def cg(
         new_r_norm = generic_vector_norm(residual)
 
         k += 1
-        if tol is not None and new_r_norm <= tol * init_norm: return x
+        if new_r_norm <= tol * init_norm: return x
         if k >= maxiter: return x
 
         beta = (new_r_norm**2) / (r_norm**2)
@@ -131,6 +161,8 @@ def nystrom_pcg(
         generator=generator,
     )
     lambd += reg
+    eps = torch.finfo(b.dtype).eps ** 2
+    if tol is None: tol = eps
 
     def A_mm_reg(x): # A_mm with regularization
         Ax = A_mm(x)
@@ -150,7 +182,7 @@ def nystrom_pcg(
     p = z.clone() # search direction
 
     init_norm = torch.linalg.vector_norm(residual) # pylint:disable=not-callable
-    if tol is not None and init_norm < tol: return x
+    if init_norm < tol: return x
     k = 0
     while True:
         Ap = A_mm_reg(p)
@@ -160,10 +192,217 @@ def nystrom_pcg(
         residual -= step_size * Ap
 
         k += 1
-        if tol is not None and torch.linalg.vector_norm(residual) <= tol * init_norm: return x # pylint:disable=not-callable
+        if torch.linalg.vector_norm(residual) <= tol * init_norm: return x # pylint:disable=not-callable
         if k >= maxiter: return x
 
         z = P_inv @ residual
         beta = residual.dot(z) / rz
         p = z + p*beta
 
+
+def _safe_clip(x: torch.Tensor):
+    """makes sure scalar tensor x is not smaller than epsilon"""
+    assert x.numel() == 1, x.shape
+    eps = torch.finfo(x.dtype).eps
+    if x.abs() < eps: return x.new_full(x.size(), eps).copysign(x)
+    return x
+
+def _trust_tau(x,d,trust_region):
+    xx = x.dot(x)
+    xd = x.dot(d)
+    dd = _safe_clip(d.dot(d))
+
+    rad = (xd**2 - dd * (xx - trust_region**2)).clip(min=0).sqrt()
+    tau = (-xd + rad) / dd
+
+    return x + tau * d
+
+
+@overload
+def steihaug_toint_cg(
+    A_mm: Callable[[torch.Tensor], torch.Tensor] | torch.Tensor,
+    b: torch.Tensor,
+    trust_region: float,
+    x0: torch.Tensor | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float = 0,
+) -> torch.Tensor: ...
+@overload
+def steihaug_toint_cg(
+    A_mm: Callable[[TensorList], TensorList],
+    b: TensorList,
+    trust_region: float,
+    x0: TensorList | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+) -> TensorList: ...
+def steihaug_toint_cg(
+    A_mm: Callable | torch.Tensor,
+    b: torch.Tensor | TensorList,
+    trust_region: float,
+    x0: torch.Tensor | TensorList | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+):
+    """
+    Solution is bounded to have L2 norm no larger than :code:`trust_region`. If solution exceeds :code:`trust_region`, CG is terminated early, so it is also faster.
+    """
+    A_mm_reg = _make_A_mm_reg(A_mm, reg)
+
+    x = x0
+    if x is None: x = generic_zeros_like(b)
+    r = b
+    d = r.clone()
+
+    eps = generic_finfo_eps(b)**2
+    if tol is None: tol = eps
+
+    if generic_vector_norm(r) < tol:
+        return x
+
+    if maxiter is None:
+        maxiter = generic_numel(b)
+
+    for _ in range(maxiter):
+        Ad = A_mm_reg(d)
+
+        d_Ad = d.dot(Ad)
+        if d_Ad <= eps:
+            return _trust_tau(x, d, trust_region)
+
+        alpha = r.dot(r) / d_Ad
+        p_next = x + alpha * d
+
+        # check if the step exceeds the trust-region boundary
+        if generic_vector_norm(p_next) >= trust_region:
+            return _trust_tau(x, d, trust_region)
+
+        # update step, residual and direction
+        x = p_next
+        r_next = r - alpha * Ad
+
+        if generic_vector_norm(r_next) < tol:
+            return x
+
+        beta = r_next.dot(r_next) / r.dot(r)
+        d = r_next + beta * d
+        r = r_next
+
+    return x
+
+
+
+# Liu, Yang, and Fred Roosta. "MINRES: From negative curvature detection to monotonicity properties." SIAM Journal on Optimization 32.4 (2022): 2636-2661.
+@overload
+def minres(
+    A_mm: Callable[[torch.Tensor], torch.Tensor] | torch.Tensor,
+    b: torch.Tensor,
+    x0: torch.Tensor | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float = 0,
+    npc_terminate: bool=True,
+    trust_region: float | None = None,
+) -> torch.Tensor: ...
+@overload
+def minres(
+    A_mm: Callable[[TensorList], TensorList],
+    b: TensorList,
+    x0: TensorList | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+    npc_terminate: bool=True,
+    trust_region: float | None = None,
+) -> TensorList: ...
+def minres(
+    A_mm,
+    b,
+    x0: torch.Tensor | TensorList | None = None,
+    tol: float | None = 1e-4,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+    npc_terminate: bool=True,
+    trust_region: float | None = None,
+):
+    A_mm_reg = _make_A_mm_reg(A_mm, reg)
+    eps = generic_finfo_eps(b)
+    if tol is None: tol = eps**2
+
+    if maxiter is None: maxiter = generic_numel(b)
+    if x0 is None:
+        R = b
+        x0 = generic_zeros_like(b)
+    else:
+        R = b - A_mm_reg(x0)
+
+    X: Any = x0
+    beta = b_norm = generic_vector_norm(b)
+    if b_norm < eps**2:
+        return generic_zeros_like(b)
+
+
+    V = b / beta
+    V_prev = generic_zeros_like(b)
+    D = generic_zeros_like(b)
+    D_prev = generic_zeros_like(b)
+
+    c = -1
+    phi = tau = beta
+    s = delta1 = e = 0
+
+
+    for _ in range(maxiter):
+
+        P = A_mm_reg(V)
+        alpha = V.dot(P)
+        P -= beta*V_prev
+        P -= alpha*V
+        beta = generic_vector_norm(P)
+
+        delta2 = c*delta1 + s*alpha
+        gamma1 = s*delta1 - c*alpha
+        e_next = s*beta
+        delta1 = -c*beta
+
+        cgamma1 = c*gamma1
+        if trust_region is not None and cgamma1 >= 0:
+            if npc_terminate: return _trust_tau(X, R, trust_region)
+            return _trust_tau(X, D, trust_region)
+
+        if npc_terminate and cgamma1 >= 0:
+            return R
+
+        gamma2 = (gamma1**2 + beta**2)**(1/2)
+
+        if abs(gamma2) <= eps: # singular system
+            # c=0; s=1; tau=0
+            if trust_region is None: return X
+            return _trust_tau(X, D, trust_region)
+
+        c = gamma1 / gamma2
+        s = beta/gamma2
+        tau = c*phi
+        phi = s*phi
+
+        D_prev = D
+        D = (V - delta2*D - e*D_prev) / gamma2
+        e = e_next
+        X = X + tau*D
+
+        if trust_region is not None:
+            if generic_vector_norm(X) > trust_region:
+                return _trust_tau(X, D, trust_region)
+
+        if (abs(beta) < eps) or (phi / b_norm <= tol):
+            # R = zeros(R)
+            return X
+
+        V_prev = V
+        V = P/beta
+        R = s**2*R - phi*c*V
+
+    return X

torchzero/utils/numberlist.py

@@ -129,4 +129,6 @@ class NumberList(list[int | float | Any]):
         return self.__class__(fn(i, *args, **kwargs) for i in self)
 
     def clamp(self, min=None, max=None):
+        return self.zipmap_args(_clamp, min, max)
+    def clip(self, min=None, max=None):
         return self.zipmap_args(_clamp, min, max)

torchzero/utils/python_tools.py

@@ -31,6 +31,16 @@ def generic_eq(x: int | float | Iterable[int | float], y: int | float | Iterable
         return all(i==y for i in x)
     return all(i==j for i,j in zip(x,y))
 
+def generic_ne(x: int | float | Iterable[int | float], y: int | float | Iterable[int | float]) -> bool:
+    """generic not equals function that supports scalars and lists of numbers. Faster than not generic_eq"""
+    if isinstance(x, (int,float)):
+        if isinstance(y, (int,float)): return x!=y
+        return any(i!=x for i in y)
+    if isinstance(y, (int,float)):
+        return any(i!=y for i in x)
+    return any(i!=j for i,j in zip(x,y))
+
+
 def zipmap(self, fn: Callable, other: Any | list | tuple, *args, **kwargs):
     """If `other` is list/tuple, applies `fn` to self zipped with `other`.
     Otherwise applies `fn` to this sequence and `other`.