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

torchzero/utils/linalg/solve.py

@@ -1,69 +1,32 @@
+# pyright: reportArgumentType=false
+import math
+from collections import deque
 from collections.abc import Callable
-from typing import overload
+from typing import Any, NamedTuple, 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_tiny,
+    generic_numel,
+    generic_vector_norm,
+    generic_zeros_like,
+)
 
-@overload
-def cg(
-    A_mm: Callable[[torch.Tensor], torch.Tensor],
-    b: torch.Tensor,
-    x0_: torch.Tensor | None = None,
-    tol: float | None = 1e-4,
-    maxiter: int | None = None,
-    reg: float = 0,
-) -> torch.Tensor: ...
-@overload
-def cg(
-    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,
-) -> TensorList: ...
 
-def cg(
-    A_mm: Callable,
-    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 _make_A_mm_reg(A_mm: Callable, reg):
     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 maxiter is None: maxiter = generic_numel(b)
-    if x0_ is None: x0_ = generic_zeros_like(b)
-
-    x = x0_
-    residual = b - A_mm_reg(x)
-    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
-    k = 0
-
-    while True:
-        Ap = A_mm_reg(p)
-        step_size = (r_norm**2) / p.dot(Ap)
-        x += step_size * p # Update solution
-        residual -= step_size * Ap # Update residual
-        new_r_norm = generic_vector_norm(residual)
-
-        k += 1
-        if tol is not None and new_r_norm <= tol * init_norm: return x
-        if k >= maxiter: return x
-
-        beta = (new_r_norm**2) / (r_norm**2)
-        p = residual + beta*p
-        r_norm = new_r_norm
+def _identity(x): return x
 
 
-# https://arxiv.org/pdf/2110.02820 algorithm 2.1 apparently supposed to be diabolical
+# https://arxiv.org/pdf/2110.02820
 def nystrom_approximation(
     A_mm: Callable[[torch.Tensor], torch.Tensor],
     ndim: int,
@@ -85,7 +48,6 @@ def nystrom_approximation(
     lambd = (S.pow(2) - v).clip(min=0) #Remove shift, compute eigs
     return U, lambd
 
-# this one works worse
 def nystrom_sketch_and_solve(
     A_mm: Callable[[torch.Tensor], torch.Tensor],
     b: torch.Tensor,
@@ -111,7 +73,6 @@ def nystrom_sketch_and_solve(
     term2 = (1.0 / reg) * (b - U @ Uᵀb)
     return (term1 + term2).squeeze(-1)
 
-# this one is insane
 def nystrom_pcg(
     A_mm: Callable[[torch.Tensor], torch.Tensor],
     b: torch.Tensor,
@@ -131,6 +92,8 @@ def nystrom_pcg(
         generator=generator,
     )
     lambd += reg
+    eps = torch.finfo(b.dtype).tiny * 2
+    if tol is None: tol = eps
 
     def A_mm_reg(x): # A_mm with regularization
         Ax = A_mm(x)
@@ -150,7 +113,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 +123,358 @@ 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 tiny"""
+    assert x.numel() == 1, x.shape
+    eps = torch.finfo(x.dtype).tiny * 2
+    if x.abs() < eps: return x.new_full(x.size(), eps).copysign(x)
+    return x
+
+def _trust_tau(x,d,trust_radius):
+    xx = x.dot(x)
+    xd = x.dot(d)
+    dd = _safe_clip(d.dot(d))
+
+    rad = (xd**2 - dd * (xx - trust_radius**2)).clip(min=0).sqrt()
+    tau = (-xd + rad) / dd
+
+    return x + tau * d
+
+
+class CG:
+    """Conjugate gradient method.
+
+    Args:
+        A_mm (Callable[[torch.Tensor], torch.Tensor] | torch.Tensor): Callable that returns matvec ``Ax``.
+        b (torch.Tensor): right hand side
+        x0 (torch.Tensor | None, optional): initial guess, defaults to zeros. Defaults to None.
+        tol (float | None, optional): tolerance for convergence. Defaults to 1e-8.
+        maxiter (int | None, optional):
+            maximum number of iterations, if None sets to number of dimensions. Defaults to None.
+        reg (float, optional): regularization. Defaults to 0.
+        trust_radius (float | None, optional):
+            CG is terminated whenever solution exceeds trust region, returning a solution modified to be within it. Defaults to None.
+        npc_terminate (bool, optional):
+            whether to terminate CG whenever negative curavture is detected. Defaults to False.
+        miniter (int, optional):
+            minimal number of iterations even if tolerance is satisfied, this ensures some progress
+            is always made.
+        history_size (int, optional):
+            number of past iterations to store, to re-use them when trust radius is decreased.
+        P_mm (Callable | torch.Tensor | None, optional):
+            Callable that returns inverse preconditioner times vector. Defaults to None.
+    """
+    def __init__(
+        self,
+        A_mm: Callable,
+        b: torch.Tensor | TensorList,
+        x0: torch.Tensor | TensorList | None = None,
+        tol: float | None = 1e-4,
+        maxiter: int | None = None,
+        reg: float = 0,
+        trust_radius: float | None = None,
+        npc_terminate: bool=False,
+        miniter: int = 0,
+        history_size: int = 0,
+        P_mm: Callable | None = None,
+):
+        # --------------------------------- set attrs -------------------------------- #
+        self.A_mm = _make_A_mm_reg(A_mm, reg)
+        self.b = b
+        if tol is None: tol = generic_finfo_tiny(b) * 2
+        self.tol = tol
+        self.eps = generic_finfo_tiny(b) * 2
+        if maxiter is None: maxiter = generic_numel(b)
+        self.maxiter = maxiter
+        self.miniter = miniter
+        self.trust_radius = trust_radius
+        self.npc_terminate = npc_terminate
+        self.P_mm = P_mm if P_mm is not None else _identity
+
+        if history_size > 0:
+            self.history = deque(maxlen = history_size)
+            """history of (x, x_norm, d)"""
+        else:
+            self.history = None
+
+        # -------------------------------- initialize -------------------------------- #
+
+        self.iter = 0
+
+        if x0 is None:
+            self.x = generic_zeros_like(b)
+            self.r = b
+        else:
+            self.x = x0
+            self.r = b - A_mm(self.x)
+
+        self.z = self.P_mm(self.r)
+        self.d = self.z
+
+        if self.history is not None:
+            self.history.append((self.x, generic_vector_norm(self.x), self.d))
+
+    def step(self) -> tuple[Any, bool]:
+        """returns ``(solution, should_terminate)``"""
+        x, b, d, r, z = self.x, self.b, self.d, self.r, self.z
+
+        if self.iter >= self.maxiter:
+            return x, True
+
+        Ad = self.A_mm(d)
+        dAd = d.dot(Ad)
+
+        # check negative curvature
+        if dAd <= self.eps:
+            if self.trust_radius is not None: return _trust_tau(x, d, self.trust_radius), True
+            if self.iter == 0: return b * (b.dot(b) / dAd).abs(), True
+            if self.npc_terminate: return x, True
+
+        rz = r.dot(z)
+        alpha = rz / dAd
+        x_next = x + alpha * d
+
+        # check if the step exceeds the trust-region boundary
+        x_next_norm = None
+        if self.trust_radius is not None:
+            x_next_norm = generic_vector_norm(x_next)
+            if x_next_norm >= self.trust_radius:
+                return _trust_tau(x, d, self.trust_radius), True
+
+        # update step, residual and direction
+        r_next = r - alpha * Ad
+
+        # check if r is sufficiently small
+        if self.iter >= self.miniter and generic_vector_norm(r_next) < self.tol:
+            return x_next, True
+
+        # update d, r, z
+        z_next = self.P_mm(r_next)
+        beta = r_next.dot(z_next) / rz
+
+        self.d = z_next + beta * d
+        self.x = x_next
+        self.r = r_next
+        self.z = z_next
+
+        # update history
+        if self.history is not None:
+            if x_next_norm is None: x_next_norm = generic_vector_norm(x_next)
+            self.history.append((self.x, x_next_norm, self.d))
+
+        self.iter += 1
+        return x, False
+
+
+    def solve(self):
+        # return initial guess if it is good enough
+        if self.miniter < 1 and generic_vector_norm(self.r) < self.tol:
+            return self.x
+
+        should_terminate = False
+        sol = None
+
+        while not should_terminate:
+            sol, should_terminate = self.step()
+
+        assert sol is not None
+        return sol
+
+def find_within_trust_radius(history, trust_radius: float):
+    """find first ``x`` in history that exceeds trust radius, if no such ``x`` exists, returns ``None``"""
+    for x, x_norm, d in reversed(tuple(history)):
+        if x_norm <= trust_radius:
+            return _trust_tau(x, d, trust_radius)
+    return None
+
+class _TensorSolution(NamedTuple):
+    x: torch.Tensor
+    solver: CG
+
+class _TensorListSolution(NamedTuple):
+    x: TensorList
+    solver: CG
+
+
+@overload
+def cg(
+    A_mm: Callable[[torch.Tensor], torch.Tensor],
+    b: torch.Tensor,
+    x0: torch.Tensor | None = None,
+    tol: float | None = 1e-8,
+    maxiter: int | None = None,
+    reg: float = 0,
+    trust_radius: float | None = None,
+    npc_terminate: bool = False,
+    miniter: int = 0,
+    history_size: int = 0,
+    P_mm: Callable[[torch.Tensor], torch.Tensor] | None = None
+) -> _TensorSolution: ...
+@overload
+def cg(
+    A_mm: Callable[[TensorList], TensorList],
+    b: TensorList,
+    x0: TensorList | None = None,
+    tol: float | None = 1e-8,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+    trust_radius: float | None = None,
+    npc_terminate: bool=False,
+    miniter: int = 0,
+    history_size: int = 0,
+    P_mm: Callable[[TensorList], TensorList] | None = None
+) -> _TensorListSolution: ...
+def cg(
+    A_mm: Callable,
+    b: torch.Tensor | TensorList,
+    x0: torch.Tensor | TensorList | None = None,
+    tol: float | None = 1e-8,
+    maxiter: int | None = None,
+    reg: float | list[float] | tuple[float] = 0,
+    trust_radius: float | None = None,
+    npc_terminate: bool = False,
+    miniter: int = 0,
+    history_size:int = 0,
+    P_mm: Callable | None = None
+):
+    solver = CG(
+        A_mm=A_mm,
+        b=b,
+        x0=x0,
+        tol=tol,
+        maxiter=maxiter,
+        reg=reg,
+        trust_radius=trust_radius,
+        npc_terminate=npc_terminate,
+        miniter=miniter,
+        history_size=history_size,
+        P_mm=P_mm,
+    )
+
+    x = solver.solve()
+
+    if isinstance(b, torch.Tensor):
+        return _TensorSolution(x, solver)
+
+    return _TensorListSolution(x, solver)
+
+
+# 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_radius: 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_radius: 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_radius: float | None = None, #trust region is experimental
+):
+    A_mm_reg = _make_A_mm_reg(A_mm, reg)
+    eps = math.sqrt(generic_finfo_tiny(b) * 2)
+    if tol is None: tol = eps
+
+    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_radius is not None and cgamma1 >= 0:
+            if npc_terminate: return _trust_tau(X, R, trust_radius)
+            return _trust_tau(X, D, trust_radius)
+
+        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_radius is None: return X
+            return _trust_tau(X, D, trust_radius)
+
+        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_radius is not None:
+            if generic_vector_norm(X) > trust_radius:
+                return _trust_tau(X, D, trust_radius)
+
+        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/metrics.py

@@ -0,0 +1,83 @@
+"""convenience submodule which allows to calculate a metric based on its string name,
+used in many places"""
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Literal, overload
+
+import torch
+
+if TYPE_CHECKING:
+    from .tensorlist import TensorList
+
+
+
+class Metric(ABC):
+    @abstractmethod
+    def evaluate_global(self, x: "TensorList") -> torch.Tensor:
+        """returns a global metric for a tensorlist"""
+
+    @abstractmethod
+    def evaluate_tensor(self, x: torch.Tensor, dim=None, keepdim=False) -> torch.Tensor:
+        """returns metric for a tensor"""
+
+    def evaluate_list(self, x: "TensorList") -> "TensorList":
+        """returns list of metrics for a tensorlist (possibly vectorized)"""
+        return x.map(self.evaluate_tensor)
+
+
+class _MAD(Metric):
+    def evaluate_global(self, x): return x.abs().global_mean()
+    def evaluate_tensor(self, x, dim=None, keepdim=False): return x.abs().mean(dim=dim, keepdim=keepdim)
+    def evaluate_list(self, x): return x.abs().mean()
+
+class _Std(Metric):
+    def evaluate_global(self, x): return x.global_std()
+    def evaluate_tensor(self, x, dim=None, keepdim=False): return x.std(dim=dim, keepdim=keepdim)
+    def evaluate_list(self, x): return x.std()
+
+class _Var(Metric):
+    def evaluate_global(self, x): return x.global_var()
+    def evaluate_tensor(self, x, dim=None, keepdim=False): return x.var(dim=dim, keepdim=keepdim)
+    def evaluate_list(self, x): return x.var()
+
+class _Sum(Metric):
+    def evaluate_global(self, x): return x.global_sum()
+    def evaluate_tensor(self, x, dim=None, keepdim=False): return x.sum(dim=dim, keepdim=keepdim)
+    def evaluate_list(self, x): return x.sum()
+
+class _Norm(Metric):
+    def __init__(self, ord): self.ord = ord
+    def evaluate_global(self, x): return x.global_vector_norm(self.ord)
+    def evaluate_tensor(self, x, dim=None, keepdim=False):
+        return torch.linalg.vector_norm(x, ord=self.ord, dim=dim, keepdim=keepdim) # pylint:disable=not-callable
+    def evaluate_list(self, x): return x.norm(self.ord)
+
+_METRIC_KEYS = Literal['mad', 'std', 'var', 'sum', 'l0', 'l1', 'l2', 'l3', 'l4', 'linf']
+_METRICS: dict[_METRIC_KEYS, Metric] = {
+    "mad": _MAD(),
+    "std": _Std(),
+    "var": _Var(),
+    "sum": _Sum(),
+    "l0": _Norm(0),
+    "l1": _Norm(1),
+    "l2": _Norm(2),
+    "l3": _Norm(3),
+    "l4": _Norm(4),
+    "linf": _Norm(torch.inf),
+}
+
+Metrics = _METRIC_KEYS | float | torch.Tensor
+def evaluate_metric(x: "torch.Tensor | TensorList", metric: Metrics) -> torch.Tensor:
+    if isinstance(metric, (int, float, torch.Tensor)):
+        if isinstance(x, torch.Tensor): return torch.linalg.vector_norm(x, ord=metric) # pylint:disable=not-callable
+        return x.global_vector_norm(ord=float(metric))
+
+    if isinstance(x, torch.Tensor): return _METRICS[metric].evaluate_tensor(x)
+    return _METRICS[metric].evaluate_global(x)
+
+
+def calculate_metric_list(x: "TensorList", metric: Metrics) -> "TensorList":
+    if isinstance(metric, (int, float, torch.Tensor)):
+        return x.norm(ord=float(metric))
+
+    return _METRICS[metric].evaluate_list(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`.
@@ -51,3 +61,9 @@ def unpack_dicts(dicts: Iterable[Mapping[str, Any]], key:str, key2: str | None =
     values = [cls(s[k] for s in dicts) for k in keys] # pyright:ignore[reportCallIssue]
     if len(values) == 1: return values[0]
     return values
+
+
+def safe_dict_update_(d1_:dict, d2:dict):
+    inter = set(d1_.keys()).intersection(d2.keys())
+    if len(inter) > 0: raise RuntimeError(f"Duplicate keys {inter}")
+    d1_.update(d2)