From 51119bfa04f0c3748f42c027c07eef4a3c30c37b Mon Sep 17 00:00:00 2001 From: mrava87 Date: Tue, 14 Apr 2026 20:57:57 +0100 Subject: [PATCH 01/24] feat: started to add class based primal solvers --- pyproximal/optimization/__init__.py | 3 +- pyproximal/optimization/basesolver.py | 136 +++++ pyproximal/optimization/cls_primal.py | 821 ++++++++++++++++++++++++++ 3 files changed, 959 insertions(+), 1 deletion(-) create mode 100644 pyproximal/optimization/basesolver.py create mode 100644 pyproximal/optimization/cls_primal.py diff --git a/pyproximal/optimization/__init__.py b/pyproximal/optimization/__init__.py index 694aac1..7505210 100644 --- a/pyproximal/optimization/__init__.py +++ b/pyproximal/optimization/__init__.py @@ -48,4 +48,5 @@ """ -from . import primal, primaldual, bregman, segmentation, sr3, palm, pnp +from . import cls_primal, primal +from . import primaldual, bregman, segmentation, sr3, palm, pnp diff --git a/pyproximal/optimization/basesolver.py b/pyproximal/optimization/basesolver.py new file mode 100644 index 0000000..d64d839 --- /dev/null +++ b/pyproximal/optimization/basesolver.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +__all__ = ["Solver"] + +import time +from abc import ABCMeta, abstractmethod +from typing import TYPE_CHECKING, Any + +from pylops.optimization.basesolver import Solver as pSolver +from pylops.optimization.callback import Callbacks + +if TYPE_CHECKING: + from pyproximal.ProxOperator import ProxOperator + +_units = {"B": 1, "KB": 1024, "MB": 1024**2, "GB": 1024**3} + + +class Solver(pSolver, metaclass=ABCMeta): + r"""Solver + + This is a template class which a user must subclass when implementing a new solver. + This class comprises of the following mandatory methods: + + - ``__init__``: initialization method the solver object is created and + callbacks (if any present) are registered + - ``memory_usage``: a method to compute upfront the memory used by each + step of the solver + - ``setup``: a method that is invoked to setup the solver, basically it will create + anything required prior to applying a step of the solver + - ``step``: a method applying a single step of the solver + - ``run``: a method applying multiple steps of the solver + - ``finalize``: a method that is invoked at the end of the optimization process. It can + be used to do some final clean-up of the properties of the operator that we want + to expose to the user + - ``solve``: a method applying the entire optimization loop of the solver for a + certain number of steps + + and optional methods: + + - ``_print_solver``: a method print on screen details of the solver (already implemented) + - ``_print_setup``: a method print on screen details of the setup process + - ``_print_step``: a method print on screen details of each step + - ``_print_finalize``: a method print on screen details of the finalize process + - ``callback``: a method implementing a callback function, which is called after + every step of the solver + + Parameters + ---------- + callbacks : :obj:`pylops.optimization.callback.Callbacks` + Callbacks object used to implement custom callbacks + + Attributes + ---------- + iiter : :obj:`int` + Iteration counter. + tstart : :obj:`float` + Time at the start of the optimization process. + tend : :obj:`float` + Time at the end of the optimization process. Available + only after ``finalize`` is called. + telapsed : :obj:`float` + Total time elapsed during the optimization process. Available + only after ``finalize`` is called. + + """ + + def __init__( + self, + callbacks: Callbacks = None, + ) -> None: + self.callbacks = callbacks + self._registercallbacks() + self.iiter = 0 + self.tstart = time.time() + + def _print_solver(self, text: str = "", nbar: int = 80) -> None: + print(f"{type(self).__name__}" + text) + print("-" * nbar) + + def _print_finalize(self, *args: Any, nbar: int = 80, **kwargs: Any) -> None: + print( + f"\nIterations = {self.iiter} Total time (s) = {self.telapsed:.2f}" + ) + print("-" * nbar + "\n") + + @abstractmethod + def setup( + self, + proxf: ProxOperator | list[ProxOperator], + proxg: ProxOperator | None = None, + *args: Any, + show: bool = False, + **kwargs: Any, + ) -> None: + """Setup solver + + This method is used to setup the solver. Users can change the function signature + by including any other input parameter required during the setup stage + + Parameters + ---------- + proxf : :obj:`pyproximal."ProxOperator"` or :obj:`list` + Proximal operator(s) to be used in the optimization + proxg : :obj:`pyproximal."ProxOperator"`, optional + Proximal operator for the regularization term (if None, no regularization is used) + show : :obj:`bool`, optional + Display setup log + + """ + pass + + @abstractmethod + def solve( + self, + proxf: ProxOperator | list[ProxOperator], + proxg: ProxOperator | None = None, + *args, + show: bool = False, + **kwargs, + ) -> Any: + """Solve + + This method is used to run the entire optimization process. Users can change the + function signature by including any other input parameter required by the solver + + Parameters + ---------- + proxf : :obj:`pyproximal."ProxOperator"` or :obj:`list` + Proximal operator(s) to be used in the optimization + proxg : :obj:`pyproximal."ProxOperator"`, optional + Proximal operator for the regularization term (if None, no regularization is used) + show : :obj:`bool`, optional + Display finalize log + + """ + pass diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py new file mode 100644 index 0000000..be1b678 --- /dev/null +++ b/pyproximal/optimization/cls_primal.py @@ -0,0 +1,821 @@ +__all__ = [ + "ProximalPoint", +] + +import time +from typing import TYPE_CHECKING, Optional + +import numpy as np +from pylops.optimization.callback import _callback_stop +from pylops.utils.backend import get_array_module, to_numpy +from pylops.utils.typing import NDArray, Tmemunit + +from pyproximal.optimization.basesolver import Solver +from pyproximal.ProxOperator import ProxOperator +from pyproximal.utils.bilinear import BilinearOperator + +if TYPE_CHECKING: + from pylops.linearoperator import LinearOperator + + +def _backtracking( + x: NDArray, + tau: float, + proxf: ProxOperator, + proxg: ProxOperator, + epsg: float, + beta: float = 0.5, + niterback: int = 10, +) -> tuple[NDArray, float]: + r"""Backtracking + + Line-search algorithm for finding step sizes in proximal algorithms when + the Lipschitz constant of the operator is unknown (or expensive to + estimate). + + """ + + def ftilde(x: NDArray, y: NDArray, f: ProxOperator, tau: float) -> float: + xy = x - y + return float( + f(y) + np.dot(f.grad(y), xy) + (1.0 / (2.0 * tau)) * np.linalg.norm(xy) ** 2 + ) + + iiterback = 0 + while iiterback < niterback: + z = proxg.prox(x - tau * proxf.grad(x), epsg * tau) + ft = ftilde(z, x, proxf, tau) + if proxf(z) <= ft: + break + tau *= beta + iiterback += 1 + return z, tau + + +def _x0z0_init( + x0: NDArray | None, + z0: NDArray | None, + Op: Optional["LinearOperator"] = None, + z0name: str | None = "z0", + Opname: str | None = "Op", +) -> tuple[NDArray, NDArray]: + r"""Initialize x0 and z0 + + Initialize x0 and z0 using the following convention. + + For ``Op=None``: + - if both are provided, they are simply returned; + - if only one is provided (the other is ``None``), the one provided + is copied to the other one. + + For ``Op!=None``, ``x0`` must be provided, and: + - if both are provided, they are simply returned; + - if ``z0`` is not provided, set to ``Op @ x0``. + + Parameters + ---------- + x0 : :obj:`numpy.ndarray` + Initial vector + z0 : :obj:`numpy.ndarray` + Initial auxiliary vector + Op : :obj:`pylops.LinearOperator`, optional + Linear Operator to apply to ``x0`` + z0name : :obj:`str`, optional + Name to display in error message instead of ``z0`` + Opname : :obj:`str`, optional + Name to display in error message instead of ``Op`` + + """ + if x0 is None and z0 is None: + msg = f"Both x0 or {z0name} are None, provide either of them or both" + raise ValueError(msg) + + if Op is None: + if x0 is None: + x0 = z0.copy() # type: ignore[union-attr] + elif z0 is None: + z0 = x0.copy() + else: + if x0 is None: + msg = f"x0 must be provided when {Opname} is also provided" + raise ValueError(msg) + elif z0 is None: + z0 = Op @ x0 + return x0, z0 + + +class ProximalPoint(Solver): + r"""Proximal point algorithm + + Solves the following minimization problem using Proximal point algorithm: + + .. math:: + + \mathbf{x} = \argmin_\mathbf{x} f(\mathbf{x}) + + where :math:`f(\mathbf{x})` is any convex function that has a known + proximal operator. + + Notes + ----- + The Proximal point algorithm can be expressed by the following recursion: + + .. math:: + + \mathbf{x}^{k+1} = \prox_{\tau f}(\mathbf{x}^k) + + """ + + def _print_setup(self, xcomplex: bool = False) -> None: + self._print_solver(nbar=60) + + strpar = f"Proximal operator: {type(self.prox).__name__}" + if self.niter is not None: + strpar1 = f"tau = {self.tau:6e}\ttol = {self.tol:6e}\tniter = {self.niter}" + else: + strpar1 = f"tau = {self.tau:6e}\ttol = {self.tol:6e}" + print(strpar) + print(strpar1) + print("-" * 60 + "\n") + if not xcomplex: + head1 = " Itn x[0] f" + else: + head1 = " Itn x[0] f" + print(head1) + + def _print_step(self, x: NDArray) -> None: + if self.tol is None: + self.pf = self.prox(x) + strx = f"{x[0]:1.2e} " if np.iscomplexobj(x) else f"{x[0]:11.4e} " + msg = f"{self.iiter:6g} " + strx + f"{self.pf:11.4e}" + print(msg) + + def memory_usage( + self, + show: bool = False, + unit: Tmemunit = "B", + ) -> float: + pass + + def setup( + self, + prox: "ProxOperator", + x0: NDArray, + tau: float, + niter: int | None = None, + tol: float = 1e-4, + show: bool = False, + ) -> NDArray: + r"""Setup solver + + Parameters + ---------- + prox : :obj:`pyproximal."ProxOperator"` + Proximal operator + x0 : :obj:`numpy.ndarray` + Initial guess + tau : :obj:`float` + Positive scalar weight + niter : :obj:`int`, optional + Number of iterations (default to ``None`` in case a user wants to + manually step over the solver) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display setup log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + + """ + self.prox = prox + self.x0 = x0 + self.tau = tau + self.niter = niter + self.tol = tol + + self.ncp = get_array_module(x0) + + # create variables to track the objective function and iterations + self.pf, self.pfold = np.inf, np.inf + self.pfs: list = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(np.iscomplexobj(x0)) + return x0 + + def step(self, x: NDArray, show: bool = False) -> NDArray: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + proximal point algorithm + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + + """ + x = self.prox.prox(x, self.tau) + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfold = self.pf + self.pf = self.prox(x) + if np.abs(1.0 - self.pf / self.pfold) < self.tol: + self.tolbreak = True + + self.iiter += 1 + if show: + self._print_step(x) + if self.tol is not None or show: + self.pfs.append(float(self.pf)) + return x + + def run( + self, + x: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> NDArray: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the proximal point algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x = self.step(x, showstep) + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x + + def finalize(self, show: bool = False) -> None: + r"""Finalize solver + + Parameters + ---------- + show : :obj:`bool`, optional + Display finalize log + + """ + self.tend = time.time() + self.telapsed = self.tend - self.tstart + if show: + self._print_finalize(nbar=60) + + def solve( + self, + prox: "ProxOperator", + x0: NDArray, + tau: float, + niter: int = 10, + tol: float = 1e-4, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + prox : :obj:`pyproximal."ProxOperator"` + Proximal operator + x0 : :obj:`numpy.ndarray`, optional + Initial guess + tau : :obj:`float` + Positive scalar weight + niter : :obj:`int`, optional + Number of iterations + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping + criterion). If ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + iiter : :obj:`int` + Number of executed iterations + pfs : :obj:`numpy.ndarray` + History of the objective function + + """ + x = self.setup(prox=prox, x0=x0, tau=tau, niter=niter, tol=tol, show=show) + x = self.run(x, niter, show=show, itershow=itershow) + self.finalize(show) + return x, self.iiter, self.pfs + + +class ProximalGradient(Solver): + r"""Proximal gradient (optionally accelerated) + + Solves the following minimization problem using (Accelerated) Proximal + gradient algorithm: + + .. math:: + + \mathbf{x} = \argmin_\mathbf{x} f(\mathbf{x}) + \epsilon g(\mathbf{x}) + + where :math:`f(\mathbf{x})` is a smooth convex function with a uniquely + defined gradient and :math:`g(\mathbf{x})` is any convex function that + has a known proximal operator. + + Notes + ----- + The Proximal gradient algorithm can be expressed by the following recursion: + + .. math:: + + \mathbf{x}^{k+1} = \mathbf{y}^k + \eta (\prox_{\tau^k \epsilon g}(\mathbf{y}^k - + \tau^k \nabla f(\mathbf{y}^k)) - \mathbf{y}^k) \\ + \mathbf{y}^{k+1} = \mathbf{x}^k + \omega^k + (\mathbf{x}^k - \mathbf{x}^{k-1}) + + where at each iteration :math:`\tau^k` can be estimated by back-tracking + as follows: + + .. math:: + + \begin{aligned} + &\tau = \tau^{k-1} &\\ + &repeat \; \mathbf{z} = \prox_{\tau \epsilon g}(\mathbf{x}^k - + \tau \nabla f(\mathbf{x}^k)), \tau = \beta \tau \quad if \; + f(\mathbf{z}) \leq \tilde{f}_\tau(\mathbf{z}, \mathbf{x}^k) \\ + &\tau^k = \tau, \quad \mathbf{x}^{k+1} = \mathbf{z} &\\ + \end{aligned} + + where :math:`\tilde{f}_\tau(\mathbf{x}, \mathbf{y}) = f(\mathbf{y}) + + \nabla f(\mathbf{y})^T (\mathbf{x} - \mathbf{y}) + + 1/(2\tau)||\mathbf{x} - \mathbf{y}||_2^2`. + + Different accelerations are provided: + + - ``acceleration=None``: :math:`\omega^k = 0`; + - ``acceleration=vandenberghe`` [1]_: :math:`\omega^k = k / (k + 3)` for ` + - ``acceleration=fista``: :math:`\omega^k = (t_{k-1}-1)/t_k` where + :math:`t_k = (1 + \sqrt{1+4t_{k-1}^{2}}) / 2` [2]_ + + .. [1] Vandenberghe, L., "Fast proximal gradient methods", 2010. + .. [2] Beck, A., and Teboulle, M. "A Fast Iterative Shrinkage-Thresholding + Algorithm for Linear Inverse Problems", SIAM Journal on + Imaging Sciences, vol. 2, pp. 183-202. 2009. + + """ + + def _print_setup(self, epsg_print, xcomplex: bool = False) -> None: + self._print_solver(nbar=81) + + strpar = ( + f"Proximal operator (f): {type(self.proxf).__name__}\n" + f"Proximal operator (g): {type(self.proxg).__name__}\n" + ) + strpar1 = f"tau = {self.tau:4.2e}\t\tbacktrack = {self.backtracking}\tbeta = {self.beta}" + strpar2 = f"epsg = {epsg_print}\t\tniter = {self.niter}\t\ttol = {self.tol}" + strpar3 = f"niterback = {self.niterback}\t\tacceleration = {self.acceleration}" + print(strpar) + print(strpar1) + print(strpar2) + print(strpar3) + print("-" * 81 + "\n") + if not xcomplex: + head1 = " Itn x[0] f g J=f+eps*g tau" + else: + head1 = " Itn x[0] f g J=f+eps*g tau" + print(head1) + + def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: + if self.tol is None: + pf, pg = self.proxf(x), self.proxg(x) + self.pfg = pf + np.sum(self.epsg[self.iiter - 1] * pg) + x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = ( + f"{self.iiter:6g} " + + strx + + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e} {self.tau:11.2e}" + ) + print(msg) + + def memory_usage( + self, + show: bool = False, + unit: Tmemunit = "B", + ) -> float: + pass + + def setup( + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + epsg: float | NDArray = 1.0, + tau: float | None = None, + backtracking: bool = False, + beta: float = 0.5, + eta: float = 1.0, + niter: int = 10, + niterback: int = 100, + acceleration: str | None = None, + tol: float | None = None, + show: bool = False, + ) -> NDArray: + r"""Setup solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function (must have ``grad`` implemented) + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector + epsg : :obj:`float` or :obj:`numpy.ndarray`, optional + Scaling factor of g function + tau : :obj:`float` or :obj:`numpy.ndarray`, optional + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. When ``tau=None``, + backtracking is used to adaptively estimate the best tau at each + iteration. Finally, note that :math:`\tau` can be chosen to be a vector + when dealing with problems with multiple right-hand-sides + backtracking : :obj:`bool`, optional + Force backtracking, even if ``tau`` is not equal to ``None``. In this case + the chosen ``tau`` will be used as the initial guess in the first + step of backtracking + beta : :obj:`float`, optional + Backtracking parameter (must be between 0 and 1) + eta : :obj:`float`, optional + Relaxation parameter (must be between 0 and 1, 0 excluded). + niter : :obj:`int`, optional + Number of iterations of iterative scheme + niterback : :obj:`int`, optional + Max number of iterations of backtracking + acceleration : :obj:`str`, optional + Acceleration (``None``, ``vandenberghe`` or ``fista``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + + """ + self.proxf = proxf + self.proxg = proxg + self.x0 = x0 + self.backtracking = backtracking + self.beta = beta + self.eta = eta + self.niter = niter + self.niterback = niterback + self.tol = tol + + self.ncp = get_array_module(x0) + + # check if epgs is a vector + self.epsg = np.asarray(epsg, dtype=float) + if self.epsg.size == 1: + self.epsg = self.epsg * np.ones(niter) + epsg_print = str(self.epsg[0]) + else: + epsg_print = "Multi" + + # set tau + self.tau = tau + if tau is None: + self.backtracking = True + self.tau = 1.0 + + # check acceleration + if acceleration in [None, "None", "vandenberghe", "fista"]: + self.acceleration = acceleration + else: + msg = "Acceleration should be None, vandenberghe or fista" + raise NotImplementedError(msg) + + # set initial vectors + x = x0.copy() + y = x.copy() + + # for accelaration + self.t = 1.0 + + # create variables to track the objective function and iterations + self.pfg, self.pfgold = np.inf, np.inf + self.pfs: list = [] + self.pfgs: list = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(epsg_print, np.iscomplexobj(x0)) + return x, y + + def step(self, x: NDArray, y: NDArray, show: bool = False) -> NDArray: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + proximal gradient algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by a step of the + proximal gradient algorithm + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + + """ + xold = x.copy() + + # proximal step + if not self.backtracking: + if self.eta == 1.0: + x = self.proxg.prox( + y - self.tau * self.proxf.grad(y), self.epsg[self.iiter] * self.tau + ) + else: + x = x + self.eta * ( + self.proxg.prox( + x - self.tau * self.proxf.grad(x), + self.epsg[self.iiter] * self.tau, + ) + - x + ) + else: + x, self.tau = _backtracking( + y, + self.tau, + self.proxf, + self.proxg, + self.epsg[self.iiter], + beta=self.beta, + niterback=self.niterback, + ) + if self.eta != 1.0: + x = x + self.eta * ( + self.proxg.prox( + x - self.tau * self.proxf.grad(x), + self.epsg[self.iiter] * self.tau, + ) + - x + ) + + # update internal parameters for bilinear operator + if isinstance(self.proxf, BilinearOperator): + self.proxf.updatexy(x) + + # update y + if self.acceleration == "vandenberghe": + omega = self.iiter / (self.iiter + 3) + elif self.acceleration == "fista": + told = self.t + self.t = (1.0 + np.sqrt(1.0 + 4.0 * self.t**2)) / 2.0 + omega = (told - 1.0) / self.t + else: + omega = 0 + y = x + omega * (x - xold) + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfgold = self.pfg + pf, pg = self.proxf(x), self.proxg(x) + self.pfg = pf + np.sum(self.epsg[self.iiter] * pg) + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pf, pg = None, None + + self.iiter += 1 + if show: + self._print_step(x, pf, pg) + if self.tol is not None or show: + self.pfgs.append(float(self.pfg)) + return x, y + + def run( + self, + x: NDArray, + y: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> NDArray: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the proximal gradient algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the proximal gradient algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, y = self.step(x, y, showstep) + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, y + + def finalize(self, show: bool = False) -> None: + r"""Finalize solver + + Parameters + ---------- + show : :obj:`bool`, optional + Display finalize log + + """ + self.tend = time.time() + self.telapsed = self.tend - self.tstart + if show: + self._print_finalize(nbar=81) + + def solve( + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + epsg: float | NDArray = 1.0, + tau: float | None = None, + backtracking: bool = False, + beta: float = 0.5, + eta: float = 1.0, + niter: int = 10, + niterback: int = 100, + acceleration: str | None = None, + tol: float | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function (must have ``grad`` implemented) + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector + epsg : :obj:`float` or :obj:`numpy.ndarray`, optional + Scaling factor of g function + tau : :obj:`float` or :obj:`numpy.ndarray`, optional + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. When ``tau=None``, + backtracking is used to adaptively estimate the best tau at each + iteration. Finally, note that :math:`\tau` can be chosen to be a vector + when dealing with problems with multiple right-hand-sides + backtracking : :obj:`bool`, optional + Force backtracking, even if ``tau`` is not equal to ``None``. In this case + the chosen ``tau`` will be used as the initial guess in the first + step of backtracking + beta : :obj:`float`, optional + Backtracking parameter (must be between 0 and 1) + eta : :obj:`float`, optional + Relaxation parameter (must be between 0 and 1, 0 excluded). + niter : :obj:`int`, optional + Number of iterations of iterative scheme + niterback : :obj:`int`, optional + Max number of iterations of backtracking + acceleration : :obj:`str`, optional + Acceleration (``None``, ``vandenberghe`` or ``fista``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping + criterion). If ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + iiter : :obj:`int` + Number of executed iterations + pfgs : :obj:`numpy.ndarray` + History of the objective function + + """ + x, y = self.setup( + proxf=proxf, + proxg=proxg, + x0=x0, + epsg=epsg, + tau=tau, + backtracking=backtracking, + beta=beta, + eta=eta, + niter=niter, + niterback=niterback, + acceleration=acceleration, + tol=tol, + show=show, + ) + + x, y = self.run(x, y, niter, show=show, itershow=itershow) + self.finalize(show) + return x, y, self.iiter, self.pfgs From 3605b1d74c9785cdb063fb39781fd8d7a9cf068c Mon Sep 17 00:00:00 2001 From: mrava87 Date: Tue, 14 Apr 2026 20:58:52 +0100 Subject: [PATCH 02/24] test: added tests for ProximalPoint --- pytests/test_solver.py | 42 ++++++++++++++++++++++++++++-------------- 1 file changed, 28 insertions(+), 14 deletions(-) diff --git a/pytests/test_solver.py b/pytests/test_solver.py index 8480e76..2c5b5dc 100644 --- a/pytests/test_solver.py +++ b/pytests/test_solver.py @@ -15,8 +15,9 @@ GeneralizedProximalGradient, LinearizedADMM, ProximalGradient, + ProximalPoint, ) -from pyproximal.proximal import L1, L2 +from pyproximal.proximal import L1, L2, Quadratic par1 = {"n": 8, "m": 10, "dtype": "float32"} # float64 par2 = {"n": 8, "m": 10, "dtype": "float64"} # float32 @@ -137,6 +138,29 @@ def test_GPG_weights(par): ) +@pytest.mark.parametrize("par", [(par1), (par2)]) +def test_ProximalPoint(par): + """Check solution of ProximalPoint for quadratic function equals the solution of the + associated system of linear equations + """ + np.random.seed(10) + m = par["m"] + + # Random mixing matrix + A = np.random.normal(0.0, 1.0, (m, m)) + A = A.T @ A + + # Model and data + x = np.linspace(-5.0, 5.0, par["m"]) + y = A @ x + + # Proximal point algorithm with quadatic function + quad = Quadratic(Op=MatrixMult(A), b=-y, niter=2) + xpp = ProximalPoint(quad, x0=np.zeros_like(x), tau=0.1, niter=1000, tol=0) + + assert_array_almost_equal(xpp, x, decimal=2) + + @pytest.mark.parametrize("par", [(par1), (par2)]) def test_PG_GPG(par): """Check equivalency of ProximalGradient and GeneralizedProximalGradient when using @@ -210,20 +234,20 @@ def test_ADMM_DRS(par): # ADMM l2 = L2(Op=Rop, b=y, niter=10, warm=True) l1 = L1(sigma=5e-1) - xadmm, zadmm = ADMM(l2, l1, x0=np.zeros(m), tau=tau, niter=100, show=True) + xadmm, zadmm = ADMM(l2, l1, x0=np.zeros(m), tau=tau, niter=100) # DRS with g first l2 = L2(Op=Rop, b=y, niter=10, warm=True) l1 = L1(sigma=5e-1) xdrs_g, ydrs_g = DouglasRachfordSplitting( - l2, l1, x0=np.zeros(m), tau=tau, niter=100, show=True, gfirst=True + l2, l1, x0=np.zeros(m), tau=tau, niter=100, gfirst=True ) # DRS with f first l2 = L2(Op=Rop, b=y, niter=10, warm=True) l1 = L1(sigma=5e-1) xdrs_f, ydrs_f = DouglasRachfordSplitting( - l2, l1, x0=np.zeros(m), tau=tau, niter=100, show=True, gfirst=False + l2, l1, x0=np.zeros(m), tau=tau, niter=100, gfirst=False ) assert_array_almost_equal(xadmm, xdrs_g, decimal=2) @@ -262,13 +286,11 @@ def test_PPXA_with_ADMM(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=2000, # niter=1500 makes this test fail for seeds 0 to 499 - show=True, ) xppxa = PPXA( [l2, l1], x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), - show=True, ) assert_array_almost_equal(xppxa, xadmm, decimal=2) @@ -312,13 +334,11 @@ def test_PPXA_with_GPG(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=200, # niter=150 makes this test fail for seeds 0 to 499 - show=True, ) xppxa = PPXA( [l2_1, l2_2, l1_1, l1_2], x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), - show=True, ) assert_array_almost_equal(xppxa, xgpg, decimal=2) @@ -356,13 +376,11 @@ def test_ConsensusADMM_with_ADMM(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=2000, # niter=1500 makes this test fail for seeds 0 to 499 - show=True, ) xcadmm = ConsensusADMM( [l2, l1], x0=np.random.normal(0.0, 1.0, m), # x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), - show=True, ) assert_array_almost_equal(xcadmm, xadmm, decimal=2) @@ -416,7 +434,6 @@ def test_ConsensusADMM_with_ADMM_for_Lasso(par: dict[str, Any]) -> None: x0=np.random.normal(0.0, 1.0, m), # x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), niter=20000, # niter=15000 makes this test fail for seeds 0 to 499 - show=True, ) # 1/2 || [R1; R2; R3] ||_2^2 + ||x||_1 @@ -426,7 +443,6 @@ def test_ConsensusADMM_with_ADMM_for_Lasso(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=15000, # niter=10000 makes this test fail for seeds 0 to 499 - show=True, ) assert_array_almost_equal(xcadmm, xadmm, decimal=2) @@ -471,13 +487,11 @@ def test_ConsensusADMM_with_GPG(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=200, # niter=150 makes this test fail for seeds 0 to 499 - show=True, ) xppxa = ConsensusADMM( [l2_1, l2_2, l1_1, l1_2], x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), - show=True, ) assert_array_almost_equal(xppxa, xgpg, decimal=2) From 057265c4228b130ce925930e0f7e460d89859a04 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Wed, 15 Apr 2026 21:34:12 +0100 Subject: [PATCH 03/24] feat: added class based AndersonProximalGradient (and fix mypy errors) --- pyproximal/optimization/basesolver.py | 6 +- pyproximal/optimization/cls_primal.py | 484 ++++++++++++++++++++++++-- 2 files changed, 467 insertions(+), 23 deletions(-) diff --git a/pyproximal/optimization/basesolver.py b/pyproximal/optimization/basesolver.py index d64d839..222979a 100644 --- a/pyproximal/optimization/basesolver.py +++ b/pyproximal/optimization/basesolver.py @@ -15,7 +15,7 @@ _units = {"B": 1, "KB": 1024, "MB": 1024**2, "GB": 1024**3} -class Solver(pSolver, metaclass=ABCMeta): +class Solver(pSolver, metaclass=ABCMeta): # type: ignore[misc] r"""Solver This is a template class which a user must subclass when implementing a new solver. @@ -114,9 +114,9 @@ def solve( self, proxf: ProxOperator | list[ProxOperator], proxg: ProxOperator | None = None, - *args, + *args: Any, show: bool = False, - **kwargs, + **kwargs: Any, ) -> Any: """Solve diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index be1b678..5a125e6 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -3,7 +3,7 @@ ] import time -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, cast import numpy as np from pylops.optimization.callback import _callback_stop @@ -126,14 +126,18 @@ class ProximalPoint(Solver): """ + pf: float + def _print_setup(self, xcomplex: bool = False) -> None: self._print_solver(nbar=60) strpar = f"Proximal operator: {type(self.prox).__name__}" if self.niter is not None: - strpar1 = f"tau = {self.tau:6e}\ttol = {self.tol:6e}\tniter = {self.niter}" + strpar1 = ( + f"tau = {self.tau:6e}\ttol = {str(self.tol)}\tniter = {self.niter}" + ) else: - strpar1 = f"tau = {self.tau:6e}\ttol = {self.tol:6e}" + strpar1 = f"tau = {self.tau:6e}\ttol = {str(self.tol)}" print(strpar) print(strpar1) print("-" * 60 + "\n") @@ -154,10 +158,10 @@ def memory_usage( self, show: bool = False, unit: Tmemunit = "B", - ) -> float: + ) -> None: pass - def setup( + def setup( # type: ignore[override] self, prox: "ProxOperator", x0: NDArray, @@ -201,7 +205,7 @@ def setup( # create variables to track the objective function and iterations self.pf, self.pfold = np.inf, np.inf - self.pfs: list = [] + self.pfs: list[float] = [] self.tolbreak = False self.iiter = 0 @@ -311,7 +315,7 @@ def finalize(self, show: bool = False) -> None: if show: self._print_finalize(nbar=60) - def solve( + def solve( # type: ignore[override] self, prox: "ProxOperator", x0: NDArray, @@ -415,7 +419,7 @@ class ProximalGradient(Solver): """ - def _print_setup(self, epsg_print, xcomplex: bool = False) -> None: + def _print_setup(self, epsg_print: str, xcomplex: bool = False) -> None: self._print_solver(nbar=81) strpar = ( @@ -423,7 +427,9 @@ def _print_setup(self, epsg_print, xcomplex: bool = False) -> None: f"Proximal operator (g): {type(self.proxg).__name__}\n" ) strpar1 = f"tau = {self.tau:4.2e}\t\tbacktrack = {self.backtracking}\tbeta = {self.beta}" - strpar2 = f"epsg = {epsg_print}\t\tniter = {self.niter}\t\ttol = {self.tol}" + strpar2 = ( + f"epsg = {epsg_print}\t\tniter = {self.niter}\t\ttol = {str(self.tol)}" + ) strpar3 = f"niterback = {self.niterback}\t\tacceleration = {self.acceleration}" print(strpar) print(strpar1) @@ -453,10 +459,10 @@ def memory_usage( self, show: bool = False, unit: Tmemunit = "B", - ) -> float: + ) -> None: pass - def setup( + def setup( # type: ignore[override] self, proxf: ProxOperator, proxg: ProxOperator, @@ -471,7 +477,7 @@ def setup( acceleration: str | None = None, tol: float | None = None, show: bool = False, - ) -> NDArray: + ) -> tuple[NDArray, NDArray]: r"""Setup solver Parameters @@ -515,6 +521,8 @@ def setup( ------- x : :obj:`numpy.ndarray` Initial guess + y : :obj:`numpy.ndarray` + Initial guess for the auxiliary variable """ self.proxf = proxf @@ -559,8 +567,8 @@ def setup( # create variables to track the objective function and iterations self.pfg, self.pfgold = np.inf, np.inf - self.pfs: list = [] - self.pfgs: list = [] + self.pfs: list[float] = [] + self.pfgs: list[float] = [] self.tolbreak = False self.iiter = 0 @@ -569,7 +577,9 @@ def setup( self._print_setup(epsg_print, np.iscomplexobj(x0)) return x, y - def step(self, x: NDArray, y: NDArray, show: bool = False) -> NDArray: + def step( + self, x: NDArray, y: NDArray, show: bool = False + ) -> tuple[NDArray, NDArray]: r"""Run one step of solver Parameters @@ -587,6 +597,8 @@ def step(self, x: NDArray, y: NDArray, show: bool = False) -> NDArray: ------- x : :obj:`numpy.ndarray` Updated model vector + y : :obj:`numpy.ndarray` + Updated additional model vector """ xold = x.copy() @@ -608,7 +620,7 @@ def step(self, x: NDArray, y: NDArray, show: bool = False) -> NDArray: else: x, self.tau = _backtracking( y, - self.tau, + cast(float, self.tau), self.proxf, self.proxg, self.epsg[self.iiter], @@ -648,7 +660,7 @@ def step(self, x: NDArray, y: NDArray, show: bool = False) -> NDArray: if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: self.tolbreak = True else: - pf, pg = None, None + pf, pg = 0.0, 0.0 self.iiter += 1 if show: @@ -664,7 +676,7 @@ def run( niter: int | None = None, show: bool = False, itershow: tuple[int, int, int] = (10, 10, 10), - ) -> NDArray: + ) -> tuple[NDArray, NDArray]: r"""Run solver Parameters @@ -730,7 +742,7 @@ def finalize(self, show: bool = False) -> None: if show: self._print_finalize(nbar=81) - def solve( + def solve( # type: ignore[override] self, proxf: ProxOperator, proxg: ProxOperator, @@ -746,7 +758,7 @@ def solve( tol: float | None = None, show: bool = False, itershow: tuple[int, int, int] = (10, 10, 10), - ) -> tuple[NDArray, int, NDArray]: + ) -> tuple[NDArray, NDArray, int, NDArray]: r"""Run entire solver Parameters @@ -794,6 +806,8 @@ def solve( ------- x : :obj:`numpy.ndarray` Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model iiter : :obj:`int` Number of executed iterations pfgs : :obj:`numpy.ndarray` @@ -819,3 +833,433 @@ def solve( x, y = self.run(x, y, niter, show=show, itershow=itershow) self.finalize(show) return x, y, self.iiter, self.pfgs + + +class AndersonProximalGradient(Solver): + r"""Proximal gradient with Anderson acceleration + + Solves the following minimization problem using the Proximal + gradient algorithm with Anderson acceleration: + + .. math:: + + \mathbf{x} = \argmin_\mathbf{x} f(\mathbf{x}) + \epsilon g(\mathbf{x}) + + where :math:`f(\mathbf{x})` is a smooth convex function with a uniquely + defined gradient and :math:`g(\mathbf{x})` is any convex function that + has a known proximal operator. + + Notes + ----- + The Proximal gradient algorithm with Anderson acceleration can be expressed by the + following recursion [1]_: + + .. math:: + m_k = min(m, k)\\ + \mathbf{g}^{k} = \mathbf{x}^{k} - \tau^k \nabla f(\mathbf{x}^k)\\ + \mathbf{r}^{k} = \mathbf{g}^{k} - \mathbf{g}^{k}\\ + \mathbf{G}^{k} = [\mathbf{g}^{k},..., \mathbf{g}^{k-m_k}]\\ + \mathbf{R}^{k} = [\mathbf{r}^{k},..., \mathbf{r}^{k-m_k}]\\ + \alpha_k = (\mathbf{R}^{kT} \mathbf{R}^{k})^{-1} \mathbf{1} / \mathbf{1}^T + (\mathbf{R}^{kT} \mathbf{R}^{k})^{-1} \mathbf{1}\\ + \mathbf{y}^{k+1} = \mathbf{G}^{k} \alpha_k\\ + \mathbf{x}^{k+1} = \prox_{\tau^{k+1} g}(\mathbf{y}^{k+1}) + + where :math:`m` equals ``nhistory``, :math:`k=1,2,...,n_{iter}`, :math:`\mathbf{y}^{0}=\mathbf{x}^{0}`, + :math:`\mathbf{y}^{1}=\mathbf{x}^{0} - \tau^0 \nabla f(\mathbf{x}^0)`, + :math:`\mathbf{x}^{1}=\prox_{\tau^k g}(\mathbf{y}^{1})`, and + :math:`\mathbf{g}^{0}=\mathbf{y}^{1}`. + + Refer to [1]_ for the guarded version of the algorithm (when ``safeguard=True``). + + .. [1] Mai, V., and Johansson, M. "Anderson Acceleration of Proximal Gradient + Methods", 2020. + + """ + + def _print_setup(self, epsg_print: str, xcomplex: bool = False) -> None: + self._print_solver(nbar=81) + + strpar = ( + f"Proximal operator (f): {type(self.proxf).__name__}\n" + f"Proximal operator (g): {type(self.proxg).__name__}\n" + ) + strpar1 = ( + f"tau = {self.tau:4.2e}\t\tepsg = {epsg_print}\t\tniter = {self.niter}" + ) + strpar2 = f"nhist = {self.nhistory}\t\tepsr = {self.epsr:4.2e}" + strpar3 = f"guard = {str(self.safeguard)}\t\ttol = {str(self.tol)}" + print(strpar) + print(strpar1) + print(strpar2) + print(strpar3) + print("-" * 81 + "\n") + if not xcomplex: + head1 = " Itn x[0] f g J=f+eps*g tau" + else: + head1 = " Itn x[0] f g J=f+eps*g tau" + print(head1) + + def _print_step(self, x: NDArray, pg: float | None) -> None: + if self.tol is None: + self.pf, pg = self.proxf(x), self.proxg(x) + self.pfg = self.pf + np.sum(self.epsg[self.iiter - 1] * pg) + x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = ( + f"{self.iiter:6g} " + + strx + + f"{self.pf:10.3e} {pg:10.3e} {self.pfg:10.3e} {self.tau:11.2e}" + ) + print(msg) + + def memory_usage( + self, + show: bool = False, + unit: Tmemunit = "B", + ) -> None: + pass + + def setup( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + epsg: float | NDArray = 1.0, + tau: float | NDArray = 1.0, + niter: int = 10, + nhistory: int = 10, + epsr: float = 1e-10, + safeguard: bool = False, + tol: float | None = None, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Setup solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function (must have ``grad`` implemented) + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector + epsg : :obj:`float` or :obj:`numpy.ndarray`, optional + Scaling factor of g function + tau : :obj:`float` or :obj:`numpy.ndarray`, optional + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. N ote that :math:`\tau` + can be chosen to be a vector when dealing with problems with + multiple right-hand-sides + niter : :obj:`int`, optional + Number of iterations of iterative scheme + nhistory : :obj:`int`, optional + Number of previous iterates to be kept in memory (to compute the scaling factors) + epsr : :obj:`float`, optional + Scaling factor for regularization added to the inverse of :math:\mathbf{R}^T \mathbf{R}` + safeguard : :obj:`bool`, optional + Apply safeguarding strategy to the update (``True``) or not (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + + """ + self.proxf = proxf + self.proxg = proxg + self.x0 = x0 + self.tau = tau + self.niter = niter + self.nhistory = nhistory + self.epsr = epsr + self.safeguard = safeguard + self.tol = tol + + self.ncp = get_array_module(x0) + + # check if epgs is a vector + self.epsg = np.asarray(epsg, dtype=float) + if self.epsg.size == 1: + self.epsg = self.epsg * np.ones(niter) + epsg_print = str(self.epsg[0]) + else: + epsg_print = "Multi" + + # set initial vectors + y = x0 - self.tau * proxf.grad(x0) + x = self.proxg.prox(y, self.epsg[0] * self.tau) + + # set history of iterates for Anderson acceleration + g = y.copy() + r = g - x0 + self.R, self.G = ( + [ + g, + ], + [ + r, + ], + ) + + # create variables to track the objective function and iterations + self.pf = self.proxf(x) + print("Initial pf: %10.3e" % self.pf) + self.pfg, self.pfgold = np.inf, np.inf + self.pfs: list[float] = [] + self.pfgs: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(epsg_print, np.iscomplexobj(x0)) + return x, y + + def step( + self, x: NDArray, y: NDArray, show: bool = False + ) -> tuple[NDArray, NDArray]: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + proximal gradient algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by a step of the + proximal gradient algorithm + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + y : :obj:`numpy.ndarray` + Updated additional model vector + + """ + + # update fix point + g = x - self.tau * self.proxf.grad(x) + r = g - y + + # update history vectors + self.R.insert(0, r) + self.G.insert(0, g) + if self.iiter >= self.nhistory - 1: + self.R.pop(-1) + self.G.pop(-1) + + # solve for alpha coefficients + Rstack = np.vstack(self.R) + Rinv = np.linalg.pinv( + Rstack @ Rstack.T + self.epsr * np.linalg.norm(Rstack) ** 2 + ) + ones = np.ones(min(self.nhistory, self.iiter + 2)) + Rinvones = Rinv @ ones + alpha = Rinvones / (ones[None] @ Rinvones) + + if not self.safeguard: + # update auxiliary variable + y = np.vstack(self.G).T @ alpha + + # update main variable + x = self.proxg.prox(y, self.epsg[self.iiter] * self.tau) + else: + # update auxiliary variable + ytest = np.vstack(self.G).T @ alpha + + # update main variable + xtest = self.proxg.prox(ytest, self.epsg[self.iiter] * self.tau) + + # check if function is decreased, otherwise do basic PG step + pfold, self.pf = self.pf, self.proxf(xtest) + if ( + self.pf + <= pfold - self.tau * np.linalg.norm(self.proxf.grad(x)) ** 2 / 2 + ): + y = ytest + x = xtest + else: + x = self.proxg.prox(g, self.epsg[self.iiter] * self.tau) + y = g + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfgold = self.pfg + self.pf, pg = self.proxf(x), self.proxg(x) + self.pfg = self.pf + np.sum(self.epsg[self.iiter] * pg) + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pg = 0.0 + + self.iiter += 1 + if show: + self._print_step(x, pg) + if self.tol is not None or show: + self.pfgs.append(float(self.pfg)) + return x, y + + def run( + self, + x: NDArray, + y: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> NDArray: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the Anderson proximal gradient algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the proximal gradient algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, y = self.step(x, y, showstep) + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, y + + def finalize(self, show: bool = False) -> None: + r"""Finalize solver + + Parameters + ---------- + show : :obj:`bool`, optional + Display finalize log + + """ + self.tend = time.time() + self.telapsed = self.tend - self.tstart + if show: + self._print_finalize(nbar=81) + + def solve( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + epsg: float | NDArray = 1.0, + tau: float | NDArray = 1.0, + niter: int = 10, + nhistory: int = 10, + epsr: float = 1e-10, + safeguard: bool = False, + tol: float | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function (must have ``grad`` implemented) + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector + epsg : :obj:`float` or :obj:`numpy.ndarray`, optional + Scaling factor of g function + tau : :obj:`float` or :obj:`numpy.ndarray`, optional + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. N ote that :math:`\tau` + can be chosen to be a vector when dealing with problems with + multiple right-hand-sides + niter : :obj:`int`, optional + Number of iterations of iterative scheme + nhistory : :obj:`int`, optional + Number of previous iterates to be kept in memory (to compute the scaling factors) + epsr : :obj:`float`, optional + Scaling factor for regularization added to the inverse of :math:\mathbf{R}^T \mathbf{R}` + safeguard : :obj:`bool`, optional + Apply safeguarding strategy to the update (``True``) or not (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping + criterion). If ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model + iiter : :obj:`int` + Number of executed iterations + pfgs : :obj:`numpy.ndarray` + History of the objective function + + """ + x, y = self.setup( + proxf=proxf, + proxg=proxg, + x0=x0, + epsg=epsg, + tau=tau, + niter=niter, + nhistory=nhistory, + epsr=epsr, + safeguard=safeguard, + tol=tol, + show=show, + ) + + x, y = self.run(x, y, niter, show=show, itershow=itershow) + self.finalize(show) + return x, y, self.iiter, self.pfgs From 8814e307a403a1bf3fbdfa16d1e0eecf53612ce0 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Wed, 15 Apr 2026 22:41:54 +0100 Subject: [PATCH 04/24] feat: added class based GeneralizedProximalGradient --- pyproximal/optimization/cls_primal.py | 428 +++++++++++++++++++++++++- pytests/test_solver.py | 65 ++++ 2 files changed, 492 insertions(+), 1 deletion(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index 5a125e6..d0f8418 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -3,10 +3,12 @@ ] import time +from collections.abc import Sequence from typing import TYPE_CHECKING, Optional, cast import numpy as np -from pylops.optimization.callback import _callback_stop +import pylops +from pylops.optimization.callback import Callbacks from pylops.utils.backend import get_array_module, to_numpy from pylops.utils.typing import NDArray, Tmemunit @@ -18,6 +20,17 @@ from pylops.linearoperator import LinearOperator +# need to check pylops version since _callback_stop +# is only available in pylops>=2.6.0 +sp_version = pylops.__version__.split(".") +if int(sp_version[0]) < 2 or (int(sp_version[0]) == 2 and int(sp_version[1]) < 6): + + def _callback_stop(callbacks: Sequence[Callbacks]) -> bool: + return False +else: + from pylops.optimization.callback import _callback_stop # type: ignore[no-redef] + + def _backtracking( x: NDArray, tau: float, @@ -1263,3 +1276,416 @@ def solve( # type: ignore[override] x, y = self.run(x, y, niter, show=show, itershow=itershow) self.finalize(show) return x, y, self.iiter, self.pfgs + + +class GeneralizedProximalGradient(Solver): + r"""Generalized Proximal gradient + + Solves the following minimization problem using Generalized Proximal + gradient algorithm: + + .. math:: + + \mathbf{x} = \argmin_\mathbf{x} \sum_{i=1}^n f_i(\mathbf{x}) + + \sum_{j=1}^m \epsilon_j g_j(\mathbf{x}),~~n,m \in \mathbb{N}^+ + + where the :math:`f_i(\mathbf{x})` are smooth convex functions with a uniquely + defined gradient and the :math:`g_j(\mathbf{x})` are any convex function that + have a known proximal operator. + + Notes + ----- + The Generalized Proximal gradient algorithm can be expressed by the + following recursion [1]_: + + .. math:: + \text{for } j=1,\cdots,n, \\ + ~~~~\mathbf z_j^{k+1} = \mathbf z_j^{k} + \eta + \left[prox_{\frac{\tau^k \epsilon_j}{w_j} g_j}\left(2 \mathbf{x}^{k} - \mathbf{z}_j^{k} + - \tau^k \sum_{i=1}^n \nabla f_i(\mathbf{x}^{k})\right) - \mathbf{x}^{k} \right] \\ + \mathbf{x}^{k+1} = \sum_{j=1}^n w_j \mathbf z_j^{k+1} \\ + + where :math:`\sum_{j=1}^n w_j=1`. In the current implementation, :math:`w_j=1/n` when + not provided. + + .. [1] Raguet, H., Fadili, J. and Peyré, G. "Generalized Forward-Backward Splitting", + arXiv, 2012. + + """ + + def _print_setup(self, epsg_print: str, xcomplex: bool = False) -> None: + self._print_solver(nbar=65) + + strpar = ( + f"Proximal operators (f): {[type(proxf).__name__ for proxf in self.proxfs]}\n" + f"Proximal operators (g): {[type(proxg).__name__ for proxg in self.proxgs]}\n" + ) + strpar1 = f"tau = {self.tau:4.2e}\tepsg = {epsg_print}\tniter = {self.niter}" + print(strpar) + print(strpar1) + print("-" * 65 + "\n") + if not xcomplex: + head1 = " Itn x[0] f g J=f+g" + else: + head1 = ( + " Itn x[0] f g J=f+g" + ) + print(head1) + + def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: + if self.tol is None: + pf = np.sum([proxf(x) for proxf in self.proxfs]) + pg = np.sum( + [ + eg * proxg(x) + for proxg, eg in zip(self.proxgs, self.epsg, strict=True) + ] + ) + self.pfg = pf + pg + x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = ( + f"{self.iiter:6g} " + + strx + + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e}" + ) + print(msg) + + def memory_usage( + self, + show: bool = False, + unit: Tmemunit = "B", + ) -> None: + pass + + def setup( # type: ignore[override] + self, + proxfs: list[ProxOperator], + proxgs: list[ProxOperator], + x0: NDArray, + tau: float, + epsg: float | NDArray = 1.0, + weights: NDArray | None = None, + eta: float = 1.0, + niter: int = 10, + acceleration: str | None = None, + tol: float | None = None, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Setup solver + + Parameters + ---------- + proxfs : :obj:`list` + Proximal operators of the :math:`f_i` functions (must have ``grad`` implemented) + proxgs : :obj:`list` + Proximal operators of the :math:`g_j` functions + x0 : :obj:`numpy.ndarray` + Initial vector + tau : :obj:`float` + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\sum_{i=1}^n \nabla f_i`. + epsg : :obj:`float` or :obj:`numpy.ndarray`, optional + Scaling factor(s) of ``g`` function(s) + weights : :obj:`float`, optional + Weighting factors of ``g`` functions. Must sum to 1. + eta : :obj:`float`, optional + Relaxation parameter (must be between 0 and 1, 0 excluded). Note that + this will be only used when ``acceleration=None``. + niter : :obj:`int`, optional + Number of iterations of iterative scheme + acceleration: :obj:`str`, optional + Acceleration (``None``, ``vandenberghe`` or ``fista``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + y : :obj:`numpy.ndarray` + Initial guess for the auxiliary variable + + """ + self.proxfs = proxfs + self.proxgs = proxgs + self.x0 = x0 + self.tau = tau + self.eta = eta + self.niter = niter + self.tol = tol + + self.ncp = get_array_module(x0) + + # check if weights sum to 1 + self.weights = ( + np.ones(len(proxgs)) / len(proxgs) if weights is None else weights + ) + if len(self.weights) != len(self.proxgs) or np.sum(self.weights) != 1.0: + msg = f"weights={self.weights} must be an array of size {len(self.proxgs)} summing to 1" + raise ValueError(msg) + + # check if epgs is a vector + self.epsg = np.asarray(epsg, dtype=float) + if self.epsg.size == 1: + self.epsg = epsg * np.ones(len(proxgs)) + epsg_print = str(self.epsg[0]) + else: + epsg_print = "Multi" + + # check acceleration + if acceleration in [None, "None", "vandenberghe", "fista"]: + self.acceleration = acceleration + else: + msg = "Acceleration should be None, vandenberghe or fista" + raise NotImplementedError(msg) + + # set initial vectors + x = x0.copy() + y = x.copy() + self.zs = [x.copy() for _ in range(len(proxgs))] + + # for accelaration + self.t = 1.0 + + # create variables to track the objective function and iterations + self.pfg, self.pfgold = np.inf, np.inf + self.pfgs: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(epsg_print, np.iscomplexobj(x0)) + return x, y + + def step( + self, x: NDArray, y: NDArray, show: bool = False + ) -> tuple[NDArray, NDArray]: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + proximal gradient algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by a step of the + proximal gradient algorithm + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + y : :obj:`numpy.ndarray` + Updated additional model vector + + """ + xold = x.copy() + + # gradient + grad = np.zeros_like(x) + for _, proxf in enumerate(self.proxfs): + grad += proxf.grad(x) + + # proximal step + x = np.zeros_like(x) + for i, proxg in enumerate(self.proxgs): + ztmp = 2 * y - self.zs[i] - self.tau * grad + ztmp = proxg.prox(ztmp, self.tau * self.epsg[i] / self.weights[i]) + self.zs[i] += self.eta * (ztmp - y) + x += self.weights[i] * self.zs[i] + + # update y + if self.acceleration == "vandenberghe": + omega = self.iiter / (self.iiter + 3) + elif self.acceleration == "fista": + told = self.t + self.t = (1.0 + np.sqrt(1.0 + 4.0 * self.t**2)) / 2.0 + omega = (told - 1.0) / self.t + else: + omega = 0 + y = x + omega * (x - xold) + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfgold = self.pfg + pf = np.sum([proxf(x) for proxf in self.proxfs]) + pg = np.sum( + [ + eg * proxg(x) + for proxg, eg in zip(self.proxgs, self.epsg, strict=True) + ] + ) + self.pfg = pf + pg + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pf, pg = 0.0, 0.0 + + self.iiter += 1 + if show: + self._print_step(x, pf, pg) + if self.tol is not None or show: + self.pfgs.append(float(self.pfg)) + return x, y + + def run( + self, + x: NDArray, + y: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray]: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the proximal gradient algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the proximal gradient algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, y = self.step(x, y, showstep) + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, y + + def finalize(self, show: bool = False) -> None: + r"""Finalize solver + + Parameters + ---------- + show : :obj:`bool`, optional + Display finalize log + + """ + self.tend = time.time() + self.telapsed = self.tend - self.tstart + if show: + self._print_finalize(nbar=65) + + def solve( # type: ignore[override] + self, + proxfs: list[ProxOperator], + proxgs: list[ProxOperator], + x0: NDArray, + tau: float, + epsg: float | NDArray = 1.0, + weights: NDArray | None = None, + eta: float = 1.0, + niter: int = 10, + acceleration: str | None = None, + tol: float | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxfs : :obj:`list` + Proximal operators of the :math:`f_i` functions (must have ``grad`` implemented) + proxgs : :obj:`list` + Proximal operators of the :math:`g_j` functions + x0 : :obj:`numpy.ndarray` + Initial vector + tau : :obj:`float` + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\sum_{i=1}^n \nabla f_i`. + epsg : :obj:`float` or :obj:`numpy.ndarray`, optional + Scaling factor(s) of ``g`` function(s) + weights : :obj:`float`, optional + Weighting factors of ``g`` functions. Must sum to 1. + eta : :obj:`float`, optional + Relaxation parameter (must be between 0 and 1, 0 excluded). Note that + this will be only used when ``acceleration=None``. + niter : :obj:`int`, optional + Number of iterations of iterative scheme + acceleration: :obj:`str`, optional + Acceleration (``None``, ``vandenberghe`` or ``fista``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping + criterion). If ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model + iiter : :obj:`int` + Number of executed iterations + pfgs : :obj:`numpy.ndarray` + History of the objective function + + """ + x, y = self.setup( + proxfs=proxfs, + proxgs=proxgs, + x0=x0, + tau=tau, + epsg=epsg, + weights=weights, + eta=eta, + niter=niter, + acceleration=acceleration, + tol=tol, + show=show, + ) + + x, y = self.run(x, y, niter, show=show, itershow=itershow) + self.finalize(show) + return x, y, self.iiter, self.pfgs diff --git a/pytests/test_solver.py b/pytests/test_solver.py index 2c5b5dc..a2105a0 100644 --- a/pytests/test_solver.py +++ b/pytests/test_solver.py @@ -4,6 +4,7 @@ import pytest from numpy.testing import assert_array_almost_equal from pylops.basicoperators import Identity, MatrixMult +from pylops.optimization.sparsity import fista, ista from pyproximal.optimization.primal import ( ADMM, @@ -23,6 +24,14 @@ par2 = {"n": 8, "m": 10, "dtype": "float64"} # float32 +def test_ProximalGradient_unknown_acceleration(): + """Check that an error is raised if an unknown acceleration + method is provided to ProximalGradient solver + """ + with pytest.raises(NotImplementedError, match="Acceleration should "): + _ = ProximalGradient(proxf=L2(), proxg=L1(), x0=None, acceleration="unknown") + + def test_HQS_noinitial(): """Check that an error is raised if no initial value is provided to HQS solver @@ -161,6 +170,62 @@ def test_ProximalPoint(par): assert_array_almost_equal(xpp, x, decimal=2) +@pytest.mark.parametrize("par", [(par1), (par2)]) +def test_PG_ISTA(par): + """Check equivalency of ProximalGradient and ISTA/FISTA (PyLops)""" + np.random.seed(0) + n, m = par["n"], par["m"] + + # Define sparse model + x = np.zeros(m) + x[2], x[4] = 1, 0.5 + + # Random mixing matrix + R = np.random.normal(0.0, 1.0, (n, m)) + Rop = MatrixMult(R) + + y = Rop @ x + + # Step size + L = (Rop.H * Rop).eigs(1).real + tau = 0.99 / L + + for solver, acceleration in zip( + [ista, fista], + [None, "fista"], + strict=True, + ): + # ISTA/FISTA + eps = 5e-1 + xista = solver( + Rop, + y, + niter=100, + alpha=tau, + eps=eps, + tol=1e-8, + monitorres=False, + show=False, + )[0] + + # PG + l2 = L2(Op=Rop, b=y) + l1 = L1() + epsg = eps * 0.5 # to compensate for 0.5 in ISTA: thresh = eps * alpha * 0.5 + xpg = ProximalGradient( + l2, + l1, + x0=np.zeros(m), + tau=tau, + epsg=epsg, + acceleration=acceleration, + niter=100, + tol=1e-8, + ) + + assert_array_almost_equal(xpg, xista, decimal=2) + + @pytest.mark.parametrize("par", [(par1), (par2)]) def test_PG_GPG(par): """Check equivalency of ProximalGradient and GeneralizedProximalGradient when using From cd4f78d32563c6ea42bbc751645dfefca5538e68 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Sat, 18 Apr 2026 20:22:47 +0100 Subject: [PATCH 05/24] feat: added class based HQS --- pyproximal/optimization/cls_primal.py | 409 ++++++++++++++++++++++++-- pyproximal/optimization/primal.py | 1 + 2 files changed, 391 insertions(+), 19 deletions(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index d0f8418..3a99929 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -218,7 +218,7 @@ def setup( # type: ignore[override] # create variables to track the objective function and iterations self.pf, self.pfold = np.inf, np.inf - self.pfs: list[float] = [] + self.cost: list[float] = [] self.tolbreak = False self.iiter = 0 @@ -258,7 +258,7 @@ def step(self, x: NDArray, show: bool = False) -> NDArray: if show: self._print_step(x) if self.tol is not None or show: - self.pfs.append(float(self.pf)) + self.cost.append(float(self.pf)) return x def run( @@ -325,6 +325,7 @@ def finalize(self, show: bool = False) -> None: """ self.tend = time.time() self.telapsed = self.tend - self.tstart + if show: self._print_finalize(nbar=60) @@ -366,14 +367,14 @@ def solve( # type: ignore[override] Estimated model iiter : :obj:`int` Number of executed iterations - pfs : :obj:`numpy.ndarray` + cost : :obj:`list` History of the objective function """ x = self.setup(prox=prox, x0=x0, tau=tau, niter=niter, tol=tol, show=show) x = self.run(x, niter, show=show, itershow=itershow) self.finalize(show) - return x, self.iiter, self.pfs + return x, self.iiter, self.cost class ProximalGradient(Solver): @@ -580,8 +581,7 @@ def setup( # type: ignore[override] # create variables to track the objective function and iterations self.pfg, self.pfgold = np.inf, np.inf - self.pfs: list[float] = [] - self.pfgs: list[float] = [] + self.cost: list[float] = [] self.tolbreak = False self.iiter = 0 @@ -679,7 +679,7 @@ def step( if show: self._print_step(x, pf, pg) if self.tol is not None or show: - self.pfgs.append(float(self.pfg)) + self.cost.append(float(self.pfg)) return x, y def run( @@ -752,6 +752,7 @@ def finalize(self, show: bool = False) -> None: """ self.tend = time.time() self.telapsed = self.tend - self.tstart + if show: self._print_finalize(nbar=81) @@ -823,7 +824,7 @@ def solve( # type: ignore[override] Additional estimated model iiter : :obj:`int` Number of executed iterations - pfgs : :obj:`numpy.ndarray` + cost : :obj:`list` History of the objective function """ @@ -845,7 +846,7 @@ def solve( # type: ignore[override] x, y = self.run(x, y, niter, show=show, itershow=itershow) self.finalize(show) - return x, y, self.iiter, self.pfgs + return x, y, self.iiter, self.cost class AndersonProximalGradient(Solver): @@ -1023,10 +1024,8 @@ def setup( # type: ignore[override] # create variables to track the objective function and iterations self.pf = self.proxf(x) - print("Initial pf: %10.3e" % self.pf) self.pfg, self.pfgold = np.inf, np.inf - self.pfs: list[float] = [] - self.pfgs: list[float] = [] + self.cost: list[float] = [] self.tolbreak = False self.iiter = 0 @@ -1120,7 +1119,7 @@ def step( if show: self._print_step(x, pg) if self.tol is not None or show: - self.pfgs.append(float(self.pfg)) + self.cost.append(float(self.pfg)) return x, y def run( @@ -1193,6 +1192,7 @@ def finalize(self, show: bool = False) -> None: """ self.tend = time.time() self.telapsed = self.tend - self.tstart + if show: self._print_finalize(nbar=81) @@ -1255,7 +1255,7 @@ def solve( # type: ignore[override] Additional estimated model iiter : :obj:`int` Number of executed iterations - pfgs : :obj:`numpy.ndarray` + cost : :obj:`list` History of the objective function """ @@ -1275,7 +1275,7 @@ def solve( # type: ignore[override] x, y = self.run(x, y, niter, show=show, itershow=itershow) self.finalize(show) - return x, y, self.iiter, self.pfgs + return x, y, self.iiter, self.cost class GeneralizedProximalGradient(Solver): @@ -1454,7 +1454,7 @@ def setup( # type: ignore[override] # create variables to track the objective function and iterations self.pfg, self.pfgold = np.inf, np.inf - self.pfgs: list[float] = [] + self.cost: list[float] = [] self.tolbreak = False self.iiter = 0 @@ -1534,7 +1534,7 @@ def step( if show: self._print_step(x, pf, pg) if self.tol is not None or show: - self.pfgs.append(float(self.pfg)) + self.cost.append(float(self.pfg)) return x, y def run( @@ -1607,6 +1607,7 @@ def finalize(self, show: bool = False) -> None: """ self.tend = time.time() self.telapsed = self.tend - self.tstart + if show: self._print_finalize(nbar=65) @@ -1668,7 +1669,7 @@ def solve( # type: ignore[override] Additional estimated model iiter : :obj:`int` Number of executed iterations - pfgs : :obj:`numpy.ndarray` + cost : :obj:`list` History of the objective function """ @@ -1688,4 +1689,374 @@ def solve( # type: ignore[override] x, y = self.run(x, y, niter, show=show, itershow=itershow) self.finalize(show) - return x, y, self.iiter, self.pfgs + return x, y, self.iiter, self.cost + + +class HQS(Solver): + r"""Half Quadratic splitting + + Solves the following minimization problem using Half Quadratic splitting + algorithm: + + .. math:: + + \mathbf{x},\mathbf{z} = \argmin_{\mathbf{x},\mathbf{z}} + f(\mathbf{x}) + g(\mathbf{z}) \\ + s.t. \; \mathbf{x}=\mathbf{z} + + where :math:`f(\mathbf{x})` and :math:`g(\mathbf{z})` are any convex + function that has a known proximal operator. + + Notes + ----- + The HQS algorithm can be expressed by the following recursion [1]_: + + .. math:: + + \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{x}^{k}) \\ + \mathbf{x}^{k+1} = \prox_{\tau f}(\mathbf{z}^{k+1}) + + for ``gfirst=False``, or + + .. math:: + + \mathbf{x}^{k+1} = \prox_{\tau f}(\mathbf{z}^{k}) \\ + \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{x}^{k+1}) + + for ``gfirst=False``. Note that ``x`` and ``z`` converge to each other, + however if iterations are stopped too early ``x`` is guaranteed to belong to + the domain of ``f`` while ``z`` is guaranteed to belong to the domain of ``g``. + Depending on the problem either of the two may be the best solution. + + .. [1] D., Geman, and C., Yang, "Nonlinear image recovery with halfquadratic + regularization", IEEE Transactions on Image Processing, + 4, 7, pp. 932-946, 1995. + + """ + + def _print_setup(self, tau_print: str, xcomplex: bool = False) -> None: + self._print_solver(nbar=65) + + strpar = ( + f"Proximal operator (f): {type(self.proxf).__name__}\n" + f"Proximal operator (g): {type(self.proxg).__name__}\n" + ) + strpar1 = f"tau = {tau_print}\tniter = {self.niter}" + print(strpar) + print(strpar1) + print("-" * 65 + "\n") + if not xcomplex: + head1 = " Itn x[0] f g J=f+g" + else: + head1 = ( + " Itn x[0] f g J=f+g" + ) + print(head1) + + def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: + if self.tol is None: + pf, pg = self.proxf(x), self.proxg(x) + self.pfg = pf + pg + x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = ( + f"{self.iiter:6g} " + + strx + + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e}" + ) + print(msg) + + def memory_usage( + self, + show: bool = False, + unit: Tmemunit = "B", + ) -> None: + pass + + def setup( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + tau: float | NDArray, + z0: NDArray | None = None, + niter: int = 10, + gfirst: bool = True, + tol: float | None = None, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Setup solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector (not required when ``gfirst=False``, can pass ``None``) + tau : :obj:`float` or :obj:`numpy.ndarray` + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. Finally note that + :math:`\tau` can be chosen to be a vector of size ``niter`` such that + different :math:`\tau` is used at different iterations (i.e., continuation + strategy) + z0 : :obj:`numpy.ndarray`, optional + Initial z vector (not required when ``gfirst=True``) + niter : :obj:`int` + Number of iterations of iterative scheme + gfirst : :obj:`bool`, optional + Apply Proximal of operator ``g`` first (``True``) or Proximal of + operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + z : :obj:`numpy.ndarray` + Initial guess for the auxiliary variable + + Raises + ------ + ValueError + If both ``x0`` and ``z0`` are set to ``None`` + + """ + self.proxf = proxf + self.proxg = proxg + self.niter = niter + self.gfirst = gfirst + self.tol = tol + + self.ncp = get_array_module(x0) + + # check if tau is a vector + self.tau = self.ncp.asarray(tau, dtype=float) + if tau.size == 1: + tau_print = str(self.tau) + self.tau = self.tau * np.ones(niter) + else: + tau_print = "Variable" + + # set initial vectors + x, z = _x0z0_init(x0, z0) + + # for accelaration + self.t = 1.0 + + # create variables to track the objective function and iterations + self.pfg, self.pfgold = np.inf, np.inf + self.cost: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(tau_print, np.iscomplexobj(x0)) + return x, z + + def step( + self, x: NDArray, z: NDArray, show: bool = False + ) -> tuple[NDArray, NDArray]: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + proximal gradient algorithm + z : :obj:`numpy.ndarray` + Additional model vector to be updated by a step of the + proximal gradient algorithm + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + z : :obj:`numpy.ndarray` + Updated additional model vector + + """ + # proximal steps + if self.gfirst: + z = self.proxg.prox(x, self.tau[self.iiter]) + x = self.proxf.prox(z, self.tau[self.iiter]) + else: + x = self.proxf.prox(z, self.tau[self.iiter]) + z = self.proxg.prox(x, self.tau[self.iiter]) + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfgold = self.pfg + pf = self.proxf(x) + pg = self.proxg(x) + self.pfg = pf + pg + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pf, pg = 0.0, 0.0 + + self.iiter += 1 + if show: + self._print_step(x, pf, pg) + if self.tol is not None or show: + self.cost.append(float(self.pfg)) + return x, z + + def run( + self, + x: NDArray, + z: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray]: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the proximal gradient algorithm + z : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the proximal gradient algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + z : :obj:`numpy.ndarray` + Additional estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, z = self.step(x, z, showstep) + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, z + + def finalize(self, show: bool = False) -> None: + r"""Finalize solver + + Parameters + ---------- + show : :obj:`bool`, optional + Display finalize log + + """ + self.tend = time.time() + self.telapsed = self.tend - self.tstart + + if show: + self._print_finalize(nbar=65) + + def solve( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + tau: float | NDArray, + z0: NDArray | None = None, + niter: int = 10, + gfirst: bool = True, + tol: float | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector (not required when ``gfirst=False``, can pass ``None``) + tau : :obj:`float` or :obj:`numpy.ndarray` + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. Finally note that + :math:`\tau` can be chosen to be a vector of size ``niter`` such that + different :math:`\tau` is used at different iterations (i.e., continuation + strategy) + z0 : :obj:`numpy.ndarray`, optional + Initial z vector (not required when ``gfirst=True``) + niter : :obj:`int` + Number of iterations of iterative scheme + gfirst : :obj:`bool`, optional + Apply Proximal of operator ``g`` first (``True``) or Proximal of + operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping + criterion). If ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model + iiter : :obj:`int` + Number of executed iterations + cost : :obj:`list` + History of the objective function + + """ + x, z = self.setup( + proxf=proxf, + proxg=proxg, + x0=x0, + tau=tau, + z0=z0, + niter=niter, + gfirst=gfirst, + tol=tol, + show=show, + ) + + x, z = self.run(x, z, niter, show=show, itershow=itershow) + self.finalize(show) + return x, z, self.iiter, self.cost diff --git a/pyproximal/optimization/primal.py b/pyproximal/optimization/primal.py index fad1fe7..12ad1e4 100644 --- a/pyproximal/optimization/primal.py +++ b/pyproximal/optimization/primal.py @@ -332,6 +332,7 @@ def ProximalGradient( if acceleration not in [None, "None", "vandenberghe", "fista"]: msg = "Acceleration should be None, vandenberghe or fista" raise NotImplementedError(msg) + if show: tstart = time.time() print( From 014158aef5b1377be86505f42a5093f7e775e583 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Sat, 18 Apr 2026 20:24:38 +0100 Subject: [PATCH 06/24] feat: added HQS tests --- pytests/test_solver.py | 52 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/pytests/test_solver.py b/pytests/test_solver.py index a2105a0..e416f94 100644 --- a/pytests/test_solver.py +++ b/pytests/test_solver.py @@ -4,6 +4,7 @@ import pytest from numpy.testing import assert_array_almost_equal from pylops.basicoperators import Identity, MatrixMult +from pylops.optimization.leastsquares import regularized_inversion from pylops.optimization.sparsity import fista, ista from pyproximal.optimization.primal import ( @@ -274,6 +275,57 @@ def test_PG_GPG(par): assert_array_almost_equal(xpg, xgpg, decimal=2) +@pytest.mark.parametrize("par", [(par1), (par2)]) +def test_HQS_ADMM_L2(par): + """Check that HQS/ADMM can be used to solved a pure L2-based objective function + (and compare with LSQR - note that despite the trajectory will be different, + they should converge to the same solution) + """ + np.random.seed(0) + n, m = par["n"], par["m"] + + # Define sparse model + x = np.random.normal(0.0, 1.0, m).astype(par["dtype"]) + + # Random mixing matrix + R = np.random.normal(0.0, 1.0, (n, m)).astype(par["dtype"]) + Rop = MatrixMult(R, dtype=par["dtype"]) + + y = Rop @ x + + # Step size + L = (Rop.H * Rop).eigs(1).real + tau = 0.99 / L + eps = 1e-1 + + # L2 + Iop = Identity(m, dtype=par["dtype"]) + xl2 = regularized_inversion( + Rop, + y, + Regs=[ + Iop, + ], + epsRs=[ + np.sqrt(eps), + ], + iter_lim=1000, + )[0] + + # HQS + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l2reg = L2(sigma=eps) + xhqs = HQS(l2, l2reg, x0=np.zeros(m), tau=tau, niter=1000)[0] + + # ADMM + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l2reg = L2(sigma=eps) + xadmm = ADMM(l2, l2reg, x0=np.zeros(m), tau=tau, niter=1000)[0] + + assert_array_almost_equal(xl2, xhqs, decimal=2) + assert_array_almost_equal(xl2, xadmm, decimal=2) + + @pytest.mark.parametrize("par", [(par1), (par2)]) def test_ADMM_DRS(par): """Check equivalency of ADMM and DouglasRachfordSplitting From c302ec6264c36bf45d9f861343d4a7be6c8cf421 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Sat, 18 Apr 2026 20:43:32 +0100 Subject: [PATCH 07/24] minor: added Tmemunit to typing --- pyproximal/optimization/cls_primal.py | 3 ++- pyproximal/utils/typing.py | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index 3a99929..2f7caf0 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -10,11 +10,12 @@ import pylops from pylops.optimization.callback import Callbacks from pylops.utils.backend import get_array_module, to_numpy -from pylops.utils.typing import NDArray, Tmemunit +from pylops.utils.typing import NDArray from pyproximal.optimization.basesolver import Solver from pyproximal.ProxOperator import ProxOperator from pyproximal.utils.bilinear import BilinearOperator +from pyproximal.utils.typing import Tmemunit if TYPE_CHECKING: from pylops.linearoperator import LinearOperator diff --git a/pyproximal/utils/typing.py b/pyproximal/utils/typing.py index 05c801c..a90e103 100644 --- a/pyproximal/utils/typing.py +++ b/pyproximal/utils/typing.py @@ -4,8 +4,11 @@ ] from collections.abc import Callable +from typing import Literal from pylops.utils.typing import NDArray FloatCallableLike = float | NDArray | Callable[[int], float | NDArray] IntCallableLike = int | Callable[[int], int] + +Tmemunit = Literal["B", "KB", "MB", "GB"] From 21949bafdd0f64c30d79de03128b227e3d8c9579 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Sat, 18 Apr 2026 20:46:09 +0100 Subject: [PATCH 08/24] minor: fix mypy error --- pyproximal/optimization/cls_primal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index 2f7caf0..fc6c842 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -1839,7 +1839,7 @@ def setup( # type: ignore[override] # check if tau is a vector self.tau = self.ncp.asarray(tau, dtype=float) - if tau.size == 1: + if self.tau.size == 1: tau_print = str(self.tau) self.tau = self.tau * np.ones(niter) else: From 63edabc3e05f2ba5dd3cc1b4442081d2bdac02a4 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Mon, 20 Apr 2026 21:51:40 +0100 Subject: [PATCH 09/24] feat: added finalize and callback to basesolver --- pyproximal/optimization/basesolver.py | 59 ++++ pyproximal/optimization/cls_primal.py | 464 +++++++++++++++++++++----- 2 files changed, 439 insertions(+), 84 deletions(-) diff --git a/pyproximal/optimization/basesolver.py b/pyproximal/optimization/basesolver.py index 222979a..954b2ab 100644 --- a/pyproximal/optimization/basesolver.py +++ b/pyproximal/optimization/basesolver.py @@ -8,6 +8,7 @@ from pylops.optimization.basesolver import Solver as pSolver from pylops.optimization.callback import Callbacks +from pylops.utils.typing import NDArray if TYPE_CHECKING: from pyproximal.ProxOperator import ProxOperator @@ -134,3 +135,61 @@ def solve( """ pass + + def finalize(self, nbar: int = 60, show: bool = False) -> None: + r"""Finalize solver + + Parameters + ---------- + nbar : :obj:`int`, optional + Number of ``-`` in the bar dividing iterations + from finalize messages in the print message of + the solver + show : :obj:`bool`, optional + Display finalize log + + """ + self.tend = time.time() + self.telapsed = self.tend - self.tstart + + if show: + self._print_finalize(nbar=nbar) + + def callback( # noqa: B027 + self, + x: NDArray, + z: NDArray | None = None, + *args, + **kwargs, + ) -> None: + """Callback routine + + This routine must be passed by the user. Its function signature must contain + either a single input that contains the current solution or two inputs + that contain the current solutions for methods that apply splitting + (when using the `solve` method it will be automatically invoked after + each step of the solve) + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current solution + z : :obj:`numpy.ndarray` + Current additional solution + + Examples + -------- + >>> import numpy as np + >>> from pyproximal.optimization.cls_primal import ADMM + >>> def callback(x, z): + ... print(f"Running callback, current solutions {x} - {z}") + ... + >>> admmsolve.callback = callback + + >>> x = np.ones(2) + >>> z = np.zeros(2) + >>> admmsolve.callback(x, z) + Running callback, current solutions [1. 1.] - [0. 0.] + + """ + pass diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index fc6c842..9019eda 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -2,7 +2,6 @@ "ProximalPoint", ] -import time from collections.abc import Sequence from typing import TYPE_CHECKING, Optional, cast @@ -315,21 +314,6 @@ def run( break return x - def finalize(self, show: bool = False) -> None: - r"""Finalize solver - - Parameters - ---------- - show : :obj:`bool`, optional - Display finalize log - - """ - self.tend = time.time() - self.telapsed = self.tend - self.tstart - - if show: - self._print_finalize(nbar=60) - def solve( # type: ignore[override] self, prox: "ProxOperator", @@ -374,7 +358,7 @@ def solve( # type: ignore[override] """ x = self.setup(prox=prox, x0=x0, tau=tau, niter=niter, tol=tol, show=show) x = self.run(x, niter, show=show, itershow=itershow) - self.finalize(show) + self.finalize(60, show) return x, self.iiter, self.cost @@ -573,7 +557,7 @@ def setup( # type: ignore[override] msg = "Acceleration should be None, vandenberghe or fista" raise NotImplementedError(msg) - # set initial vectors + # initialize solver x = x0.copy() y = x.copy() @@ -742,21 +726,6 @@ def run( break return x, y - def finalize(self, show: bool = False) -> None: - r"""Finalize solver - - Parameters - ---------- - show : :obj:`bool`, optional - Display finalize log - - """ - self.tend = time.time() - self.telapsed = self.tend - self.tstart - - if show: - self._print_finalize(nbar=81) - def solve( # type: ignore[override] self, proxf: ProxOperator, @@ -846,7 +815,7 @@ def solve( # type: ignore[override] ) x, y = self.run(x, y, niter, show=show, itershow=itershow) - self.finalize(show) + self.finalize(81, show) return x, y, self.iiter, self.cost @@ -1007,7 +976,7 @@ def setup( # type: ignore[override] else: epsg_print = "Multi" - # set initial vectors + # initialize solver y = x0 - self.tau * proxf.grad(x0) x = self.proxg.prox(y, self.epsg[0] * self.tau) @@ -1182,21 +1151,6 @@ def run( break return x, y - def finalize(self, show: bool = False) -> None: - r"""Finalize solver - - Parameters - ---------- - show : :obj:`bool`, optional - Display finalize log - - """ - self.tend = time.time() - self.telapsed = self.tend - self.tstart - - if show: - self._print_finalize(nbar=81) - def solve( # type: ignore[override] self, proxf: ProxOperator, @@ -1275,7 +1229,7 @@ def solve( # type: ignore[override] ) x, y = self.run(x, y, niter, show=show, itershow=itershow) - self.finalize(show) + self.finalize(81, show) return x, y, self.iiter, self.cost @@ -1445,7 +1399,7 @@ def setup( # type: ignore[override] msg = "Acceleration should be None, vandenberghe or fista" raise NotImplementedError(msg) - # set initial vectors + # initialize solver x = x0.copy() y = x.copy() self.zs = [x.copy() for _ in range(len(proxgs))] @@ -1597,21 +1551,6 @@ def run( break return x, y - def finalize(self, show: bool = False) -> None: - r"""Finalize solver - - Parameters - ---------- - show : :obj:`bool`, optional - Display finalize log - - """ - self.tend = time.time() - self.telapsed = self.tend - self.tstart - - if show: - self._print_finalize(nbar=65) - def solve( # type: ignore[override] self, proxfs: list[ProxOperator], @@ -1689,7 +1628,7 @@ def solve( # type: ignore[override] ) x, y = self.run(x, y, niter, show=show, itershow=itershow) - self.finalize(show) + self.finalize(65, show) return x, y, self.iiter, self.cost @@ -1743,8 +1682,10 @@ def _print_setup(self, tau_print: str, xcomplex: bool = False) -> None: f"Proximal operator (g): {type(self.proxg).__name__}\n" ) strpar1 = f"tau = {tau_print}\tniter = {self.niter}" + strpar2 = f"gfirst = {self.gfirst}\ttol = {self.tol}" print(strpar) print(strpar1) + print(strpar2) print("-" * 65 + "\n") if not xcomplex: head1 = " Itn x[0] f g J=f+g" @@ -1840,12 +1781,12 @@ def setup( # type: ignore[override] # check if tau is a vector self.tau = self.ncp.asarray(tau, dtype=float) if self.tau.size == 1: - tau_print = str(self.tau) + tau_print = str(np.round(self.tau, 6)) self.tau = self.tau * np.ones(niter) else: tau_print = "Variable" - # set initial vectors + # initialize solver x, z = _x0z0_init(x0, z0) # for accelaration @@ -1972,20 +1913,375 @@ def run( break return x, z - def finalize(self, show: bool = False) -> None: - r"""Finalize solver + def solve( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + tau: float | NDArray, + z0: NDArray | None = None, + niter: int = 10, + gfirst: bool = True, + tol: float | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector (not required when ``gfirst=False``, can pass ``None``) + tau : :obj:`float` or :obj:`numpy.ndarray` + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. Finally note that + :math:`\tau` can be chosen to be a vector of size ``niter`` such that + different :math:`\tau` is used at different iterations (i.e., continuation + strategy) + z0 : :obj:`numpy.ndarray`, optional + Initial z vector (not required when ``gfirst=True``) + niter : :obj:`int` + Number of iterations of iterative scheme + gfirst : :obj:`bool`, optional + Apply Proximal of operator ``g`` first (``True``) or Proximal of + operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping + criterion). If ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model + iiter : :obj:`int` + Number of executed iterations + cost : :obj:`list` + History of the objective function + + """ + x, z = self.setup( + proxf=proxf, + proxg=proxg, + x0=x0, + tau=tau, + z0=z0, + niter=niter, + gfirst=gfirst, + tol=tol, + show=show, + ) + + x, z = self.run(x, z, niter, show=show, itershow=itershow) + self.finalize(65, show) + return x, z, self.iiter, self.cost + + +class ADMM(Solver): + r"""Alternating Direction Method of Multipliers + + Solves the following minimization problem using Alternating Direction + Method of Multipliers: + + .. math:: + + \mathbf{x},\mathbf{z} = \argmin_{\mathbf{x},\mathbf{z}} + f(\mathbf{x}) + g(\mathbf{z}) \\ + s.t. \; \mathbf{x}=\mathbf{z} + + where :math:`f(\mathbf{x})` and :math:`g(\mathbf{z})` are any convex + function that has a known proximal operator. + + ADMM can also solve the problem of the form above with a more general + constraint: :math:`\mathbf{Ax}+\mathbf{Bz}=\mathbf{c}`. This routine implements + the special case where :math:`\mathbf{A}=\mathbf{I}`, :math:`\mathbf{B}=-\mathbf{I}`, + and :math:`\mathbf{c}=\mathbf{0}`, as a general algorithm can be obtained for any choice of + :math:`f` and :math:`g` provided they have a known proximal operator. + + On the other hand, for more general choice of :math:`\mathbf{A}`, :math:`\mathbf{B}`, + and :math:`\mathbf{c}`, the iterations are not generalizable, i.e. they depend on the choice of + the :math:`f` and :math:`g` functions. For this reason, we currently only provide an additional + solver for the special case where :math:`f` is a :class:`pyproximal.proximal.L2` + operator with a linear operator :math:`\mathbf{G}` and data :math:`\mathbf{y}`, + :math:`\mathbf{B}=-\mathbf{I}` and :math:`\mathbf{c}=\mathbf{0}`, + called :func:`pyproximal.optimization.primal.ADMML2`. Note that for the very same choice + of :math:`\mathbf{B}` and :math:`\mathbf{c}`, the :func:`pyproximal.optimization.primal.LinearizedADMM` + can also be used (and this does not require a specific choice of :math:`f`). + + See Also + -------- + ADMML2: ADMM with L2 misfit function + LinearizedADMM: Linearized ADMM + + Notes + ----- + The ADMM algorithm can be expressed by the following recursion [1]_: + + .. math:: + + \mathbf{x}^{k+1} = \prox_{\tau f}(\mathbf{z}^{k} - \mathbf{u}^{k})\\ + \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{x}^{k+1} + \mathbf{u}^{k})\\ + \mathbf{u}^{k+1} = \mathbf{u}^{k} + \mathbf{x}^{k+1} - \mathbf{z}^{k+1} + + Note that ``x`` and ``z`` converge to each other, however if iterations are + stopped too early ``x`` is guaranteed to belong to the domain of ``f`` + while ``z`` is guaranteed to belong to the domain of ``g``. Depending on + the problem either of the two may be the best solution. + + .. [1] S. Boyd, N. Parikh, E. Chu, B. Peleato, and J. Eckstein. 2011. + Distributed optimization and statistical learning via the alternating + direction method of multipliers. Foundations and Trends in Machine + Learning, 3 (1), 1-122. https://doi.org/10.1561/2200000016. + + """ + + def _print_setup(self, xcomplex: bool = False) -> None: + self._print_solver(nbar=65) + + strpar = ( + f"Proximal operator (f): {type(self.proxf).__name__}\n" + f"Proximal operator (g): {type(self.proxg).__name__}\n" + ) + strpar1 = f"tau = {self.tau:6e}\tniter = {self.niter}" + strpar2 = f"gfirst = {self.gfirst}\t\ttol = {self.tol}" + print(strpar) + print(strpar1) + print(strpar2) + + print("-" * 65 + "\n") + if not xcomplex: + head1 = " Itn x[0] f g J=f+g" + else: + head1 = ( + " Itn x[0] f g J=f+g" + ) + print(head1) + + def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: + if self.tol is None: + pf, pg = self.proxf(x), self.proxg(x) + self.pfg = pf + pg + x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = ( + f"{self.iiter:6g} " + + strx + + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e}" + ) + print(msg) + + def memory_usage( + self, + show: bool = False, + unit: Tmemunit = "B", + ) -> None: + pass + + def setup( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + tau: float | NDArray, + z0: NDArray | None = None, + niter: int = 10, + gfirst: bool = False, + tol: float | None = None, + callbackz: bool = False, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Setup solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector (not required when ``gfirst=False``, can pass ``None``) + tau : :obj:`float` + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. + z0 : :obj:`numpy.ndarray`, optional + Initial z vector (not required when ``gfirst=True``) + niter : :obj:`int` + Number of iterations of iterative scheme + gfirst : :obj:`bool`, optional + Apply Proximal of operator ``g`` first (``True``) or Proximal of + operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + callbackz : :obj:`bool`, optional + Modify callback signature to (``callback(x, z)``) when ``callbackz=True`` + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + z : :obj:`numpy.ndarray` + Initial guess for the auxiliary variable + + Raises + ------ + ValueError + If both ``x0`` and ``z0`` are set to ``None`` + + """ + self.proxf = proxf + self.proxg = proxg + self.tau = tau + self.niter = niter + self.gfirst = gfirst + self.tol = tol + self.callbackz = callbackz + + self.ncp = get_array_module(x0) + + # initialize solver + x, z = _x0z0_init(x0, z0) + self.u = self.ncp.zeros_like(x) + + # create variables to track the objective function and iterations + self.pfg, self.pfgold = np.inf, np.inf + self.cost: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(np.iscomplexobj(x0)) + return x, z + + def step( + self, x: NDArray, z: NDArray, show: bool = False + ) -> tuple[NDArray, NDArray]: + r"""Run one step of solver Parameters ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + proximal gradient algorithm + z : :obj:`numpy.ndarray` + Additional model vector to be updated by a step of the + proximal gradient algorithm show : :obj:`bool`, optional - Display finalize log + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + z : :obj:`numpy.ndarray` + Updated additional model vector """ - self.tend = time.time() - self.telapsed = self.tend - self.tstart + # proximal steps + if self.gfirst: + z = self.proxg.prox(x + self.u, self.tau) + x = self.proxf.prox(z - self.u, self.tau) + else: + x = self.proxf.prox(z - self.u, self.tau) + z = self.proxg.prox(x + self.u, self.tau) + self.u = self.u + x - z + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfgold = self.pfg + pf = self.proxf(x) + pg = self.proxg(x) + self.pfg = pf + pg + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pf, pg = 0.0, 0.0 + self.iiter += 1 if show: - self._print_finalize(nbar=65) + self._print_step(x, pf, pg) + if self.tol is not None or show: + self.cost.append(float(self.pfg)) + return x, z + + def run( + self, + x: NDArray, + z: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray]: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the proximal gradient algorithm + z : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the proximal gradient algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + z : :obj:`numpy.ndarray` + Additional estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, z = self.step(x, z, showstep) + if self.callbackz: + self.callback(x, z) + else: + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, z def solve( # type: ignore[override] self, @@ -1995,8 +2291,9 @@ def solve( # type: ignore[override] tau: float | NDArray, z0: NDArray | None = None, niter: int = 10, - gfirst: bool = True, + gfirst: bool = False, tol: float | None = None, + callbackz: bool = False, show: bool = False, itershow: tuple[int, int, int] = (10, 10, 10), ) -> tuple[NDArray, NDArray, int, NDArray]: @@ -2010,13 +2307,10 @@ def solve( # type: ignore[override] Proximal operator of g function x0 : :obj:`numpy.ndarray` Initial vector (not required when ``gfirst=False``, can pass ``None``) - tau : :obj:`float` or :obj:`numpy.ndarray` + tau : :obj:`float` Positive scalar weight, which should satisfy the following condition to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is - the Lipschitz constant of :math:`\nabla f`. Finally note that - :math:`\tau` can be chosen to be a vector of size ``niter`` such that - different :math:`\tau` is used at different iterations (i.e., continuation - strategy) + the Lipschitz constant of :math:`\nabla f`. z0 : :obj:`numpy.ndarray`, optional Initial z vector (not required when ``gfirst=True``) niter : :obj:`int` @@ -2027,6 +2321,8 @@ def solve( # type: ignore[override] tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached + callbackz : :obj:`bool`, optional + Modify callback signature to (``callback(x, z)``) when ``callbackz=True`` show : :obj:`bool`, optional Display logs itershow : :obj:`tuple`, optional @@ -2038,7 +2334,7 @@ def solve( # type: ignore[override] ------- x : :obj:`numpy.ndarray` Estimated model - y : :obj:`numpy.ndarray` + z : :obj:`numpy.ndarray` Additional estimated model iiter : :obj:`int` Number of executed iterations @@ -2059,5 +2355,5 @@ def solve( # type: ignore[override] ) x, z = self.run(x, z, niter, show=show, itershow=itershow) - self.finalize(show) + self.finalize(65, show) return x, z, self.iiter, self.cost From e6d59d11389853355762939ef1ca3eabaf457902 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Sun, 31 May 2026 21:07:49 +0100 Subject: [PATCH 10/24] minor: fix mypy error --- pyproximal/optimization/basesolver.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproximal/optimization/basesolver.py b/pyproximal/optimization/basesolver.py index 954b2ab..b710ab5 100644 --- a/pyproximal/optimization/basesolver.py +++ b/pyproximal/optimization/basesolver.py @@ -159,8 +159,8 @@ def callback( # noqa: B027 self, x: NDArray, z: NDArray | None = None, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: """Callback routine From c3842b2155d5cbb3624273ed74dd7e6c2b56ce82 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Mon, 8 Jun 2026 20:50:59 +0100 Subject: [PATCH 11/24] doc: added page for adding solvers --- docs/source/addingsolver.rst | 300 +++++++++++++++++++++++++++++++++++ docs/source/index.rst | 1 + 2 files changed, 301 insertions(+) create mode 100755 docs/source/addingsolver.rst diff --git a/docs/source/addingsolver.rst b/docs/source/addingsolver.rst new file mode 100755 index 0000000..dbafb03 --- /dev/null +++ b/docs/source/addingsolver.rst @@ -0,0 +1,300 @@ +.. _addingsolver: + +|:fire:| Implementing new solvers +################################# + +Users are welcome to create new solvers and add them to the PyProximal library. + +In this tutorial, we will go through the key steps in the definition of a solver, using a +sligthly simplified version of :py:class:`pyproximal.optimization.cls_primal.ProximalGradient` +as an example. + +.. note:: + In case you are already familiar with PyLops and have created a solver before, these guidelines closely + follow those in the PyLops documentation. We aim to highlight every step where the creation + of a PyProximal solver differs from that of a PyLops server or steps that are completely new + for this kind of solvers. + +Creating the solver +------------------- +The first thing we need to do is to locate a file containing solvers in the same family of the solver we plan to +include, or create a new file with the name of the solver we would like to implement (or preferably its family). +Note that as the solver will be a class, we need to follow the UpperCaseCamelCase convention for the class itself +but not for the filename. + +At this point we can start by importing the modules that will be needed by the solver. +This varies from solver to solver, however you will always need to import the +:py:class:`pyproximal.optimization.basesolver.Solver` which will be used as *parent* class for any of our solvers. +Moreover, we always recommend to import :py:func:`pyproximal.utils.backend.get_array_module` as solvers should be written +in such a way that it can work both with ``numpy`` and ``cupy`` arrays. See later for details. + +.. code-block:: python + + import time + + import numpy as np + + from pyproximal.optimization.basesolver import Solver + from pyproximal.utils.backend import get_array_module + + +After that we define our new object: + +.. code-block:: python + + class ProximalGradient(Solver): + +followed by a `numpydoc docstring `__ +(starting with ``r"""`` and ending with ``"""``) containing the documentation of the solver. Such docstring should +contain at least a short description of the solver, a ``Parameters`` section with a description of the +input parameters of the associated ``_init__`` method and a ``Notes`` section providing a reference to the original +solver and possibly a concise mathematical explanation of the solver. Take a look at some of the core solver of PyProximal +to get a feeling of the level of details of the mathematical explanation. + +As for any Python class, our solver will need an ``__init__`` method. In this case, however, we will just rely on that +of the base class. A single input parameter is passed to the ``__init__`` method and saved as members of our class, +namely the :class:`pyproximal.optimization.callback.Callbacks` object that contains any callback we wish to attach +to the solver. Moreover, two additional parameters are created that contains the counter of the iterations (which +will be incremented every time the ``step`` method is called) and the current time (this is used later to report +the execution time of the solver). Here is the ``__init__`` method of the base class: + +.. code-block:: python + + def __init__(self, callbacks=None): + self.callbacks = callbacks + self._registercallbacks() + self.iiter = 0 + self.tstart = time.time() + +We can now move onto writing the *setup* of the solver in the method ``setup``. We will need to write +a piece of code that prepares the solver prior to being able to apply a step. As most solvers in PyProximal are designed +as the sum of two terms (usually defined as ``f`` and ``g``), the setup method usually takes two proximal operators +(generally represented by ``proxf`` and ``proxg``) plus additional parameters as explained below. However, there are some +special cases in which solvers operate on a single proximal, a list of proximals, or a single proximal and some linear +operator plus a data vector; in that case the number and type of the first few inputs of the ``setup`` method may be +different. + +Next the method usually requires the initial guess of the solver ``x0``, alongside various hyperparameters of the solver +--- e.g., step sizes and parameters involved in the stopping criterion. For example in this case we only have nine parameters. +Two of them, ``niter``, which refers to the maximum allowed number of iterations, and ``tol`` to set the +tolerance on the residual norm (the solver will be stopped if this is smaller than the chosen tolerance) are always present. +Moreover, we always have the possibility to decide whether we want to operate the solver (in this case its setup part) in verbose +or silent mode. This is driven by the ``show`` parameter. We will soon discuss how to choose what to print on screen in +case of verbose mode (``show=True``). + +The setup method can be loosely seen as composed of four parts. First, the inputs are stored as members of the class. +Moreover the type of the ``x0`` vector is checked to evaluate whether to use ``numpy`` or ``cupy`` for algebraic operations +(this is done by ``self.ncp = get_array_module(x0)``). As some of the optional parameters hyperparameters may require +some specialized checking and initialization, this is carried out in this second part (see specific solver implementations +for more details). Next, the starting guess or guesses are initialized (see also note below). Finally, a number +of variables are initialized to be used inside the ``step`` method to keep track of the optimization process. Moreover, +note that the ``setup`` method returns the created starting guess ``x`` or guesses (they are not stored as member of the +class). + +.. code-block:: python + + def setup(self, proxf, proxg, x0, epsg, tau, + niter=10, tol=None): + self.proxf = proxf + self.proxg = proxg + self.x0 = x0 + self.epsg = epsg + self.tau = tau + self.niter = niter + self.tol = tol + + self.ncp = get_array_module(x0) + + # more specialized initializations if needed + # ... + # + + # initialize solver + x = x0.copy() + y = x.copy() + + # create variables to track the objective function and iterations + self.pfg, self.pfgold = np.inf, np.inf + self.cost: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(epsg_print, np.iscomplexobj(x0)) + return x, y + +.. note:: + For most solvers with multiple proximal operators, more than one initial guess is usually required. + Alongside ``x0``, such solvers may have an additional (optional) input parameter called ``z0``. In + this case, we provide a method called ``_x0z0_init`` which takes ``x0`` and ``z0`` and provides + checks as well as default initialization of ``z0`` when not provided (i.e., ``z0=None``). + +At this point, we need to implement the core of the solver, the ``step`` method. Here, we take the input (or inputs for solvers +with multiple variables) at the previous iterate, update it following the rule of the solver of choice, and return it (or them). +The other input parameter required by this method is ``show`` to choose whether we want to print a report of the step on screen or +not. However, if appropriate, a user can add additional input parameters. A simplified version of the step method for +ProximalGradient is: + +.. code-block:: python + + def step(self, x, y, show=False): + xold = x.copy() + + # proximal step + x = self.proxg.prox( + x - self.tau * self.proxf.grad(x), + self.epsg * self.tau + ) + + # tolerance check + if self.tol is not None: + self.pfgold = self.pfg + pf, pg = self.proxf(x), self.proxg(x) + self.pfg = pf + self.epsg * pg + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pf, pg = 0.0, 0.0 + + self.iiter += 1 + if show: + self._print_step(x, pf, pg) + if self.tol is not None or show: + self.cost.append(float(self.pfg)) + return x, y + +Similarly, we also implement a ``run`` method that is in charge of running a number of iterations by repeatedly +calling the ``step`` method. + +.. code-block:: python + + def run(self, x, y, niter, show, itershow): + while self.iiter < niter and not self.tolbreak: + x, y = self.step(x, y, showstep) + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, y + +It is worth noting that any number of callbacks can be attached to the solver; some of these +callbacks can implement a stopping criterion and set the ``stop`` member to True when a given +condition is met. The ``_callback_stop`` method is in change of checking if any of the callbacks +has set ``stop`` to True and in the case break the iterations. + +Finally, it is also usually convenient to implement a ``finalize`` method; this method can do any required post-processing that should +not be applied at the end of each step, rather at the end of the entire optimization process. For ProximalGradient this is not implemented +as the method of the baseclass already performs the needed steps. + +Last but not least, we can wrap it all up in the ``solve`` method. This method takes as input the data, the initial +model and the same hyperparameters of the setup method and runs the entire optimization process. For CG: + +.. code-block:: python + + def solve(self, proxf, proxg, x0, epsg, tau, + niter=10, tol=None, show=False, + itershow=(10, 10, 10)): + x, y = self.setup( + proxf=proxf, + proxg=proxg, + x0=x0, + epsg=epsg, + tau=tau, + niter=niter, + tol=tol, + show=show, + ) + + x, y = self.run(x, y, niter, show=show, itershow=itershow) + self.finalize(81, show) + return x, y, self.iiter, self.cost + +And that's it, we have implemented our first solver operator! + +Although the methods that we just described are enough to implement any solver of choice, we find important to provide +users with feedback during the inversion process. Imagine that the modelling operator is very expensive and can take +minutes (or even hours to run), we don't want to leave a user waiting for hours before they can tell if the solver has +done something meaningful. To avoid such scenario, we can implement so called `_print_*` methods where +``*=solver, setup, step, finalize`` that print on screen some useful information (e.g., first value of the current +estimate, norm of residual, etc.). The ``solver`` and ``finalize`` print are already implemented in the base class, +the other two must be implemented when creating a new solver. When these methods are implemented and a user passes +``show=True`` to the associated method, our solver will provide such information on screen throughout the inverse +process. To better understand how to write such methods, we suggest to look into the source code of the CG method. + +Finally, to be backward compatible with versions of PyProximal, we also want to create a function with the same +name of the class-based solver (but in small letters) which simply instantiates the solver and runs it. This function +is usually placed in the same file of the class-based solver and snake_case should be used for its name. +This function generally takes all the mandatory and optional parameters of the solver as +input and returns some of the most valuable properties of the class-based solver object. An example +for `ProximalGradient` is: + +.. code-block:: python + + def ProximalGradient(proxf, proxg, x0, epsg, tau, niter=10, + tol=None, callback=None, show=False): + pgsolve = cProximalGradient(callbacks=[CostToInitialCallback(rtol), ]) + if callback is not None: + cgsolve.callback = callback + x, _, _, _ = cgsolve.solve( + proxf=proxf, proxg=proxg, x0=x0, epsg=epsg, tau=tau, + niter=niter, tol=tol, show=show, + ) + return x + + +Testing the solver +------------------ +Being able to write a solver is not yet a guarantee of the fact that the solver is correct, or in other words +that the solver can converge to a correct solution. + +We encourage to create a new test within an existing ``test_*.py`` file in the ``pytests`` folder (or in a new file). +We also encourage to test the function-bases solver, as this will implicitly test the underlying class-based solver. + +Generally a test file will start with a number of dictionaries containing different parameters we would like to +use in the testing of one or more solvers. The test itself starts with a *decorator* that contains a list +of all (or some) of dictionaries that will would like to use for our specific operator, followed by +the definition of the test: + +.. code-block:: python + + @pytest.mark.parametrize("par", [(par1),(par2)]) + def test_ProximalGradient(par): + +However, for proximal solvers we often find it difficult to create tests that verify a single solver in isolation +against an expected (or ground-truth) solution due to the convex nature of the objective function. So a strategy +that is commonly used in our test suite is to test two solvers that are capable of solving the same objective +function against each other and verify that they converge to the same solution within some expected tolerance. +We encourage to check our existing tests for inspiration. + + +Documenting the solver +---------------------- +Once the solver has been created, we can add it to the documentation of PyProximal. To do so, simply add the name of +the operator within the ``index.rst`` file in ``docs/source/api`` directory. + +Moreover, in order to facilitate the user of your operator by other users, each test should be included in at least +one of the tutorials as part of the Sphinx-gallery of the documentation of the PyProximal library. The directory +``tutorials`` contains currently available tutorials; if the new solver does not fit any of them, create a new tutorial. + + +Final checklist +--------------- +Before submitting your new solver for review, use the following **checklist** to ensure that your code +adheres to the guidelines of PyProximal: + +- you have added the new solver to a new or existing file in the ``optimization`` directory within the ``pyproximal`` + package. + +- the new class contains at least ``__init__``, ``setup``, ``step``, ``run``, ``finalize``, and ``solve`` methods. + +- each of the above methods have a `numpydoc docstring `__ documenting + at least the input ``Parameters`` and the ``__init__`` method contains also a ``Notes`` section providing a + mathematical explanation of the solver. + +- a new test has been added to an existing ``test_*.py`` file within the ``pytests`` folder. The test should verify + that the new solver converges to the true solution for a well-designed inverse problem (i.e., full rank operator). + +- the new solver is used within at least one *tutorial* (in ``tutorials`` directory). + diff --git a/docs/source/index.rst b/docs/source/index.rst index c5ace64..f854709 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -117,6 +117,7 @@ subclassing the :py:class:`pyproximal.ProxOperator` class and by implementing :caption: Getting involved: adding.rst + addingsolver.rst contributing.rst changelog.rst credits.rst From 7a7b01bf8d78aec0b6245a6484db236fff12b2bd Mon Sep 17 00:00:00 2001 From: mrava87 Date: Mon, 8 Jun 2026 20:51:44 +0100 Subject: [PATCH 12/24] minor: moved memory_usage to basesolver --- pyproximal/optimization/basesolver.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/pyproximal/optimization/basesolver.py b/pyproximal/optimization/basesolver.py index b710ab5..1d1b73d 100644 --- a/pyproximal/optimization/basesolver.py +++ b/pyproximal/optimization/basesolver.py @@ -6,10 +6,13 @@ from abc import ABCMeta, abstractmethod from typing import TYPE_CHECKING, Any +import numpy as np from pylops.optimization.basesolver import Solver as pSolver from pylops.optimization.callback import Callbacks from pylops.utils.typing import NDArray +from pyproximal.utils.typing import Tmemunit + if TYPE_CHECKING: from pyproximal.ProxOperator import ProxOperator @@ -84,6 +87,17 @@ def _print_finalize(self, *args: Any, nbar: int = 80, **kwargs: Any) -> None: ) print("-" * nbar + "\n") + def memory_usage( + self, + show: bool = False, + unit: Tmemunit = "B", + ) -> float: + """Compute memory usage of the solver + + Currently not implemented in PyProximal and therefore just overwritten here. + """ + pass + @abstractmethod def setup( self, @@ -92,7 +106,7 @@ def setup( *args: Any, show: bool = False, **kwargs: Any, - ) -> None: + ) -> NDArray | tuple[NDArray, NDArray]: """Setup solver This method is used to setup the solver. Users can change the function signature @@ -151,7 +165,7 @@ def finalize(self, nbar: int = 60, show: bool = False) -> None: """ self.tend = time.time() self.telapsed = self.tend - self.tstart - + self.cost = np.array(self.cost) if show: self._print_finalize(nbar=nbar) From d50866036359272965bc289a58bfbd81e7bc1602 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Mon, 8 Jun 2026 21:16:40 +0100 Subject: [PATCH 13/24] test: added more tests for solvers --- pytests/test_solver.py | 156 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 141 insertions(+), 15 deletions(-) diff --git a/pytests/test_solver.py b/pytests/test_solver.py index e416f94..c5f38c6 100644 --- a/pytests/test_solver.py +++ b/pytests/test_solver.py @@ -18,11 +18,13 @@ LinearizedADMM, ProximalGradient, ProximalPoint, + TwIST, ) from pyproximal.proximal import L1, L2, Quadratic -par1 = {"n": 8, "m": 10, "dtype": "float32"} # float64 -par2 = {"n": 8, "m": 10, "dtype": "float64"} # float32 +par1 = {"n": 10, "m": 10, "dtype": "float64"} # square, float64 +par2 = {"n": 8, "m": 10, "dtype": "float64"} # underdetermined, float64 +par3 = {"n": 8, "m": 10, "dtype": "float32"} # underdetermined, float32 def test_ProximalGradient_unknown_acceleration(): @@ -63,7 +65,7 @@ def test_ADMM_noinitial(): def test_ADMML2_noinitial(): """Check that an error is raised if no initial value - is provided to PrimalDual solver + is provided to ADMML2 solver """ with pytest.raises(ValueError, match="Both x0 or"): # Both None @@ -117,7 +119,7 @@ def test_LinearizedADMM_noinitial(): ) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_GPG_weights(par): """Check GPG raises error if weight is not summing to 1""" with pytest.raises(ValueError, match="must be an array of size"): @@ -148,7 +150,7 @@ def test_GPG_weights(par): ) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_ProximalPoint(par): """Check solution of ProximalPoint for quadratic function equals the solution of the associated system of linear equations @@ -171,7 +173,7 @@ def test_ProximalPoint(par): assert_array_almost_equal(xpp, x, decimal=2) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_PG_ISTA(par): """Check equivalency of ProximalGradient and ISTA/FISTA (PyLops)""" np.random.seed(0) @@ -227,7 +229,7 @@ def test_PG_ISTA(par): assert_array_almost_equal(xpg, xista, decimal=2) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_PG_GPG(par): """Check equivalency of ProximalGradient and GeneralizedProximalGradient when using a single regularization term @@ -275,9 +277,55 @@ def test_PG_GPG(par): assert_array_almost_equal(xpg, xgpg, decimal=2) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) +def test_PG_TwiST(par): + """Check that PG/TwiST can be used to solve a sparsity regularized objective function + (note that despite the trajectory will be different, they should converge to the + same solution) + """ + np.random.seed(0) + n, m = par["n"], par["m"] + + # Define sparse model + x = np.zeros(m) + x[2], x[4] = 1, 0.5 + + # Random mixing matrix + R = np.random.normal(0.0, 1.0, (n, m)) + Rop = MatrixMult(R) + + y = Rop @ x + + # Step size + L = (Rop.H * Rop).eigs(1).real + tau = 0.99 / L + + # PG + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l1 = L1(sigma=5e-1) + xpg = ProximalGradient( + l2, l1, x0=np.zeros(m), tau=tau, niter=100, acceleration="fista" + ) + + # TwiST + l1 = L1(sigma=5e-1) + eigs = np.linalg.eig(R.T @ R)[0] + eigs = (np.abs(eigs[0]), max(1e-1, np.abs(eigs[-1]))) + xtwist = TwIST( + l1, + Rop, + y, + x0=np.zeros(m), + eigs=eigs, + niter=100, + ) + + assert_array_almost_equal(xpg, xtwist, decimal=2) + + +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_HQS_ADMM_L2(par): - """Check that HQS/ADMM can be used to solved a pure L2-based objective function + """Check that HQS/ADMM can be used to solve a pure L2-based objective function (and compare with LSQR - note that despite the trajectory will be different, they should converge to the same solution) """ @@ -326,7 +374,85 @@ def test_HQS_ADMM_L2(par): assert_array_almost_equal(xl2, xadmm, decimal=2) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) +def test_ADMM_ADMML2(par): + """Check equivalency of ADMM and ADMML2 + when the f function is a L2 term + """ + np.random.seed(0) + n, m = par["n"], par["m"] + + # Define sparse model + x = np.random.normal(0.0, 1.0, m).astype(par["dtype"]) + + # Random mixing matrix + R = np.random.normal(0.0, 1.0, (n, m)).astype(par["dtype"]) + Rop = MatrixMult(R, dtype=par["dtype"]) + + y = Rop @ x + + # Step size + Aop = Identity(m) + L = 1.0 # Lipshitz constant of Aop + tau = 0.99 / L + eps = 1e-1 + + # ADMM + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l2reg = L2(sigma=eps) + xadmm = ADMM(l2, l2reg, x0=np.zeros(m), tau=tau, niter=100)[0] + + # ADMML2 + l2reg = L2(sigma=eps) + xadmml2 = ADMML2( + l2reg, Rop, y, Aop, x0=np.zeros(m), tau=tau, niter=100, iter_lim=10 + )[0] + + assert_array_almost_equal(xadmm, xadmml2, decimal=2) + + +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) +def test_ADMM_LinearizedADMM(par): + """Check equivalency of ADMM and LinearizedADMM + when the f function is a L2 term + """ + np.random.seed(0) + n, m = par["n"], par["m"] + + # Define sparse model + x = np.random.normal(0.0, 1.0, m).astype(par["dtype"]) + + # Random mixing matrix + R = np.random.normal(0.0, 1.0, (n, m)).astype(par["dtype"]) + Rop = MatrixMult(R, dtype=par["dtype"]) + + y = Rop @ x + + # Step size + Aop = Identity(m) + L = 1.0 # Lipshitz constant of Aop + tau = 0.99 / L + mu = 0.99 / L # optimal mu<=tau/maxeig(Dop^H Dop) + + eps = 1e-1 + + # ADMM + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l2reg = L2(sigma=eps) + xadmm = ADMM(l2, l2reg, x0=np.zeros(m), tau=tau, niter=100)[0] + + # LinearizedADMM + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l2reg = L2(sigma=eps) + Aop = Identity(m) + xladmm = LinearizedADMM(l2, l2reg, Aop, x0=np.zeros(m), tau=tau, mu=mu, niter=100)[ + 0 + ] + + assert_array_almost_equal(xadmm, xladmm, decimal=2) + + +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_ADMM_DRS(par): """Check equivalency of ADMM and DouglasRachfordSplitting when using a single regularization term @@ -371,7 +497,7 @@ def test_ADMM_DRS(par): assert_array_almost_equal(xadmm, xdrs_f, decimal=2) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_PPXA_with_ADMM(par: dict[str, Any]) -> None: """Check equivalency of PPXA and ADMM when using a single regularization term @@ -413,7 +539,7 @@ def test_PPXA_with_ADMM(par: dict[str, Any]) -> None: assert_array_almost_equal(xppxa, xadmm, decimal=2) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_PPXA_with_GPG(par: dict[str, Any]) -> None: """Check equivalency of PPXA and GeneralizedProximalGradient""" np.random.seed(0) @@ -461,7 +587,7 @@ def test_PPXA_with_GPG(par: dict[str, Any]) -> None: assert_array_almost_equal(xppxa, xgpg, decimal=2) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_ConsensusADMM_with_ADMM(par: dict[str, Any]) -> None: """Check equivalency of ConsensusADMM and ADMM when two proximable functions @@ -503,7 +629,7 @@ def test_ConsensusADMM_with_ADMM(par: dict[str, Any]) -> None: assert_array_almost_equal(xcadmm, xadmm, decimal=2) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_ConsensusADMM_with_ADMM_for_Lasso(par: dict[str, Any]) -> None: """Check equivalency of ConsensusADMM and ADMM when more than two proximable functions for lasso @@ -565,7 +691,7 @@ def test_ConsensusADMM_with_ADMM_for_Lasso(par: dict[str, Any]) -> None: assert_array_almost_equal(xcadmm, xadmm, decimal=2) -@pytest.mark.parametrize("par", [(par1), (par2)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_ConsensusADMM_with_GPG(par: dict[str, Any]) -> None: """Check equivalency of ConsensusADMM and GeneralizedProximalGradient""" From c4709a14350b2e9017f3ee3449827078c7b180c6 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Mon, 8 Jun 2026 21:24:39 +0100 Subject: [PATCH 14/24] minor: fix mypy issue --- pyproximal/optimization/basesolver.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/pyproximal/optimization/basesolver.py b/pyproximal/optimization/basesolver.py index 1d1b73d..97739a8 100644 --- a/pyproximal/optimization/basesolver.py +++ b/pyproximal/optimization/basesolver.py @@ -11,8 +11,6 @@ from pylops.optimization.callback import Callbacks from pylops.utils.typing import NDArray -from pyproximal.utils.typing import Tmemunit - if TYPE_CHECKING: from pyproximal.ProxOperator import ProxOperator @@ -87,14 +85,11 @@ def _print_finalize(self, *args: Any, nbar: int = 80, **kwargs: Any) -> None: ) print("-" * nbar + "\n") - def memory_usage( - self, - show: bool = False, - unit: Tmemunit = "B", - ) -> float: + def memory_usage(self) -> None: """Compute memory usage of the solver - Currently not implemented in PyProximal and therefore just overwritten here. + Currently not implemented in PyProximal's solvers + and therefore just overwritten here. """ pass @@ -165,7 +160,7 @@ def finalize(self, nbar: int = 60, show: bool = False) -> None: """ self.tend = time.time() self.telapsed = self.tend - self.tstart - self.cost = np.array(self.cost) + self.cost: NDArray = np.array(self.cost) if show: self._print_finalize(nbar=nbar) From 810a3f859dfcd3887b5b32a5f403e08991aa90f4 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Mon, 8 Jun 2026 21:25:35 +0100 Subject: [PATCH 15/24] feat: added more class solvers --- pyproximal/optimization/cls_primal.py | 1228 +++++++++++++++++++++++-- 1 file changed, 1165 insertions(+), 63 deletions(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index 9019eda..2860059 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -1,20 +1,30 @@ __all__ = [ "ProximalPoint", + "ProximalGradient", + "AndersonProximalGradient", + "GeneralizedProximalGradient", + "HQS", + "ADMM", + "ADMML2", + "LinearizedADMM", + "TwIST", ] from collections.abc import Sequence -from typing import TYPE_CHECKING, Optional, cast +from math import sqrt +from typing import TYPE_CHECKING, Any, Optional, cast import numpy as np import pylops from pylops.optimization.callback import Callbacks +from pylops.optimization.leastsquares import regularized_inversion from pylops.utils.backend import get_array_module, to_numpy from pylops.utils.typing import NDArray from pyproximal.optimization.basesolver import Solver +from pyproximal.proximal import L2 from pyproximal.ProxOperator import ProxOperator from pyproximal.utils.bilinear import BilinearOperator -from pyproximal.utils.typing import Tmemunit if TYPE_CHECKING: from pylops.linearoperator import LinearOperator @@ -167,13 +177,6 @@ def _print_step(self, x: NDArray) -> None: msg = f"{self.iiter:6g} " + strx + f"{self.pf:11.4e}" print(msg) - def memory_usage( - self, - show: bool = False, - unit: Tmemunit = "B", - ) -> None: - pass - def setup( # type: ignore[override] self, prox: "ProxOperator", @@ -454,13 +457,6 @@ def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: ) print(msg) - def memory_usage( - self, - show: bool = False, - unit: Tmemunit = "B", - ) -> None: - pass - def setup( # type: ignore[override] self, proxf: ProxOperator, @@ -897,13 +893,6 @@ def _print_step(self, x: NDArray, pg: float | None) -> None: ) print(msg) - def memory_usage( - self, - show: bool = False, - unit: Tmemunit = "B", - ) -> None: - pass - def setup( # type: ignore[override] self, proxf: ProxOperator, @@ -1013,10 +1002,10 @@ def step( ---------- x : :obj:`numpy.ndarray` Current model vector to be updated by a step of the - proximal gradient algorithm + Andersonproximal gradient algorithm y : :obj:`numpy.ndarray` Additional model vector to be updated by a step of the - proximal gradient algorithm + Anderson proximal gradient algorithm show : :obj:`bool`, optional Display iteration log @@ -1109,7 +1098,7 @@ def run( the Anderson proximal gradient algorithm y : :obj:`numpy.ndarray` Additional model vector to be updated by multiple steps of - the proximal gradient algorithm + the Anderson proximal gradient algorithm niter : :obj:`int`, optional Number of iterations. Can be set to ``None`` if already provided in the setup call @@ -1306,13 +1295,6 @@ def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: ) print(msg) - def memory_usage( - self, - show: bool = False, - unit: Tmemunit = "B", - ) -> None: - pass - def setup( # type: ignore[override] self, proxfs: list[ProxOperator], @@ -1427,10 +1409,10 @@ def step( ---------- x : :obj:`numpy.ndarray` Current model vector to be updated by a step of the - proximal gradient algorithm + generalized proximal gradient algorithm y : :obj:`numpy.ndarray` Additional model vector to be updated by a step of the - proximal gradient algorithm + generalized proximal gradient algorithm show : :obj:`bool`, optional Display iteration log @@ -1506,10 +1488,10 @@ def run( ---------- x : :obj:`numpy.ndarray` Current model vector to be updated by multiple steps of - the proximal gradient algorithm + the generalized proximal gradient algorithm y : :obj:`numpy.ndarray` Additional model vector to be updated by multiple steps of - the proximal gradient algorithm + the generalized proximal gradient algorithm niter : :obj:`int`, optional Number of iterations. Can be set to ``None`` if already provided in the setup call @@ -1699,7 +1681,7 @@ def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: if self.tol is None: pf, pg = self.proxf(x), self.proxg(x) self.pfg = pf + pg - x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) + x0 = to_numpy(x[0]) strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " msg = ( f"{self.iiter:6g} " @@ -1708,13 +1690,6 @@ def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: ) print(msg) - def memory_usage( - self, - show: bool = False, - unit: Tmemunit = "B", - ) -> None: - pass - def setup( # type: ignore[override] self, proxf: ProxOperator, @@ -1812,10 +1787,10 @@ def step( ---------- x : :obj:`numpy.ndarray` Current model vector to be updated by a step of the - proximal gradient algorithm + HQS algorithm z : :obj:`numpy.ndarray` Additional model vector to be updated by a step of the - proximal gradient algorithm + HQS algorithm show : :obj:`bool`, optional Display iteration log @@ -1868,10 +1843,10 @@ def run( ---------- x : :obj:`numpy.ndarray` Current model vector to be updated by multiple steps of - the proximal gradient algorithm + the HQS algorithm z : :obj:`numpy.ndarray` Additional model vector to be updated by multiple steps of - the proximal gradient algorithm + the HQS algorithm niter : :obj:`int`, optional Number of iterations. Can be set to ``None`` if already provided in the setup call @@ -2073,7 +2048,7 @@ def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: if self.tol is None: pf, pg = self.proxf(x), self.proxg(x) self.pfg = pf + pg - x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) + x0 = to_numpy(x[0]) strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " msg = ( f"{self.iiter:6g} " @@ -2082,19 +2057,12 @@ def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: ) print(msg) - def memory_usage( - self, - show: bool = False, - unit: Tmemunit = "B", - ) -> None: - pass - def setup( # type: ignore[override] self, proxf: ProxOperator, proxg: ProxOperator, x0: NDArray, - tau: float | NDArray, + tau: float, z0: NDArray | None = None, niter: int = 10, gfirst: bool = False, @@ -2178,10 +2146,10 @@ def step( ---------- x : :obj:`numpy.ndarray` Current model vector to be updated by a step of the - proximal gradient algorithm + ADDM algorithm z : :obj:`numpy.ndarray` Additional model vector to be updated by a step of the - proximal gradient algorithm + ADDM algorithm show : :obj:`bool`, optional Display iteration log @@ -2235,10 +2203,10 @@ def run( ---------- x : :obj:`numpy.ndarray` Current model vector to be updated by multiple steps of - the proximal gradient algorithm + the ADDM algorithm z : :obj:`numpy.ndarray` Additional model vector to be updated by multiple steps of - the proximal gradient algorithm + the ADDM algorithm niter : :obj:`int`, optional Number of iterations. Can be set to ``None`` if already provided in the setup call @@ -2288,7 +2256,7 @@ def solve( # type: ignore[override] proxf: ProxOperator, proxg: ProxOperator, x0: NDArray, - tau: float | NDArray, + tau: float, z0: NDArray | None = None, niter: int = 10, gfirst: bool = False, @@ -2351,9 +2319,1143 @@ def solve( # type: ignore[override] niter=niter, gfirst=gfirst, tol=tol, + callbackz=callbackz, show=show, ) x, z = self.run(x, z, niter, show=show, itershow=itershow) self.finalize(65, show) return x, z, self.iiter, self.cost + + +class ADMML2(Solver): + r"""Alternating Direction Method of Multipliers for L2 misfit term + + Solves the following minimization problem using Alternating Direction + Method of Multipliers: + + .. math:: + + \mathbf{x},\mathbf{z} = \argmin_{\mathbf{x},\mathbf{z}} + \frac{1}{2}||\mathbf{Op}\mathbf{x} - \mathbf{b}||_2^2 + g(\mathbf{z}) \\ + s.t. \; \mathbf{Ax}=\mathbf{z} + + where :math:`g(\mathbf{z})` is any convex function that has a known proximal operator. + + See Also + -------- + ADMM: ADMM + LinearizedADMM: Linearized ADMM + + Notes + ----- + The ADMM algorithm with L2 misfit term can be expressed by the following recursion: + + .. math:: + + \mathbf{x}^{k+1} = \argmin_{\mathbf{x}} \frac{1}{2}||\mathbf{Op}\mathbf{x} + - \mathbf{b}||_2^2 + \frac{1}{2\tau} ||\mathbf{Ax} - \mathbf{z}^k + \mathbf{u}^k||_2^2\\ + \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{Ax}^{k+1} + \mathbf{u}^{k})\\ + \mathbf{u}^{k+1} = \mathbf{u}^{k} + \mathbf{Ax}^{k+1} - \mathbf{z}^{k+1} + + .. [1] S. Boyd, N. Parikh, E. Chu, B. Peleato, and J. Eckstein. 2011. + Distributed optimization and statistical learning via the alternating + direction method of multipliers. Foundations and Trends in Machine + Learning, 3 (1), 1-122. https://doi.org/10.1561/2200000016. + + """ + + def _print_setup(self, xcomplex: bool = False) -> None: + self._print_solver(nbar=65) + + strpar = ( + f"Proximal operator (g): {type(self.proxg).__name__}\n" + f"Linear operator (Op): {type(self.Op).__name__}\n" + f"Linear operator (A): {type(self.A).__name__}\n" + ) + strpar1 = f"tau = {self.tau:6e}\tniter = {self.niter}\ttol = {self.tol}" + print(strpar) + print(strpar1) + + print("-" * 65 + "\n") + if not xcomplex: + head1 = " Itn x[0] f g J=f+g" + else: + head1 = ( + " Itn x[0] f g J=f+g" + ) + print(head1) + + def _print_step( + self, x: NDArray, Ax: NDArray, pf: float | None, pg: float | None + ) -> None: + if self.tol is None: + pf, pg = ( + 0.5 * self.ncp.linalg.norm(self.Op @ x - self.b) ** 2, + self.proxg(Ax), + ) + self.pfg = pf + pg + x0 = to_numpy(x[0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = ( + f"{self.iiter:6g} " + + strx + + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e}" + ) + print(msg) + + def setup( # type: ignore[override] + self, + proxg: ProxOperator, + Op: "LinearOperator", + b: NDArray, + A: "LinearOperator", + x0: NDArray, + tau: float, + z0: NDArray | None = None, + niter: int = 10, + gfirst: bool = False, + tol: float | None = None, + callbackz: bool = False, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Setup solver + + Parameters + ---------- + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + Op : :obj:`pylops.LinearOperator` + Linear operator of data misfit term + b : :obj:`numpy.ndarray` + Data + A : :obj:`pylops.LinearOperator` + Linear operator of regularization term + x0 : :obj:`numpy.ndarray` + Initial vector + tau : :obj:`float` + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/\lambda_{max}(\mathbf{A}^H\mathbf{A})]`. + z0 : :obj:`numpy.ndarray` + Initial auxiliary vector. If ``None``, initialized to ``A @ x0``. + niter : :obj:`int`, optional + Number of iterations of iterative scheme + gfirst : :obj:`bool`, optional + Apply Proximal of operator ``g`` first (``True``) or Proximal of + operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + callbackz : :obj:`bool`, optional + Modify callback signature to (``callback(x, z)``) when ``callbackz=True`` + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + z : :obj:`numpy.ndarray` + Initial guess for the auxiliary variable + + Raises + ------ + ValueError + If both ``x0`` and ``z0`` are set to ``None`` or ``x0`` is set to None + + """ + self.proxg = proxg + self.Op = Op + self.b = b + self.A = A + self.tau = tau + self.niter = niter + self.gfirst = gfirst + self.tol = tol + self.callbackz = callbackz + + self.ncp = get_array_module(x0) + + # initialize solver + x, z = _x0z0_init(x0, z0, A, Opname="A") + self.u = self.ncp.zeros_like(x) + + # other parameters + self.sqrttau = 1.0 / sqrt(self.tau) + + # create variables to track the objective function and iterations + self.pfg, self.pfgold = np.inf, np.inf + self.cost: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(np.iscomplexobj(x0)) + return x, z + + def step( + self, + x: NDArray, + z: NDArray, + show: bool = False, + **kwargs_solver: dict[str, Any], + ) -> tuple[NDArray, NDArray]: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + ADMML2 algorithm + z : :obj:`numpy.ndarray` + Additional model vector to be updated by a step of the + ADMML2 algorithm + show : :obj:`bool`, optional + Display iteration log + **kwargs_solver + Arbitrary keyword arguments for :py:func:`scipy.sparse.linalg.lsqr` used + to solve the x-update + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + z : :obj:`numpy.ndarray` + Updated additional model vector + + """ + if self.gfirst: + Ax = self.A @ x + z = self.proxg.prox(Ax + self.u, self.tau) + + # solve augumented system + x = regularized_inversion( + self.Op, + self.b, + [ + self.A, + ], + x0=x, + dataregs=[ + z - self.u, + ], + epsRs=[ + self.sqrttau, + ], + **kwargs_solver, + )[0] + else: + # solve augumented system + x = regularized_inversion( + self.Op, + self.b, + [ + self.A, + ], + x0=x, + dataregs=[ + z - self.u, + ], + epsRs=[ + self.sqrttau, + ], + **kwargs_solver, + )[0] + Ax = self.A @ x + z = self.proxg.prox(Ax + self.u, self.tau) + self.u = self.u + Ax - z + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfgold = self.pfg + pf = 0.5 * self.ncp.linalg.norm(self.Op @ x - self.b) ** 2 + pg = self.proxg(Ax) + self.pfg = pf + pg + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pf, pg = 0.0, 0.0 + + self.iiter += 1 + if show: + self._print_step(x, Ax, pf, pg) + if self.tol is not None or show: + self.cost.append(float(self.pfg)) + return x, z + + def run( + self, + x: NDArray, + z: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + **kwargs_solver: dict[str, Any], + ) -> tuple[NDArray, NDArray]: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the ADDML2 algorithm + z : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the ADDML2 algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + **kwargs_solver + Arbitrary keyword arguments for :py:func:`scipy.sparse.linalg.lsqr` used + to solve the x-update + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + z : :obj:`numpy.ndarray` + Additional estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, z = self.step(x, z, showstep, **kwargs_solver) + if self.callbackz: + self.callback(x, z) + else: + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, z + + def solve( # type: ignore[override] + self, + proxg: ProxOperator, + Op: "LinearOperator", + b: NDArray, + A: "LinearOperator", + x0: NDArray, + tau: float, + z0: NDArray | None = None, + niter: int = 10, + gfirst: bool = False, + tol: float | None = None, + callbackz: bool = False, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + **kwargs_solver: dict[str, Any], + ) -> tuple[NDArray, NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + Op : :obj:`pylops.LinearOperator` + Linear operator of data misfit term + b : :obj:`numpy.ndarray` + Data + A : :obj:`pylops.LinearOperator` + Linear operator of regularization term + x0 : :obj:`numpy.ndarray` + Initial vector (not required when ``gfirst=False``, can pass ``None``) + tau : :obj:`float` + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. + z0 : :obj:`numpy.ndarray`, optional + Initial z vector (not required when ``gfirst=True``) + niter : :obj:`int` + Number of iterations of iterative scheme + gfirst : :obj:`bool`, optional + Apply Proximal of operator ``g`` first (``True``) or Proximal of + operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping + criterion). If ``tol=None``, run until ``niter`` is reached + callbackz : :obj:`bool`, optional + Modify callback signature to (``callback(x, z)``) when ``callbackz=True`` + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + **kwargs_solver + Arbitrary keyword arguments for :py:func:`scipy.sparse.linalg.lsqr` used + to solve the x-update + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + z : :obj:`numpy.ndarray` + Additional estimated model + iiter : :obj:`int` + Number of executed iterations + cost : :obj:`list` + History of the objective function + + """ + x, z = self.setup( + proxg=proxg, + Op=Op, + b=b, + A=A, + x0=x0, + tau=tau, + z0=z0, + niter=niter, + gfirst=gfirst, + tol=tol, + callbackz=callbackz, + show=show, + ) + + x, z = self.run( + x, + z, + niter, + show=show, + itershow=itershow, + **kwargs_solver, + ) + self.finalize(65, show) + return x, z, self.iiter, self.cost + + +class LinearizedADMM(Solver): + r"""Linearized Alternating Direction Method of Multipliers + + Solves the following minimization problem using Linearized Alternating + Direction Method of Multipliers: + + .. math:: + + \mathbf{x} = \argmin_\mathbf{x} f(\mathbf{x}) + g(\mathbf{A}\mathbf{x}) + + where :math:`f(\mathbf{x})` and :math:`g(\mathbf{x})` are any convex + function that has a known proximal operator and :math:`\mathbf{A}` is a + linear operator. + + See Also + -------- + ADMM: ADMM + ADMML2: ADMM with L2 misfit function + + Notes + ----- + The Linearized-ADMM algorithm can be expressed by the following recursion [1]_: + + .. math:: + + \mathbf{x}^{k+1} = \prox_{\mu f}(\mathbf{x}^{k} - \frac{\mu}{\tau} + \mathbf{A}^H(\mathbf{A} \mathbf{x}^k - \mathbf{z}^k + \mathbf{u}^k))\\ + \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{A} \mathbf{x}^{k+1} + + \mathbf{u}^k)\\ + \mathbf{u}^{k+1} = \mathbf{u}^{k} + \mathbf{A}\mathbf{x}^{k+1} - + \mathbf{z}^{k+1} + + .. [1] N., Parikh, "Proximal Algorithms", Foundations and Trends + in Optimization. 2013. + + """ + + def _print_setup(self, xcomplex: bool = False) -> None: + self._print_solver(nbar=65) + + strpar = ( + f"Proximal operator (f): {type(self.proxf).__name__}\n" + f"Proximal operator (g): {type(self.proxg).__name__}\n" + f"Linear operator (A): {type(self.A).__name__}\n" + ) + strpar1 = f"tau = {self.tau:6e}\tmu = {self.mu:6e}" + strpar2 = f"niter = {self.niter}\ttol = {self.tol}" + print(strpar) + print(strpar1) + print(strpar2) + + print("-" * 65 + "\n") + if not xcomplex: + head1 = " Itn x[0] f g J=f+g" + else: + head1 = ( + " Itn x[0] f g J=f+g" + ) + print(head1) + + def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: + if self.tol is None: + pf, pg = self.proxf(x), self.proxg(self.Ax) + self.pfg = pf + pg + x0 = to_numpy(x[0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = ( + f"{self.iiter:6g} " + + strx + + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e}" + ) + print(msg) + + def setup( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + A: "LinearOperator", + x0: NDArray, + tau: float, + mu: float, + z0: NDArray | None = None, + niter: int = 10, + tol: float | None = None, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Setup solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + A : :obj:`pylops.LinearOperator` + Linear operator + x0 : :obj:`numpy.ndarray` + Initial vector + tau : :obj:`float`, optional + Positive scalar weight, which should satisfy the following + condition to guarantee convergence: :math:`\mu \in (0, + \tau/\lambda_{max}(\mathbf{A}^H\mathbf{A})]`. + mu : :obj:`float`, optional + Second positive scalar weight, which should satisfy the following + condition to guarantees convergence: :math:`\mu \in (0, + \tau/\lambda_{max}(\mathbf{A}^H\mathbf{A})]`. + z0 : :obj:`numpy.ndarray` + Initial auxiliary vector. If ``None``, initialized to ``A @ x0``. + niter : :obj:`int`, optional + Number of iterations of iterative scheme + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + z : :obj:`numpy.ndarray` + Initial guess for the auxiliary variable + + Raises + ------ + ValueError + If both ``x0`` and ``z0`` are set to ``None`` or ``x0`` is set to None + + """ + self.proxf = proxf + self.proxg = proxg + self.A = A + self.tau = tau + self.mu = mu + + self.niter = niter + self.tol = tol + + self.ncp = get_array_module(x0) + + # initialize solver + x, z = _x0z0_init(x0, z0, A, Opname="A") + self.Ax = A.matvec(x) if z0 is None else z + self.u = self.ncp.zeros_like(x) + + # create variables to track the objective function and iterations + self.pfg, self.pfgold = np.inf, np.inf + self.cost: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(np.iscomplexobj(x0)) + return x, z + + def step( + self, + x: NDArray, + z: NDArray, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + LinearizedADMM algorithm + z : :obj:`numpy.ndarray` + Additional model vector to be updated by a step of the + LinearizedADMM algorithm + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + z : :obj:`numpy.ndarray` + Updated additional model vector + + """ + x = self.proxf.prox( + x - self.mu / self.tau * self.A.rmatvec(self.Ax - z + self.u), self.mu + ) + + self.Ax = self.A.matvec(x) + z = self.proxg.prox(self.Ax + self.u, self.tau) + self.u = self.u + self.Ax - z + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfgold = self.pfg + pf = self.proxf(x) + pg = self.proxg(self.Ax) + self.pfg = pf + pg + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pf, pg = 0.0, 0.0 + + self.iiter += 1 + if show: + self._print_step(x, pf, pg) + if self.tol is not None or show: + self.cost.append(float(self.pfg)) + return x, z + + def run( + self, + x: NDArray, + z: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray]: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the LinearizedADDM algorithm + z : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the LinearizedADDM algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + z : :obj:`numpy.ndarray` + Additional estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, z = self.step(x, z, showstep) + + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, z + + def solve( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + A: "LinearOperator", + x0: NDArray, + tau: float, + mu: float, + z0: NDArray | None = None, + niter: int = 10, + tol: float | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + A : :obj:`pylops.LinearOperator` + Linear operator + x0 : :obj:`numpy.ndarray` + Initial vector + tau : :obj:`float`, optional + Positive scalar weight, which should satisfy the following + condition to guarantee convergence: :math:`\mu \in (0, + \tau/\lambda_{max}(\mathbf{A}^H\mathbf{A})]`. + mu : :obj:`float`, optional + Second positive scalar weight, which should satisfy the following + condition to guarantees convergence: :math:`\mu \in (0, + \tau/\lambda_{max}(\mathbf{A}^H\mathbf{A})]`. + z0 : :obj:`numpy.ndarray` + Initial auxiliary vector. If ``None``, initialized to ``A @ x0``. + niter : :obj:`int`, optional + Number of iterations of iterative scheme + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + z : :obj:`numpy.ndarray` + Additional estimated model + iiter : :obj:`int` + Number of executed iterations + cost : :obj:`list` + History of the objective function + + """ + x, z = self.setup( + proxf=proxf, + proxg=proxg, + A=A, + x0=x0, + tau=tau, + mu=mu, + z0=z0, + niter=niter, + tol=tol, + show=show, + ) + + x, z = self.run( + x, + z, + niter, + show=show, + itershow=itershow, + ) + self.finalize(65, show) + return x, z, self.iiter, self.cost + + +class TwIST(Solver): + r"""Two-step Iterative Shrinkage/Threshold + + Solves the following minimization problem using Two-step Iterative + Shrinkage/Threshold: + + .. math:: + + \mathbf{x} = \argmin_\mathbf{x} \frac{1}{2} + ||\mathbf{b} - \mathbf{Ax}||_2^2 + g(\mathbf{x}) + + where :math:`\mathbf{A}` is a linear operator and :math:`g(\mathbf{x})` + is any convex function that has a known proximal operator. + + Notes + ----- + The TwIST algorithm can be expressed by the following recursion: + + .. math:: + + \mathbf{x}^{k+1} = (1-\alpha) \mathbf{x}^{k-1} + + (\alpha-\beta) \mathbf{x}^k + + \beta \prox_{g} (\mathbf{x}^k + \mathbf{A}^H + (\mathbf{b} - \mathbf{A}\mathbf{x}^k)). + + where :math:`\mathbf{x}^{1} = \prox_{g} (\mathbf{x}^0 + \mathbf{A}^T + (\mathbf{b} - \mathbf{A}\mathbf{x}^0))`. + + The optimal weighting parameters :math:`\alpha` and :math:`\beta` are + linked to the smallest and largest eigenvalues of + :math:`\mathbf{A}^H\mathbf{A}` as follows: + + .. math:: + + \alpha = 1 + \rho^2 \\ + \beta =\frac{2 \alpha}{\Lambda_{max} + \lambda_{min}} + + where :math:`\rho=\frac{1-\sqrt{k}}{1+\sqrt{k}}` with + :math:`k=\frac{\lambda_{min}}{\Lambda_{max}}` and + :math:`\Lambda_{max}=max(1, \lambda_{max})`. + + Experimentally, it has been observed that TwIST is robust to the + choice of such parameters. Finally, note that in the case of + :math:`\alpha=1` and :math:`\beta=1`, TwIST is identical to IST. + + """ + + def _print_setup(self, xcomplex: bool = False) -> None: + self._print_solver(nbar=65) + + strpar = ( + f"Proximal operator (f): {type(self.proxf).__name__}\n" + f"Linear operator (A): {type(self.A).__name__}\n" + ) + strpar1 = f"alpha = {self.alpha:6e}\tbeta = {self.beta:6e}" + strpar2 = f"niter = {self.niter}\ttol = {self.tol}" + print(strpar) + print(strpar1) + print(strpar2) + + print("-" * 65 + "\n") + if not xcomplex: + head1 = " Itn x[0] f g J=f+g" + else: + head1 = ( + " Itn x[0] f g J=f+g" + ) + print(head1) + + def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: + if self.tol is None: + pf, pg = self.proxf(x), self.proxg(x) + self.pfg = pf + pg + x0 = to_numpy(x[0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = ( + f"{self.iiter:6g} " + + strx + + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e}" + ) + print(msg) + + def setup( # type: ignore[override] + self, + proxg: ProxOperator, + A: "LinearOperator", + b: NDArray, + x0: NDArray, + alpha: float | None = None, + beta: float | None = None, + eigs: tuple[float, float] | None = None, + niter: int = 10, + tol: float | None = None, + show: bool = False, + ) -> NDArray: + r"""Setup solver + + Parameters + ---------- + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + A : :obj:`pylops.LinearOperator` + Linear operator + b : :obj:`numpy.ndarray` + Data + x0 : :obj:`numpy.ndarray` + Initial vector + alpha : :obj:`float`, optional + Positive scalar weight (if ``None``, estimated based on the + eigenvalues of :math:`\mathbf{A}`, see Notes for details) + beta : :obj:`float`, optional + Positive scalar weight (if ``None``, estimated based on the + eigenvalues of :math:`\mathbf{A}`, see Notes for details) + eigs : :obj:`tuple`, optional + Largest and smallest eigenvalues of :math:`\mathbf{A}^H \mathbf{A}`. + If passed, computes `alpha` and `beta` based on them. + niter : :obj:`int`, optional + Number of iterations of iterative scheme + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + + """ + self.proxf = L2(Op=A, b=b) + self.proxg = proxg + self.A = A + self.niter = niter + self.tol = tol + + self.ncp = get_array_module(x0) + + # find alpha and beta based on the eigenvalues of A if not provided + if alpha is None or beta is None: + if eigs is None: + emin = A.eigs(neigs=1, which="SM") + emax = max([1, A.eigs(neigs=1, which="LM")]) + else: + emax, emin = eigs + k = emin / emax + rho = (1 - sqrt(k)) / (1 + sqrt(k)) + self.alpha = 1 + rho**2 + self.beta = 2 * self.alpha / (emax + emin) + else: + self.alpha, self.beta = alpha, beta + + # compute proximal of g on initial guess (x_1) + x = self.proxg.prox(x0 - self.proxf.grad(x0), 1.0) + + # create variables to track the objective function and iterations + self.pfg, self.pfgold = np.inf, np.inf + self.cost: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(np.iscomplexobj(x0)) + return x + + def step( + self, + x: NDArray, + xold: NDArray, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + TwisT algorithm + xold : :obj:`numpy.ndarray` + Previous model vector + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + xold : :obj:`numpy.ndarray` + Previous model vector + + """ + # compute new x + xnew = ( + (1 - self.alpha) * xold + + (self.alpha - self.beta) * x + + self.beta * self.proxg.prox(x - self.proxf.grad(x), 1.0) + ) + # save current x as old (x_i -> x_i-1) + xold = x.copy() + # save new x as current (x_i+1 -> x_i) + x = xnew.copy() + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfgold = self.pfg + pf = self.proxf(x) + pg = self.proxg(self.Ax) + self.pfg = pf + pg + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pf, pg = 0.0, 0.0 + + self.iiter += 1 + if show: + self._print_step(x, pf, pg) + if self.tol is not None or show: + self.cost.append(float(self.pfg)) + return x, xold + + def run( + self, + x: NDArray, + xold: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> NDArray: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the TwisT algorithm + xold : :obj:`numpy.ndarray` + Previous estimated model + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, xold = self.step(x, xold, showstep) + + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x + + def solve( # type: ignore[override] + self, + proxg: ProxOperator, + A: "LinearOperator", + b: NDArray, + x0: NDArray, + alpha: float | None = None, + beta: float | None = None, + eigs: tuple[float, float] | None = None, + niter: int = 10, + tol: float | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + A : :obj:`pylops.LinearOperator` + Linear operator + b : :obj:`numpy.ndarray` + Data + x0 : :obj:`numpy.ndarray` + Initial vector + alpha : :obj:`float`, optional + Positive scalar weight (if ``None``, estimated based on the + eigenvalues of :math:`\mathbf{A}`, see Notes for details) + beta : :obj:`float`, optional + Positive scalar weight (if ``None``, estimated based on the + eigenvalues of :math:`\mathbf{A}`, see Notes for details) + eigs : :obj:`tuple`, optional + Largest and smallest eigenvalues of :math:`\mathbf{A}^H \mathbf{A}`. + If passed, computes `alpha` and `beta` based on them. + niter : :obj:`int`, optional + Number of iterations of iterative scheme + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + iiter : :obj:`int` + Number of executed iterations + cost : :obj:`list` + History of the objective function + + """ + x = self.setup( + proxg=proxg, + A=A, + b=b, + x0=x0, + alpha=alpha, + beta=beta, + eigs=eigs, + niter=niter, + tol=tol, + show=show, + ) + + x = self.run( + x, + x0, + niter, + show=show, + itershow=itershow, + ) + self.finalize(65, show) + return x, self.iiter, self.cost From 4c6940a104a54cbe3b944173bbe83ed5b24376b4 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Wed, 17 Jun 2026 23:09:23 +0100 Subject: [PATCH 16/24] feat: added class-based DouglasRachfordSplitting and PPXA --- pyproximal/optimization/cls_primal.py | 730 +++++++++++++++++++++++++- 1 file changed, 723 insertions(+), 7 deletions(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index 2860059..dc7f1c4 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -8,6 +8,7 @@ "ADMML2", "LinearizedADMM", "TwIST", + "DouglasRachfordSplitting", ] from collections.abc import Sequence @@ -197,8 +198,8 @@ def setup( # type: ignore[override] tau : :obj:`float` Positive scalar weight niter : :obj:`int`, optional - Number of iterations (default to ``None`` in case a user wants to - manually step over the solver) + Number of iterations of iterative scheme (default to ``None`` + in case a user wants to manually step over the solver) tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -428,10 +429,14 @@ def _print_setup(self, epsg_print: str, xcomplex: bool = False) -> None: f"Proximal operator (f): {type(self.proxf).__name__}\n" f"Proximal operator (g): {type(self.proxg).__name__}\n" ) - strpar1 = f"tau = {self.tau:4.2e}\t\tbacktrack = {self.backtracking}\tbeta = {self.beta}" - strpar2 = ( - f"epsg = {epsg_print}\t\tniter = {self.niter}\t\ttol = {str(self.tol)}" - ) + + strpar1 = f"tau = {self.tau:4.2e}\t\tbacktrack = {self.backtracking}" + if self.niter is not None: + strpar2 = f"beta = {self.beta}\tepsg = {epsg_print}\t\tniter = {self.niter}\t\ttol = {str(self.tol)}" + else: + strpar2 = ( + f"beta = {self.beta}\t epsg = {epsg_print}\t\ttol = {str(self.tol)}" + ) strpar3 = f"niterback = {self.niterback}\t\tacceleration = {self.acceleration}" print(strpar) print(strpar1) @@ -535,7 +540,7 @@ def setup( # type: ignore[override] # check if epgs is a vector self.epsg = np.asarray(epsg, dtype=float) if self.epsg.size == 1: - self.epsg = self.epsg * np.ones(niter) + self.epsg = epsg * np.ones(niter) epsg_print = str(self.epsg[0]) else: epsg_print = "Multi" @@ -3459,3 +3464,714 @@ def solve( # type: ignore[override] ) self.finalize(65, show) return x, self.iiter, self.cost + + +class DouglasRachfordSplitting(Solver): + r"""Douglas-Rachford Splitting + + Solves the following minimization problem using Douglas-Rachford Splitting + algorithm: + + .. math:: + + \mathbf{x} = \argmin_\mathbf{x} f(\mathbf{x}) + g(\mathbf{x}) + + where :math:`f(\mathbf{x})` and :math:`g(\mathbf{x})` are any convex + functions that has known proximal operators. + + Notes + ----- + The Douglas-Rachford Splitting algorithm can be expressed by the following + recursion [1]_, [2]_, [3]_, [4]_: + + .. math:: + + \mathbf{x}^{k} &= \prox_{\tau g}(\mathbf{y}^k) \\ + \mathbf{y}^{k+1} &= \mathbf{y}^{k} + + \eta (\prox_{\tau f}(2 \mathbf{x}^{k} - \mathbf{y}^{k}) + - \mathbf{x}^{k}) + + .. [1] Patrick L. Combettes and Jean-Christophe Pesquet. 2011. Proximal + Splitting Methods in Signal Processing. In Fixed-Point Algorithms for + Inverse Problems in Science and Engineering, Springer, pp. 185-212. + Algorithm 10.15. + https://doi.org/10.1007/978-1-4419-9569-8_10 + .. [2] Scott B. Lindstrom and Brailey Sims. 2021. Survey: Sixty Years of + Douglas-Rachford. Journal of the Australian Mathematical Society, 110, + 3, 333-370. Eq.(15). https://doi.org/10.1017/S1446788719000570 + https://arxiv.org/abs/1809.07181 + .. [3] Ryu, E.K., Yin, W., 2022. Large-Scale Convex Optimization: Algorithms + & Analyses via Monotone Operators. Cambridge University Press, + Cambridge. Eq.(2.18). https://doi.org/10.1017/9781009160865 + https://large-scale-book.mathopt.com/ + .. [4] Combettes, P.L., Pesquet, J.-C., 2008. A proximal decomposition + method for solving convex variational inverse problems. Inverse Problems + 24, 065014. Proposition 3.2. https://doi.org/10.1088/0266-5611/24/6/065014 + https://arxiv.org/abs/0807.2617 + + """ + + def _print_setup(self, xcomplex: bool = False) -> None: + self._print_solver(nbar=65) + + strpar = ( + f"Proximal operator (f): {type(self.proxf).__name__}\n" + f"Proximal operator (g): {type(self.proxg).__name__}\n" + ) + strpar1 = f"tau = {self.tau:6e}\teta = {self.eta:6e}\tniter = {self.niter}" + strpar2 = f"gfirst = {self.gfirst}\ttol = {self.tol}" + print(strpar) + print(strpar1) + print(strpar2) + + print("-" * 65 + "\n") + if not xcomplex: + head1 = " Itn x[0] f g J=f+g" + else: + head1 = ( + " Itn x[0] f g J=f+g" + ) + print(head1) + + def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: + if self.tol is None: + pf, pg = self.proxf(x), self.proxg(x) + self.pfg = pf + pg + x0 = to_numpy(x[0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = ( + f"{self.iiter:6g} " + + strx + + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e}" + ) + print(msg) + + def setup( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + tau: float, + eta: float = 1.0, + niter: int = 10, + gfirst: bool = True, + tol: float | None = None, + callbacky: bool = False, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Setup solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector + tau : :obj:`float` + Positive scalar weight + eta : :obj:`float`, optional + Relaxation parameter (must be between 0 and 2, 0 excluded). + niter : :obj:`int`, optional + Number of iterations of iterative scheme + gfirst : :obj:`bool`, optional + Apply Proximal of operator ``g`` first (``True``) or Proximal of + operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + callbacky : :obj:`bool`, optional + Modify callback signature to (``callback(x, y)``) + when ``callbacky=True`` + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + y : :obj:`numpy.ndarray` + Initial guess for the auxiliary variable + + """ + self.proxf = proxf + self.proxg = proxg + self.tau = tau + self.eta = eta + self.niter = niter + self.gfirst = gfirst + self.tol = tol + self.callbacky = callbacky + + self.ncp = get_array_module(x0) + + # initialize solver + x = x0.copy() + y = x0.copy() + + # create variables to track the objective function and iterations + self.pfg, self.pfgold = np.inf, np.inf + self.cost: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(np.iscomplexobj(x0)) + return x, y + + def step( + self, x: NDArray, y: NDArray, show: bool = False + ) -> tuple[NDArray, NDArray]: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + DouglasRachfordSplitting algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by a step of the + DouglasRachfordSplitting algorithm + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + y : :obj:`numpy.ndarray` + Updated additional model vector + + """ + # proximal steps + if self.gfirst: + x = self.proxg.prox(y, self.tau) + y = y + self.eta * (self.proxf.prox(2 * x - y, self.tau) - x) + else: + x = self.proxf.prox(y, self.tau) + y = y + self.eta * (self.proxg.prox(2 * x - y, self.tau) - x) + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfgold = self.pfg + pf = self.proxf(x) + pg = self.proxg(x) + self.pfg = pf + pg + if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: + self.tolbreak = True + else: + pf, pg = 0.0, 0.0 + + self.iiter += 1 + if show: + self._print_step(x, pf, pg) + if self.tol is not None or show: + self.cost.append(float(self.pfg)) + return x, y + + def run( + self, + x: NDArray, + y: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray]: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the DouglasRachfordSplitting algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the DouglasRachfordSplitting algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, y = self.step(x, y, showstep) + if self.callbacky: + self.callback(x, y) + else: + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, y + + def solve( # type: ignore[override] + self, + proxf: ProxOperator, + proxg: ProxOperator, + x0: NDArray, + tau: float, + eta: float = 1.0, + niter: int = 10, + gfirst: bool = True, + tol: float | None = None, + callbacky: bool = False, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxf : :obj:`pyproximal.ProxOperator` + Proximal operator of f function + proxg : :obj:`pyproximal.ProxOperator` + Proximal operator of g function + x0 : :obj:`numpy.ndarray` + Initial vector + tau : :obj:`float` + Positive scalar weight + eta : :obj:`float`, optional + Relaxation parameter (must be between 0 and 2, 0 excluded). + niter : :obj:`int`, optional + Number of iterations of iterative scheme + gfirst : :obj:`bool`, optional + Apply Proximal of operator ``g`` first (``True``) or Proximal of + operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + callbacky : :obj:`bool`, optional + Modify callback signature to (``callback(x, y)``) + when ``callbacky=True`` + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model + iiter : :obj:`int` + Number of executed iterations + cost : :obj:`list` + History of the objective function + + """ + x, y = self.setup( + proxf=proxf, + proxg=proxg, + x0=x0, + tau=tau, + eta=eta, + niter=niter, + gfirst=gfirst, + tol=tol, + callbacky=callbacky, + show=show, + ) + + x, y = self.run(x, y, niter, show=show, itershow=itershow) + self.finalize(65, show) + return x, y, self.iiter, self.cost + + +class PPXA(Solver): + r"""Parallel Proximal Algorithm (PPXA) + + Solves the following minimization problem using + Parallel Proximal Algorithm (PPXA): + + .. math:: + + \mathbf{x} = \argmin_\mathbf{x} \sum_{i=1}^m f_i(\mathbf{x}) + + where :math:`f_i(\mathbf{x})` are any convex + functions that has known proximal operators. + + See Also + -------- + ConsensusADMM: Consensus ADMM + + Notes + ----- + The Parallel Proximal Algorithm (PPXA) can be expressed by the following + recursion [1]_, [2]_, [3]_, [4]_: + + * :math:`\mathbf{y}_{i}^{0} = \mathbf{x}` or :math:`\mathbf{y}_{i}^{0} = \mathbf{x}_{i}` for :math:`i=1,\ldots,m` + * :math:`\mathbf{x}^{0} = \sum_{i=1}^m w_i \mathbf{y}_{i}^{0}` + * for :math:`k = 1, \ldots` + + * for :math:`i = 1, \ldots, m` + + * :math:`\mathbf{p}_{i}^{k} = \prox_{\frac{\tau}{w_i} f_i} (\mathbf{y}_{i}^{k})` + + * :math:`\mathbf{p}^{k} = \sum_{i=1}^{m} w_i \mathbf{p}_{i}^{k}` + * for :math:`i = 1, \ldots, m` + + * :math:`\mathbf{y}_{i}^{k+1} = \mathbf{y}_{i}^{k} + \eta (2 \mathbf{p}^{k} - \mathbf{x}^{k} - \mathbf{p}_i^{k})` + + * :math:`\mathbf{x}^{k+1} = \mathbf{x}^{k} + \eta (\mathbf{p}^{k} - \mathbf{x}^{k})` + + where :math:`0 < \eta < 2` and + :math:`\sum_{i=1}^m w_i = 1, \ 0 < w_i < 1`. + In the current implementation, :math:`w_i = 1 / m` when not provided. + + References + ---------- + .. [1] Combettes, P.L., Pesquet, J.-C., 2008. A proximal decomposition + method for solving convex variational inverse problems. Inverse Problems + 24, 065014. Algorithm 3.1. https://doi.org/10.1088/0266-5611/24/6/065014 + https://arxiv.org/abs/0807.2617 + .. [2] Combettes, P.L., Pesquet, J.-C., 2011. Proximal Splitting Methods in + Signal Processing, in Fixed-Point Algorithms for Inverse Problems in + Science and Engineering, Springer, pp. 185-212. Algorithm 10.27. + https://doi.org/10.1007/978-1-4419-9569-8_10 + .. [3] Bauschke, H.H., Combettes, P.L., 2011. Convex Analysis and Monotone + Operator Theory in Hilbert Spaces, 1st ed, CMS Books in Mathematics. + Springer, New York, NY. Proposition 27.8. + https://doi.org/10.1007/978-1-4419-9467-7 + .. [4] Ryu, E.K., Yin, W., 2022. Large-Scale Convex Optimization: Algorithms + & Analyses via Monotone Operators. Cambridge University Press, + Cambridge. Exercise 2.38 https://doi.org/10.1017/9781009160865 + https://large-scale-book.mathopt.com/ + + """ + + def _print_setup(self, xcomplex: bool = False) -> None: + self._print_solver(nbar=65) + + strpar = "\n".join( + [ + f"Proximal operator (f{i}): {type(proxf).__name__}" + for i, proxf in enumerate(self.proxfs) + ] + ) + strpar1 = f"tau = {self.tau:6e}\teta = {self.eta:6e}" + strpar2 = f"weights = {self.weights}" + strpar3 = f"niter = {self.niter}\ttol = {self.tol}" + print(strpar) + print(strpar1) + print(strpar2) + print(strpar3) + + print("-" * 65 + "\n") + if not xcomplex: + head1 = " Itn x[0] J=sum_i f_i" + else: + head1 = " Itn x[0] J=sum_i f_i" + print(head1) + + def _print_step(self, x: NDArray) -> None: + if self.tol is None: + self.pf = self.ncp.sum([self.proxfs[i](x) for i in range(self.nprox)]) + self.pf = self.ncp.sum([self.proxfs[i](x) for i in range(self.nprox)]) + x0 = to_numpy(x[0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = f"{self.iiter:6g} " + strx + f"{self.pf:10.3e}" + print(msg) + + def setup( # type: ignore[override] + self, + proxfs: list[ProxOperator], + x0: NDArray | list[NDArray], + tau: float, + eta: float = 1.0, + weights: NDArray | list[float] | None = None, + niter: int = 1000, + tol: float | None = 1e-7, + show: bool = False, + ) -> tuple[NDArray, NDArray]: + r"""Setup solver + + Parameters + ---------- + proxfs : :obj:`list` + A list of proximable functions :math:`f_1, \ldots, f_m`. + x0 : :obj:`numpy.ndarray` or :obj:`list` + Initial vector :math:`\mathbf{x}` for all :math:`f_i` if 1-dimensional array + is provided, or initial vectors :math:`\mathbf{x}_{i}` for each :math:`f_i` + for :math:`i=1,\ldots,m` if a :obj:`list` of 1-dimensional arrays or a 2-dimensional + array of size ``(m, d)`` is provided, where ``d`` is the dimension of :math:`\mathbf{x}_{i}`. + tau : :obj:`float` + Positive scalar weight + eta : :obj:`float`, optional + Relaxation parameter (must be between 0 and 2, 0 excluded). + weights : :obj:`numpy.ndarray` or :obj:`list` or :obj:`None`, optional + Weights :math:`\sum_{i=1}^m w_i = 1, \ 0 < w_i < 1`, + Defaults to None, which means :math:`w_1 = \cdots = w_m = \frac{1}{m}.` + niter : :obj:`int`, optional + Number of iterations of iterative scheme. + tol : :obj:`float`, optional + Tolerance on change of the solution (used as stopping criterion). + If ``tol=0``, run until ``niter`` is reached. + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + y : :obj:`numpy.ndarray` + Initial guess for the auxiliary variable(s) + + """ + self.proxfs = proxfs + self.tau = tau + self.eta = eta + self.weights = weights + self.niter = niter + self.tol = tol + + self.ncp = get_array_module(x0) + + # initialize solver + self.nprox = len(proxfs) + if weights is None: + self.w = self.ncp.full(self.nprox, 1.0 / self.nprox) + else: + self.w = self.ncp.asarray(weights) + + if isinstance(x0, list) or x0.ndim == 2: + y = self.ncp.asarray(x0) # yi_0 = xi_0, for i = 1, ..., m + else: + y = self.ncp.full( + (self.nprox, x0.size), x0 + ) # y1_0 = y2_0 = ... = ym_0 = x0 + + x = self.ncp.mean(y, axis=0) + + # create variables to track the objective function and iterations + self.pf, self.pfold = np.inf, np.inf + self.cost: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(np.iscomplexobj(x0)) + return x, y + + def step( + self, x: NDArray, y: NDArray, show: bool = False + ) -> tuple[NDArray, NDArray]: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + PPXA algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by a step of the + PPXA algorithm + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vector + y : :obj:`numpy.ndarray` + Updated additional model vector + + """ + x_old = x.copy() + + # proximal steps + p = self.ncp.stack( + [self.proxfs[i].prox(y[i], self.tau / self.w[i]) for i in range(self.nprox)] + ) + pn = self.ncp.sum(self.w[:, None] * p, axis=0) + y = y + self.eta * (2 * pn - x - p) + x = x + self.eta * (pn - x) + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + # if self.tol is not None: + # self.pfold = self.pf + # self.pf = self.ncp.sum([self.proxfs[i](x) for i in range(self.nprox)]) + # if np.abs(1.0 - self.pf / self.pfold) < self.tol: + # self.tolbreak = True + # else: + # self.pf = 0.0 + + # tolerance check: break iterations if solution does + # not changeabove tolerance + if self.tol is not None: + if self.ncp.abs(x - x_old).max() < self.tol: + self.tolbreak = True + + self.iiter += 1 + if show: + self._print_step(x) + if self.tol is not None or show: + self.cost.append(float(self.pf)) + return x, y + + def run( + self, + x: NDArray, + y: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray]: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the PPXA algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the PPXA algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model(s) + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, y = self.step(x, y, showstep) + self.callback(x, y) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, y + + def solve( # type: ignore[override] + self, + proxfs: list[ProxOperator], + x0: NDArray | list[NDArray], + tau: float, + eta: float = 1.0, + weights: NDArray | list[float] | None = None, + niter: int = 1000, + tol: float | None = 1e-7, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxfs : :obj:`list` + A list of proximable functions :math:`f_1, \ldots, f_m`. + x0 : :obj:`numpy.ndarray` or :obj:`list` + Initial vector :math:`\mathbf{x}` for all :math:`f_i` if 1-dimensional array + is provided, or initial vectors :math:`\mathbf{x}_{i}` for each :math:`f_i` + for :math:`i=1,\ldots,m` if a :obj:`list` of 1-dimensional arrays or a 2-dimensional + array of size ``(m, d)`` is provided, where ``d`` is the dimension of :math:`\mathbf{x}_{i}`. + tau : :obj:`float` + Positive scalar weight + eta : :obj:`float`, optional + Relaxation parameter (must be between 0 and 2, 0 excluded). + weights : :obj:`numpy.ndarray` or :obj:`list` or :obj:`None`, optional + Weights :math:`\sum_{i=1}^m w_i = 1, \ 0 < w_i < 1`, + Defaults to None, which means :math:`w_1 = \cdots = w_m = \frac{1}{m}.` + niter : :obj:`int`, optional + Number of iterations of iterative scheme. + tol : :obj:`float`, optional + Tolerance on change of the solution (used as stopping criterion). + If ``tol=0``, run until ``niter`` is reached. + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + y : :obj:`numpy.ndarray` + Additional estimated model(s) + iiter : :obj:`int` + Number of executed iterations + cost : :obj:`list` + History of the objective function + + """ + x, y = self.setup( + proxfs=proxfs, + x0=x0, + tau=tau, + eta=eta, + weights=weights, + niter=niter, + tol=tol, + show=show, + ) + + x, y = self.run(x, y, niter, show=show, itershow=itershow) + self.finalize(65, show) + return x, y, self.iiter, self.cost + + +# TOD0: +# - make niter optional in setup and fix and prints +# in _print_setup to not include if not provided From d66b485ee06c7d435631174e36337de51fc4cc30 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Sat, 20 Jun 2026 14:55:30 +0100 Subject: [PATCH 17/24] test: improved test coverage of primal solvers --- pytests/test_solver.py | 241 +++++++++++++++++++++++++++++++++++------ 1 file changed, 210 insertions(+), 31 deletions(-) diff --git a/pytests/test_solver.py b/pytests/test_solver.py index c5f38c6..408ac5a 100644 --- a/pytests/test_solver.py +++ b/pytests/test_solver.py @@ -1,5 +1,3 @@ -from typing import Any - import numpy as np import pytest from numpy.testing import assert_array_almost_equal @@ -12,6 +10,8 @@ ADMML2, HQS, PPXA, + AcceleratedProximalGradient, + AndersonProximalGradient, ConsensusADMM, DouglasRachfordSplitting, GeneralizedProximalGradient, @@ -168,7 +168,9 @@ def test_ProximalPoint(par): # Proximal point algorithm with quadatic function quad = Quadratic(Op=MatrixMult(A), b=-y, niter=2) - xpp = ProximalPoint(quad, x0=np.zeros_like(x), tau=0.1, niter=1000, tol=0) + xpp = ProximalPoint( + quad, x0=np.zeros_like(x), tau=0.1, niter=1000, tol=0, show=True + ) assert_array_almost_equal(xpp, x, decimal=2) @@ -201,14 +203,7 @@ def test_PG_ISTA(par): # ISTA/FISTA eps = 5e-1 xista = solver( - Rop, - y, - niter=100, - alpha=tau, - eps=eps, - tol=1e-8, - monitorres=False, - show=False, + Rop, y, niter=100, alpha=tau, eps=eps, tol=1e-8, monitorres=False )[0] # PG @@ -224,6 +219,7 @@ def test_PG_ISTA(par): acceleration=acceleration, niter=100, tol=1e-8, + show=True, ) assert_array_almost_equal(xpg, xista, decimal=2) @@ -255,7 +251,13 @@ def test_PG_GPG(par): l2 = L2(Op=Rop, b=y, niter=10, warm=True) l1 = L1(sigma=5e-1) xpg = ProximalGradient( - l2, l1, x0=np.zeros(m), tau=tau, niter=100, acceleration="fista" + l2, + l1, + x0=np.zeros(m), + tau=tau, + niter=100, + acceleration="fista", + show=True, ) # GPG @@ -272,11 +274,162 @@ def test_PG_GPG(par): tau=tau, niter=100, acceleration="fista", + show=True, ) assert_array_almost_equal(xpg, xgpg, decimal=2) +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) +@pytest.mark.parametrize("eta", [1, 0.8]) +@pytest.mark.parametrize("acceleration", [None, "fista", "vandenberghe"]) +def test_PG_backtracking(par, eta, acceleration): + """Check equivalency of ProximalGradient with and without tau (aka backtracking)""" + np.random.seed(0) + n, m = par["n"], par["m"] + + # Define sparse model + x = np.zeros(m) + x[2], x[4] = 1, 0.5 + + # Random mixing matrix + R = np.random.normal(0.0, 1.0, (n, m)) + Rop = MatrixMult(R) + + y = Rop @ x + + # Step size + + # PG with tau + L = (Rop.H * Rop).eigs(1).real + tau = 0.99 / L + + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l1 = L1(sigma=5e-1) + xpg = ProximalGradient( + l2, + l1, + x0=np.zeros(m), + tau=tau, + eta=eta, + niter=100, + acceleration=acceleration, + show=True, + ) + + # PG without tau + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l1 = L1(sigma=5e-1) + xpgback = ProximalGradient( + l2, + l1, + x0=np.zeros(m), + backtracking=True, + eta=eta, + niter=100, + acceleration=acceleration, + show=True, + ) + + assert_array_almost_equal(xpg, xpgback, decimal=2) + + +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) +def test_PG_AcceleratedPG(par): + """Check equivalency of ProximalGradient and AcceleratedProximalGradient""" + np.random.seed(0) + n, m = par["n"], par["m"] + + # Define sparse model + x = np.zeros(m) + x[2], x[4] = 1, 0.5 + + # Random mixing matrix + R = np.random.normal(0.0, 1.0, (n, m)) + Rop = MatrixMult(R) + + y = Rop @ x + + # Step size + L = (Rop.H * Rop).eigs(1).real + tau = 0.99 / L + + # PG + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l1 = L1(sigma=5e-1) + xpg = ProximalGradient( + l2, + l1, + x0=np.zeros(m), + tau=tau, + niter=100, + acceleration="fista", + show=True, + ) + + # APG + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l1 = L1(sigma=5e-1) + xapg = AcceleratedProximalGradient( + l2, + l1, + x0=np.zeros(m), + tau=tau, + niter=100, + ) + + assert_array_almost_equal(xpg, xapg, decimal=2) + + +@pytest.mark.parametrize("par", [(par1), (par2), (par3)]) +def test_PG_AndersonPG(par): + """Check equivalency of ProximalGradient and AndersonProximalGradient""" + np.random.seed(0) + n, m = par["n"], par["m"] + + # Define sparse model + x = np.zeros(m) + x[2], x[4] = 1, 0.5 + + # Random mixing matrix + R = np.random.normal(0.0, 1.0, (n, m)) + Rop = MatrixMult(R) + + y = Rop @ x + + # Step size + L = (Rop.H * Rop).eigs(1).real + tau = 0.99 / L + + # PG + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l1 = L1(sigma=5e-1) + xpg = ProximalGradient( + l2, + l1, + x0=np.zeros(m), + tau=tau, + niter=100, + acceleration="fista", + show=True, + ) + + # AndersonPG + l2 = L2(Op=Rop, b=y, niter=10, warm=True) + l1 = L1(sigma=5e-1) + xapg = AndersonProximalGradient( + l2, + l1, + x0=np.zeros(m), + tau=tau, + niter=100, + nhistory=5, + show=True, + ) + + assert_array_almost_equal(xpg, xapg, decimal=2) + + @pytest.mark.parametrize("par", [(par1), (par2), (par3)]) def test_PG_TwiST(par): """Check that PG/TwiST can be used to solve a sparsity regularized objective function @@ -304,7 +457,7 @@ def test_PG_TwiST(par): l2 = L2(Op=Rop, b=y, niter=10, warm=True) l1 = L1(sigma=5e-1) xpg = ProximalGradient( - l2, l1, x0=np.zeros(m), tau=tau, niter=100, acceleration="fista" + l2, l1, x0=np.zeros(m), tau=tau, niter=100, acceleration="fista", show=True ) # TwiST @@ -318,13 +471,15 @@ def test_PG_TwiST(par): x0=np.zeros(m), eigs=eigs, niter=100, + show=True, ) assert_array_almost_equal(xpg, xtwist, decimal=2) @pytest.mark.parametrize("par", [(par1), (par2), (par3)]) -def test_HQS_ADMM_L2(par): +@pytest.mark.parametrize("gfirst", [False, True]) +def test_HQS_ADMM_L2(par, gfirst): """Check that HQS/ADMM can be used to solve a pure L2-based objective function (and compare with LSQR - note that despite the trajectory will be different, they should converge to the same solution) @@ -363,19 +518,22 @@ def test_HQS_ADMM_L2(par): # HQS l2 = L2(Op=Rop, b=y, niter=10, warm=True) l2reg = L2(sigma=eps) - xhqs = HQS(l2, l2reg, x0=np.zeros(m), tau=tau, niter=1000)[0] + xhqs = HQS( + l2, l2reg, x0=np.zeros(m), tau=tau, gfirst=gfirst, niter=1000, show=True + )[0] # ADMM l2 = L2(Op=Rop, b=y, niter=10, warm=True) l2reg = L2(sigma=eps) - xadmm = ADMM(l2, l2reg, x0=np.zeros(m), tau=tau, niter=1000)[0] + xadmm = ADMM(l2, l2reg, x0=np.zeros(m), tau=tau, niter=1000, show=True)[0] assert_array_almost_equal(xl2, xhqs, decimal=2) assert_array_almost_equal(xl2, xadmm, decimal=2) @pytest.mark.parametrize("par", [(par1), (par2), (par3)]) -def test_ADMM_ADMML2(par): +@pytest.mark.parametrize("gfirst", [False, True]) +def test_ADMM_ADMML2(par, gfirst): """Check equivalency of ADMM and ADMML2 when the f function is a L2 term """ @@ -400,12 +558,21 @@ def test_ADMM_ADMML2(par): # ADMM l2 = L2(Op=Rop, b=y, niter=10, warm=True) l2reg = L2(sigma=eps) - xadmm = ADMM(l2, l2reg, x0=np.zeros(m), tau=tau, niter=100)[0] + xadmm = ADMM(l2, l2reg, x0=np.zeros(m), tau=tau, niter=100, show=True)[0] # ADMML2 l2reg = L2(sigma=eps) xadmml2 = ADMML2( - l2reg, Rop, y, Aop, x0=np.zeros(m), tau=tau, niter=100, iter_lim=10 + l2reg, + Rop, + y, + Aop, + x0=np.zeros(m), + tau=tau, + gfirst=gfirst, + niter=100, + iter_lim=10, + show=True, )[0] assert_array_almost_equal(xadmm, xadmml2, decimal=2) @@ -439,15 +606,15 @@ def test_ADMM_LinearizedADMM(par): # ADMM l2 = L2(Op=Rop, b=y, niter=10, warm=True) l2reg = L2(sigma=eps) - xadmm = ADMM(l2, l2reg, x0=np.zeros(m), tau=tau, niter=100)[0] + xadmm = ADMM(l2, l2reg, x0=np.zeros(m), tau=tau, niter=100, show=True)[0] # LinearizedADMM l2 = L2(Op=Rop, b=y, niter=10, warm=True) l2reg = L2(sigma=eps) Aop = Identity(m) - xladmm = LinearizedADMM(l2, l2reg, Aop, x0=np.zeros(m), tau=tau, mu=mu, niter=100)[ - 0 - ] + xladmm = LinearizedADMM( + l2, l2reg, Aop, x0=np.zeros(m), tau=tau, mu=mu, niter=100, show=True + )[0] assert_array_almost_equal(xadmm, xladmm, decimal=2) @@ -477,20 +644,20 @@ def test_ADMM_DRS(par): # ADMM l2 = L2(Op=Rop, b=y, niter=10, warm=True) l1 = L1(sigma=5e-1) - xadmm, zadmm = ADMM(l2, l1, x0=np.zeros(m), tau=tau, niter=100) + xadmm, zadmm = ADMM(l2, l1, x0=np.zeros(m), tau=tau, niter=100, show=True) # DRS with g first l2 = L2(Op=Rop, b=y, niter=10, warm=True) l1 = L1(sigma=5e-1) xdrs_g, ydrs_g = DouglasRachfordSplitting( - l2, l1, x0=np.zeros(m), tau=tau, niter=100, gfirst=True + l2, l1, x0=np.zeros(m), tau=tau, niter=100, gfirst=True, show=True ) # DRS with f first l2 = L2(Op=Rop, b=y, niter=10, warm=True) l1 = L1(sigma=5e-1) xdrs_f, ydrs_f = DouglasRachfordSplitting( - l2, l1, x0=np.zeros(m), tau=tau, niter=100, gfirst=False + l2, l1, x0=np.zeros(m), tau=tau, niter=100, gfirst=False, show=True ) assert_array_almost_equal(xadmm, xdrs_g, decimal=2) @@ -498,7 +665,8 @@ def test_ADMM_DRS(par): @pytest.mark.parametrize("par", [(par1), (par2), (par3)]) -def test_PPXA_with_ADMM(par: dict[str, Any]) -> None: +@pytest.mark.parametrize("weights", [None, (0.5, 0.5)]) +def test_PPXA_with_ADMM(par, weights) -> None: """Check equivalency of PPXA and ADMM when using a single regularization term """ @@ -529,18 +697,21 @@ def test_PPXA_with_ADMM(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=2000, # niter=1500 makes this test fail for seeds 0 to 499 + show=True, ) xppxa = PPXA( [l2, l1], x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), + weights=weights, + show=True, ) assert_array_almost_equal(xppxa, xadmm, decimal=2) @pytest.mark.parametrize("par", [(par1), (par2), (par3)]) -def test_PPXA_with_GPG(par: dict[str, Any]) -> None: +def test_PPXA_with_GPG(par) -> None: """Check equivalency of PPXA and GeneralizedProximalGradient""" np.random.seed(0) @@ -577,18 +748,20 @@ def test_PPXA_with_GPG(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=200, # niter=150 makes this test fail for seeds 0 to 499 + show=True, ) xppxa = PPXA( [l2_1, l2_2, l1_1, l1_2], x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), + show=True, ) assert_array_almost_equal(xppxa, xgpg, decimal=2) @pytest.mark.parametrize("par", [(par1), (par2), (par3)]) -def test_ConsensusADMM_with_ADMM(par: dict[str, Any]) -> None: +def test_ConsensusADMM_with_ADMM(par) -> None: """Check equivalency of ConsensusADMM and ADMM when two proximable functions """ @@ -619,18 +792,20 @@ def test_ConsensusADMM_with_ADMM(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=2000, # niter=1500 makes this test fail for seeds 0 to 499 + show=True, ) xcadmm = ConsensusADMM( [l2, l1], x0=np.random.normal(0.0, 1.0, m), # x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), + show=True, ) assert_array_almost_equal(xcadmm, xadmm, decimal=2) @pytest.mark.parametrize("par", [(par1), (par2), (par3)]) -def test_ConsensusADMM_with_ADMM_for_Lasso(par: dict[str, Any]) -> None: +def test_ConsensusADMM_with_ADMM_for_Lasso(par) -> None: """Check equivalency of ConsensusADMM and ADMM when more than two proximable functions for lasso """ @@ -677,6 +852,7 @@ def test_ConsensusADMM_with_ADMM_for_Lasso(par: dict[str, Any]) -> None: x0=np.random.normal(0.0, 1.0, m), # x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), niter=20000, # niter=15000 makes this test fail for seeds 0 to 499 + show=True, ) # 1/2 || [R1; R2; R3] ||_2^2 + ||x||_1 @@ -686,13 +862,14 @@ def test_ConsensusADMM_with_ADMM_for_Lasso(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=15000, # niter=10000 makes this test fail for seeds 0 to 499 + show=True, ) assert_array_almost_equal(xcadmm, xadmm, decimal=2) @pytest.mark.parametrize("par", [(par1), (par2), (par3)]) -def test_ConsensusADMM_with_GPG(par: dict[str, Any]) -> None: +def test_ConsensusADMM_with_GPG(par) -> None: """Check equivalency of ConsensusADMM and GeneralizedProximalGradient""" np.random.seed(0) @@ -730,11 +907,13 @@ def test_ConsensusADMM_with_GPG(par: dict[str, Any]) -> None: x0=np.zeros(m), tau=tau, niter=200, # niter=150 makes this test fail for seeds 0 to 499 + show=True, ) xppxa = ConsensusADMM( [l2_1, l2_2, l1_1, l1_2], x0=np.zeros(m), tau=np.random.uniform(3 * tau, 5 * tau), + show=True, ) assert_array_almost_equal(xppxa, xgpg, decimal=2) From 6464eaed3156cd8a49f9453df8897141cf972e35 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Sat, 20 Jun 2026 14:59:07 +0100 Subject: [PATCH 18/24] feat: added ConsensusADMM and improved overall handling of tol --- pyproximal/optimization/cls_primal.py | 369 ++++++++++++++++++++++++-- 1 file changed, 341 insertions(+), 28 deletions(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index dc7f1c4..b44b8d7 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -9,6 +9,8 @@ "LinearizedADMM", "TwIST", "DouglasRachfordSplitting", + "PPXA", + "ConsensusADMM", ] from collections.abc import Sequence @@ -184,7 +186,7 @@ def setup( # type: ignore[override] x0: NDArray, tau: float, niter: int | None = None, - tol: float = 1e-4, + tol: float | None = None, show: bool = False, ) -> NDArray: r"""Setup solver @@ -324,7 +326,7 @@ def solve( # type: ignore[override] x0: NDArray, tau: float, niter: int = 10, - tol: float = 1e-4, + tol: float | None = None, show: bool = False, itershow: tuple[int, int, int] = (10, 10, 10), ) -> tuple[NDArray, int, NDArray]: @@ -431,13 +433,11 @@ def _print_setup(self, epsg_print: str, xcomplex: bool = False) -> None: ) strpar1 = f"tau = {self.tau:4.2e}\t\tbacktrack = {self.backtracking}" + strpar2 = f"beta = {self.beta}\t\tepsg = {epsg_print}\t\tacceleration = {self.acceleration}" if self.niter is not None: - strpar2 = f"beta = {self.beta}\tepsg = {epsg_print}\t\tniter = {self.niter}\t\ttol = {str(self.tol)}" + strpar3 = f"niter = {self.niter}\t\tniterback = {self.niterback}\t\ttol = {str(self.tol)}" else: - strpar2 = ( - f"beta = {self.beta}\t epsg = {epsg_print}\t\ttol = {str(self.tol)}" - ) - strpar3 = f"niterback = {self.niterback}\t\tacceleration = {self.acceleration}" + strpar3 = f"niterback = {self.niterback}\t\ttol = {str(self.tol)}" print(strpar) print(strpar1) print(strpar2) @@ -3899,7 +3899,6 @@ def _print_setup(self, xcomplex: bool = False) -> None: def _print_step(self, x: NDArray) -> None: if self.tol is None: self.pf = self.ncp.sum([self.proxfs[i](x) for i in range(self.nprox)]) - self.pf = self.ncp.sum([self.proxfs[i](x) for i in range(self.nprox)]) x0 = to_numpy(x[0]) strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " msg = f"{self.iiter:6g} " + strx + f"{self.pf:10.3e}" @@ -3913,7 +3912,7 @@ def setup( # type: ignore[override] eta: float = 1.0, weights: NDArray | list[float] | None = None, niter: int = 1000, - tol: float | None = 1e-7, + tol: float | None = None, show: bool = False, ) -> tuple[NDArray, NDArray]: r"""Setup solver @@ -3937,8 +3936,8 @@ def setup( # type: ignore[override] niter : :obj:`int`, optional Number of iterations of iterative scheme. tol : :obj:`float`, optional - Tolerance on change of the solution (used as stopping criterion). - If ``tol=0``, run until ``niter`` is reached. + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached show : :obj:`bool`, optional Display iterations log @@ -4010,8 +4009,6 @@ def step( Updated additional model vector """ - x_old = x.copy() - # proximal steps p = self.ncp.stack( [self.proxfs[i].prox(y[i], self.tau / self.w[i]) for i in range(self.nprox)] @@ -4022,19 +4019,13 @@ def step( # tolerance check: break iterations if overall # objective does not decrease below tolerance - # if self.tol is not None: - # self.pfold = self.pf - # self.pf = self.ncp.sum([self.proxfs[i](x) for i in range(self.nprox)]) - # if np.abs(1.0 - self.pf / self.pfold) < self.tol: - # self.tolbreak = True - # else: - # self.pf = 0.0 - - # tolerance check: break iterations if solution does - # not changeabove tolerance if self.tol is not None: - if self.ncp.abs(x - x_old).max() < self.tol: + self.pfold = self.pf + self.pf = self.ncp.sum([self.proxfs[i](x) for i in range(self.nprox)]) + if np.abs(1.0 - self.pf / self.pfold) < self.tol: self.tolbreak = True + else: + self.pf = 0.0 self.iiter += 1 if show: @@ -4095,7 +4086,7 @@ def run( else False ) x, y = self.step(x, y, showstep) - self.callback(x, y) + self.callback(x) # check if any callback has raised a stop flag stop = _callback_stop(self.callbacks) if stop: @@ -4110,7 +4101,7 @@ def solve( # type: ignore[override] eta: float = 1.0, weights: NDArray | list[float] | None = None, niter: int = 1000, - tol: float | None = 1e-7, + tol: float | None = None, show: bool = False, itershow: tuple[int, int, int] = (10, 10, 10), ) -> tuple[NDArray, NDArray, int, NDArray]: @@ -4135,8 +4126,8 @@ def solve( # type: ignore[override] niter : :obj:`int`, optional Number of iterations of iterative scheme. tol : :obj:`float`, optional - Tolerance on change of the solution (used as stopping criterion). - If ``tol=0``, run until ``niter`` is reached. + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached show : :obj:`bool`, optional Display logs itershow : :obj:`tuple`, optional @@ -4172,6 +4163,328 @@ def solve( # type: ignore[override] return x, y, self.iiter, self.cost +class ConsensusADMM(Solver): + r"""Consensus ADMM + + Solves the following global consensus problem using ADMM: + + .. math:: + + \argmin_{\mathbf{x_1}, \mathbf{x_2}, \ldots, \mathbf{x_m}} + \sum_{i=1}^m f_i(\mathbf{x}_i) \quad \text{s.t.} + \quad \mathbf{x_1} = \mathbf{x_2} = \cdots = \mathbf{x_m} + + where :math:`f_i(\mathbf{x})` are any convex + functions that has known proximal operators. + + See Also + -------- + ADMM: Alternating Direction Method of Multipliers + PPXA: Parallel Proximal Algorithm + + Notes + ----- + The ADMM for the consensus problem can be expressed by the following + recursion [1]_, [2]_: + + * :math:`\bar{\mathbf{x}}^{0} = \mathbf{x}` + * for :math:`k = 1, \ldots` + + * for :math:`i = 1, \ldots, m` + + * :math:`\mathbf{x}_i^{k+1} = \mathrm{prox}_{\tau f_i} \left(\bar{\mathbf{x}}^{k} - \mathbf{y}_i^{k}\right)` + + * :math:`\bar{\mathbf{x}}^{k+1} = \frac{1}{m} \sum_{i=1}^m \mathbf{x}_i^{k}` + + * for :math:`i = 1, \ldots, m` + + * :math:`\mathbf{y}_i^{k+1} = \mathbf{y}_i^{k} + \mathbf{x}_i^{k+1} - \bar{\mathbf{x}}^{k+1}` + + The current implementation returns :math:`\bar{\mathbf{x}}`. + + References + ---------- + .. [1] Boyd, S., Parikh, N., Chu, E., Peleato, B., Eckstein, J., 2011. + Distributed Optimization and Statistical Learning via the Alternating + Direction Method of Multipliers. Foundations and Trends in Machine Learning, + Vol. 3, No. 1, pp 1-122. Section 7.1. https://doi.org/10.1561/2200000016 + https://stanford.edu/~boyd/papers/pdf/admm_distr_stats.pdf + .. [2] Parikh, N., Boyd, S., 2014. Proximal Algorithms. Foundations and + Trends in Optimization, Vol. 1, No. 3, pp 127-239. + Section 5.2.1. https://doi.org/10.1561/2400000003 + https://web.stanford.edu/~boyd/papers/pdf/prox_algs.pdf + + """ + + def _print_setup(self, xcomplex: bool = False) -> None: + self._print_solver(nbar=65) + + strpar = "\n".join( + [ + f"Proximal operator (f{i}): {type(proxf).__name__}" + for i, proxf in enumerate(self.proxfs) + ] + ) + strpar1 = f"tau = {self.tau:6e}\tniter = {self.niter}\ttol = {self.tol}" + print(strpar) + print(strpar1) + + print("-" * 65 + "\n") + if not xcomplex: + head1 = " Itn x[0] J=sum_i f_i" + else: + head1 = " Itn x[0] J=sum_i f_i" + print(head1) + + def _print_step(self, x: NDArray) -> None: + if self.tol is None: + self.pf = self.ncp.sum([self.proxfs[i](x) for i in range(self.nprox)]) + x0 = to_numpy(x[0]) + strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + msg = f"{self.iiter:6g} " + strx + f"{self.pf:10.3e}" + print(msg) + + def setup( # type: ignore[override] + self, + proxfs: list[ProxOperator], + x0: NDArray, + tau: float, + niter: int = 1000, + tol: float | None = None, + show: bool = False, + ) -> tuple[NDArray, NDArray, NDArray]: + r"""Setup solver + + Parameters + ---------- + proxfs : :obj:`list` + A list of proximable functions :math:`f_1, \ldots, f_m`. + x0 : :obj:`numpy.ndarray` + Initial vector + tau : :obj:`float` + Positive scalar weight + niter : :obj:`int`, optional + Number of iterations of iterative scheme. + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`numpy.ndarray` + Initial guess + x_bar : :obj:`numpy.ndarray` + Initial guess for averaged model vector + y : :obj:`numpy.ndarray` + Initial guess for additional model vector + + """ + self.proxfs = proxfs + self.tau = tau + self.niter = niter + self.tol = tol + + self.ncp = get_array_module(x0) + + # initialize solver + self.nprox = len(proxfs) + x = x0.copy() + x_bar = x0.copy() + y = self.ncp.zeros((self.nprox, x0.size), dtype=x0.dtype) + + # create variables to track the objective function and iterations + self.pf, self.pfold = np.inf, np.inf + self.cost: list[float] = [] + self.tolbreak = False + self.iiter = 0 + + # print setup + if show: + self._print_setup(np.iscomplexobj(x0)) + return x, x_bar, y + + def step( + self, x: NDArray, x_bar: NDArray, y: NDArray, show: bool = False + ) -> tuple[NDArray, NDArray, NDArray]: + r"""Run one step of solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by a step of the + ConsensusADMM algorithm + x_bar : :obj:`numpy.ndarray` + Current averaged model vector to be updated by a step of the + ConsensusADMM algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the ConsensusADMM algorithm + show : :obj:`bool`, optional + Display iteration log + + Returns + ------- + x : :obj:`numpy.ndarray` + Updated model vectors + x_bar : :obj:`numpy.ndarray` + Updated averaged model vector + y : :obj:`numpy.ndarray` + Updated additional model vector + + """ + # proximal steps + x = self.ncp.stack( + [self.proxfs[i].prox(x_bar - y[i], self.tau) for i in range(self.nprox)] + ) + x_bar = self.ncp.mean(x, axis=0) + y = y + x - x_bar + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if self.tol is not None: + self.pfold = self.pf + self.pf = self.ncp.sum([self.proxfs[i](x_bar) for i in range(self.nprox)]) + if np.abs(1.0 - self.pf / self.pfold) < self.tol: + self.tolbreak = True + else: + self.pf = 0.0 + + self.iiter += 1 + if show: + self._print_step(x_bar) + if self.tol is not None or show: + self.cost.append(float(self.pf)) + return x, x_bar, y + + def run( + self, + x: NDArray, + x_bar: NDArray, + y: NDArray, + niter: int | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray, NDArray]: + r"""Run solver + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Current model vector to be updated by multiple steps of + the PPXA algorithm + x_bar : :obj:`numpy.ndarray` + Current averaged model vector to be updated by multiple steps of + the PPXA algorithm + y : :obj:`numpy.ndarray` + Additional model vector to be updated by multiple steps of + the PPXA algorithm + niter : :obj:`int`, optional + Number of iterations. Can be set to ``None`` if already + provided in the setup call + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + x_bar : :obj:`numpy.ndarray` + Estimated averaged model + y : :obj:`numpy.ndarray` + Additional estimated model + + """ + niter = self.niter if niter is None else niter + if niter is None: + msg = "`niter` must not be None" + raise ValueError(msg) + while self.iiter < niter and not self.tolbreak: + showstep = ( + True + if show + and ( + self.iiter < itershow[0] + or niter - self.iiter < itershow[1] + or self.iiter % itershow[2] == 0 + ) + else False + ) + x, x_bar, y = self.step(x, x_bar, y, showstep) + self.callback(x_bar) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x, x_bar, y + + def solve( # type: ignore[override] + self, + proxfs: list[ProxOperator], + x0: NDArray, + tau: float, + niter: int = 1000, + tol: float | None = None, + show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), + ) -> tuple[NDArray, NDArray, NDArray, int, NDArray]: + r"""Run entire solver + + Parameters + ---------- + proxfs : :obj:`list` + A list of proximable functions :math:`f_1, \ldots, f_m`. + x0 : :obj:`numpy.ndarray` + Initial vector + tau : :obj:`float` + Positive scalar weight + niter : :obj:`int`, optional + Number of iterations of iterative scheme. + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + show : :obj:`bool`, optional + Display logs + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. + + Returns + ------- + x : :obj:`numpy.ndarray` + Estimated model + x_bar : :obj:`numpy.ndarray` + Estimated averaged model + y : :obj:`numpy.ndarray` + Additional estimated model + iiter : :obj:`int` + Number of executed iterations + cost : :obj:`list` + History of the objective function + + """ + x, x_bar, y = self.setup( + proxfs=proxfs, + x0=x0, + tau=tau, + niter=niter, + tol=tol, + show=show, + ) + + x, x_bar, y = self.run(x, x_bar, y, niter, show=show, itershow=itershow) + self.finalize(65, show) + return x, x_bar, y, self.iiter, self.cost + + # TOD0: # - make niter optional in setup and fix and prints # in _print_setup to not include if not provided +# - add the tolerance check on the solution for PPXA and +# ConsensunADMM for the function based solvers via callbacks From 8de09003819a4759f91129eecdbddf557d1165b7 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Fri, 26 Jun 2026 10:50:24 +0100 Subject: [PATCH 19/24] minor: removed unused callback method from Solver --- pyproximal/optimization/basesolver.py | 41 --------------------------- 1 file changed, 41 deletions(-) diff --git a/pyproximal/optimization/basesolver.py b/pyproximal/optimization/basesolver.py index 97739a8..008b6cd 100644 --- a/pyproximal/optimization/basesolver.py +++ b/pyproximal/optimization/basesolver.py @@ -14,8 +14,6 @@ if TYPE_CHECKING: from pyproximal.ProxOperator import ProxOperator -_units = {"B": 1, "KB": 1024, "MB": 1024**2, "GB": 1024**3} - class Solver(pSolver, metaclass=ABCMeta): # type: ignore[misc] r"""Solver @@ -163,42 +161,3 @@ def finalize(self, nbar: int = 60, show: bool = False) -> None: self.cost: NDArray = np.array(self.cost) if show: self._print_finalize(nbar=nbar) - - def callback( # noqa: B027 - self, - x: NDArray, - z: NDArray | None = None, - *args: Any, - **kwargs: Any, - ) -> None: - """Callback routine - - This routine must be passed by the user. Its function signature must contain - either a single input that contains the current solution or two inputs - that contain the current solutions for methods that apply splitting - (when using the `solve` method it will be automatically invoked after - each step of the solve) - - Parameters - ---------- - x : :obj:`numpy.ndarray` - Current solution - z : :obj:`numpy.ndarray` - Current additional solution - - Examples - -------- - >>> import numpy as np - >>> from pyproximal.optimization.cls_primal import ADMM - >>> def callback(x, z): - ... print(f"Running callback, current solutions {x} - {z}") - ... - >>> admmsolve.callback = callback - - >>> x = np.ones(2) - >>> z = np.zeros(2) - >>> admmsolve.callback(x, z) - Running callback, current solutions [1. 1.] - [0. 0.] - - """ - pass From 2de1345eccb9542a5760b94598d4148f27c26e7f Mon Sep 17 00:00:00 2001 From: mrava87 Date: Fri, 26 Jun 2026 10:52:06 +0100 Subject: [PATCH 20/24] feat: various improvements to class-based solvers to match old functional behaviour --- pyproximal/optimization/cls_primal.py | 272 ++++++++++++++++---------- 1 file changed, 170 insertions(+), 102 deletions(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index b44b8d7..0a39745 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -427,12 +427,17 @@ class ProximalGradient(Solver): def _print_setup(self, epsg_print: str, xcomplex: bool = False) -> None: self._print_solver(nbar=81) + tau_str = ( + f"{self.tau[0]:4.2e}" + if self.tau.size == 1 + else ", ".join(f"{tau:4.2e}" for tau in self.tau) + ) + strpar = ( f"Proximal operator (f): {type(self.proxf).__name__}\n" f"Proximal operator (g): {type(self.proxg).__name__}\n" ) - - strpar1 = f"tau = {self.tau:4.2e}\t\tbacktrack = {self.backtracking}" + strpar1 = f"tau = {tau_str}\t\tbacktrack = {self.backtracking}" strpar2 = f"beta = {self.beta}\t\tepsg = {epsg_print}\t\tacceleration = {self.acceleration}" if self.niter is not None: strpar3 = f"niter = {self.niter}\t\tniterback = {self.niterback}\t\ttol = {str(self.tol)}" @@ -455,10 +460,15 @@ def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: self.pfg = pf + np.sum(self.epsg[self.iiter - 1] * pg) x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + tau_str = ( + f"{self.tau[0]:4.2e}" + if self.tau.size == 1 + else ", ".join(f"{tau:4.2e}" for tau in self.tau) + ) msg = ( f"{self.iiter:6g} " + strx - + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e} {self.tau:11.2e}" + + f"{pf:10.3e} {pg:10.3e} {self.pfg:10.3e} {tau_str}" ) print(msg) @@ -472,9 +482,9 @@ def setup( # type: ignore[override] backtracking: bool = False, beta: float = 0.5, eta: float = 1.0, - niter: int = 10, - niterback: int = 100, acceleration: str | None = None, + niterback: int = 100, + niter: int = 10, tol: float | None = None, show: bool = False, ) -> tuple[NDArray, NDArray]: @@ -505,12 +515,12 @@ def setup( # type: ignore[override] Backtracking parameter (must be between 0 and 1) eta : :obj:`float`, optional Relaxation parameter (must be between 0 and 1, 0 excluded). - niter : :obj:`int`, optional - Number of iterations of iterative scheme - niterback : :obj:`int`, optional - Max number of iterations of backtracking acceleration : :obj:`str`, optional Acceleration (``None``, ``vandenberghe`` or ``fista``) + niterback : :obj:`int`, optional + Max number of iterations of backtracking + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -531,25 +541,25 @@ def setup( # type: ignore[override] self.backtracking = backtracking self.beta = beta self.eta = eta - self.niter = niter self.niterback = niterback + self.niter = niter self.tol = tol self.ncp = get_array_module(x0) # check if epgs is a vector - self.epsg = np.asarray(epsg, dtype=float) + self.epsg = self.ncp.asarray(epsg, dtype=np.float32) if self.epsg.size == 1: - self.epsg = epsg * np.ones(niter) + self.epsg = epsg * self.ncp.ones(niter, dtype=np.float32) epsg_print = str(self.epsg[0]) else: epsg_print = "Multi" # set tau - self.tau = tau if tau is None: self.backtracking = True - self.tau = 1.0 + tau = 1.0 + self.tau = self.ncp.atleast_1d(self.ncp.asarray(tau, dtype=np.float32)) # check acceleration if acceleration in [None, "None", "vandenberghe", "fista"]: @@ -566,8 +576,12 @@ def setup( # type: ignore[override] self.t = 1.0 # create variables to track the objective function and iterations - self.pfg, self.pfgold = np.inf, np.inf + pf, pg = self.proxf(x), self.proxg(x) + pfg = pf + np.sum(self.epsg[self.iiter] * pg) + self.pfg, self.pfgold = pfg, pfg + self.cost: list[float] = [] + self.cost.append(float(self.pfg)) self.tolbreak = False self.iiter = 0 @@ -577,7 +591,10 @@ def setup( # type: ignore[override] return x, y def step( - self, x: NDArray, y: NDArray, show: bool = False + self, + x: NDArray, + y: NDArray, + show: bool = False, ) -> tuple[NDArray, NDArray]: r"""Run one step of solver @@ -737,9 +754,9 @@ def solve( # type: ignore[override] backtracking: bool = False, beta: float = 0.5, eta: float = 1.0, - niter: int = 10, - niterback: int = 100, acceleration: str | None = None, + niterback: int = 100, + niter: int = 10, tol: float | None = None, show: bool = False, itershow: tuple[int, int, int] = (10, 10, 10), @@ -771,12 +788,12 @@ def solve( # type: ignore[override] Backtracking parameter (must be between 0 and 1) eta : :obj:`float`, optional Relaxation parameter (must be between 0 and 1, 0 excluded). - niter : :obj:`int`, optional - Number of iterations of iterative scheme - niterback : :obj:`int`, optional - Max number of iterations of backtracking acceleration : :obj:`str`, optional Acceleration (``None``, ``vandenberghe`` or ``fista``) + niterback : :obj:`int`, optional + Max number of iterations of backtracking + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -808,9 +825,9 @@ def solve( # type: ignore[override] backtracking=backtracking, beta=beta, eta=eta, - niter=niter, - niterback=niterback, acceleration=acceleration, + niterback=niterback, + niter=niter, tol=tol, show=show, ) @@ -865,13 +882,17 @@ class AndersonProximalGradient(Solver): def _print_setup(self, epsg_print: str, xcomplex: bool = False) -> None: self._print_solver(nbar=81) + tau_str = ( + f"{self.tau[0]:4.2e}" + if self.tau.size == 1 + else ", ".join(f"{tau:4.2e}" for tau in self.tau) + ) + strpar = ( f"Proximal operator (f): {type(self.proxf).__name__}\n" f"Proximal operator (g): {type(self.proxg).__name__}\n" ) - strpar1 = ( - f"tau = {self.tau:4.2e}\t\tepsg = {epsg_print}\t\tniter = {self.niter}" - ) + strpar1 = f"tau = {tau_str}\t\tepsg = {epsg_print}\t\tniter = {self.niter}" strpar2 = f"nhist = {self.nhistory}\t\tepsr = {self.epsr:4.2e}" strpar3 = f"guard = {str(self.safeguard)}\t\ttol = {str(self.tol)}" print(strpar) @@ -891,10 +912,15 @@ def _print_step(self, x: NDArray, pg: float | None) -> None: self.pfg = self.pf + np.sum(self.epsg[self.iiter - 1] * pg) x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " + tau_str = ( + f"{self.tau[0]:4.2e}" + if self.tau.size == 1 + else ", ".join(f"{tau:4.2e}" for tau in self.tau) + ) msg = ( f"{self.iiter:6g} " + strx - + f"{self.pf:10.3e} {pg:10.3e} {self.pfg:10.3e} {self.tau:11.2e}" + + f"{self.pf:10.3e} {pg:10.3e} {self.pfg:10.3e} {tau_str}" ) print(msg) @@ -905,10 +931,10 @@ def setup( # type: ignore[override] x0: NDArray, epsg: float | NDArray = 1.0, tau: float | NDArray = 1.0, - niter: int = 10, - nhistory: int = 10, epsr: float = 1e-10, safeguard: bool = False, + nhistory: int = 10, + niter: int = 10, tol: float | None = None, show: bool = False, ) -> tuple[NDArray, NDArray]: @@ -927,17 +953,17 @@ def setup( # type: ignore[override] tau : :obj:`float` or :obj:`numpy.ndarray`, optional Positive scalar weight, which should satisfy the following condition to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is - the Lipschitz constant of :math:`\nabla f`. N ote that :math:`\tau` + the Lipschitz constant of :math:`\nabla f`. Note that :math:`\tau` can be chosen to be a vector when dealing with problems with multiple right-hand-sides - niter : :obj:`int`, optional - Number of iterations of iterative scheme - nhistory : :obj:`int`, optional - Number of previous iterates to be kept in memory (to compute the scaling factors) epsr : :obj:`float`, optional Scaling factor for regularization added to the inverse of :math:\mathbf{R}^T \mathbf{R}` safeguard : :obj:`bool`, optional Apply safeguarding strategy to the update (``True``) or not (``False``) + nhistory : :obj:`int`, optional + Number of previous iterates to be kept in memory (to compute the scaling factors) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -953,23 +979,25 @@ def setup( # type: ignore[override] self.proxf = proxf self.proxg = proxg self.x0 = x0 - self.tau = tau - self.niter = niter - self.nhistory = nhistory self.epsr = epsr self.safeguard = safeguard + self.nhistory = nhistory + self.niter = niter self.tol = tol self.ncp = get_array_module(x0) # check if epgs is a vector - self.epsg = np.asarray(epsg, dtype=float) + self.epsg = self.ncp.asarray(epsg, dtype=np.float32) if self.epsg.size == 1: - self.epsg = self.epsg * np.ones(niter) + self.epsg = self.epsg * self.ncp.ones(niter, dtype=np.float32) epsg_print = str(self.epsg[0]) else: epsg_print = "Multi" + # set tau + self.tau = self.ncp.atleast_1d(self.ncp.asarray(tau, dtype=np.float32)) + # initialize solver y = x0 - self.tau * proxf.grad(x0) x = self.proxg.prox(y, self.epsg[0] * self.tau) @@ -999,7 +1027,10 @@ def setup( # type: ignore[override] return x, y def step( - self, x: NDArray, y: NDArray, show: bool = False + self, + x: NDArray, + y: NDArray, + show: bool = False, ) -> tuple[NDArray, NDArray]: r"""Run one step of solver @@ -1152,10 +1183,10 @@ def solve( # type: ignore[override] x0: NDArray, epsg: float | NDArray = 1.0, tau: float | NDArray = 1.0, - niter: int = 10, - nhistory: int = 10, epsr: float = 1e-10, safeguard: bool = False, + nhistory: int = 10, + niter: int = 10, tol: float | None = None, show: bool = False, itershow: tuple[int, int, int] = (10, 10, 10), @@ -1175,17 +1206,17 @@ def solve( # type: ignore[override] tau : :obj:`float` or :obj:`numpy.ndarray`, optional Positive scalar weight, which should satisfy the following condition to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is - the Lipschitz constant of :math:`\nabla f`. N ote that :math:`\tau` + the Lipschitz constant of :math:`\nabla f`. Note that :math:`\tau` can be chosen to be a vector when dealing with problems with multiple right-hand-sides - niter : :obj:`int`, optional - Number of iterations of iterative scheme - nhistory : :obj:`int`, optional - Number of previous iterates to be kept in memory (to compute the scaling factors) epsr : :obj:`float`, optional Scaling factor for regularization added to the inverse of :math:\mathbf{R}^T \mathbf{R}` safeguard : :obj:`bool`, optional Apply safeguarding strategy to the update (``True``) or not (``False``) + nhistory : :obj:`int`, optional + Number of previous iterates to be kept in memory (to compute the scaling factors) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -1214,10 +1245,10 @@ def solve( # type: ignore[override] x0=x0, epsg=epsg, tau=tau, - niter=niter, - nhistory=nhistory, epsr=epsr, safeguard=safeguard, + nhistory=nhistory, + niter=niter, tol=tol, show=show, ) @@ -1269,7 +1300,9 @@ def _print_setup(self, epsg_print: str, xcomplex: bool = False) -> None: f"Proximal operators (f): {[type(proxf).__name__ for proxf in self.proxfs]}\n" f"Proximal operators (g): {[type(proxg).__name__ for proxg in self.proxgs]}\n" ) - strpar1 = f"tau = {self.tau:4.2e}\tepsg = {epsg_print}\tniter = {self.niter}" + strpar1 = ( + f"tau = {float(self.tau):4.2e}\tepsg = {epsg_print}\tniter = {self.niter}" + ) print(strpar) print(strpar1) print("-" * 65 + "\n") @@ -1309,8 +1342,8 @@ def setup( # type: ignore[override] epsg: float | NDArray = 1.0, weights: NDArray | None = None, eta: float = 1.0, - niter: int = 10, acceleration: str | None = None, + niter: int = 10, tol: float | None = None, show: bool = False, ) -> tuple[NDArray, NDArray]: @@ -1335,10 +1368,10 @@ def setup( # type: ignore[override] eta : :obj:`float`, optional Relaxation parameter (must be between 0 and 1, 0 excluded). Note that this will be only used when ``acceleration=None``. - niter : :obj:`int`, optional - Number of iterations of iterative scheme acceleration: :obj:`str`, optional Acceleration (``None``, ``vandenberghe`` or ``fista``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -1372,9 +1405,9 @@ def setup( # type: ignore[override] raise ValueError(msg) # check if epgs is a vector - self.epsg = np.asarray(epsg, dtype=float) + self.epsg = np.asarray(epsg, dtype=np.float32) if self.epsg.size == 1: - self.epsg = epsg * np.ones(len(proxgs)) + self.epsg = epsg * self.ncp.ones(len(proxgs), dtype=np.float32) epsg_print = str(self.epsg[0]) else: epsg_print = "Multi" @@ -1406,7 +1439,10 @@ def setup( # type: ignore[override] return x, y def step( - self, x: NDArray, y: NDArray, show: bool = False + self, + x: NDArray, + y: NDArray, + show: bool = False, ) -> tuple[NDArray, NDArray]: r"""Run one step of solver @@ -1547,8 +1583,8 @@ def solve( # type: ignore[override] epsg: float | NDArray = 1.0, weights: NDArray | None = None, eta: float = 1.0, - niter: int = 10, acceleration: str | None = None, + niter: int = 10, tol: float | None = None, show: bool = False, itershow: tuple[int, int, int] = (10, 10, 10), @@ -1574,10 +1610,10 @@ def solve( # type: ignore[override] eta : :obj:`float`, optional Relaxation parameter (must be between 0 and 1, 0 excluded). Note that this will be only used when ``acceleration=None``. - niter : :obj:`int`, optional - Number of iterations of iterative scheme acceleration: :obj:`str`, optional Acceleration (``None``, ``vandenberghe`` or ``fista``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -1608,8 +1644,8 @@ def solve( # type: ignore[override] epsg=epsg, weights=weights, eta=eta, - niter=niter, acceleration=acceleration, + niter=niter, tol=tol, show=show, ) @@ -1702,9 +1738,10 @@ def setup( # type: ignore[override] x0: NDArray, tau: float | NDArray, z0: NDArray | None = None, - niter: int = 10, gfirst: bool = True, + niter: int = 10, tol: float | None = None, + callbackz: bool = False, show: bool = False, ) -> tuple[NDArray, NDArray]: r"""Setup solver @@ -1726,14 +1763,16 @@ def setup( # type: ignore[override] strategy) z0 : :obj:`numpy.ndarray`, optional Initial z vector (not required when ``gfirst=True``) - niter : :obj:`int` - Number of iterations of iterative scheme gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached + callbackz : :obj:`bool`, optional + Modify callback signature to (``callback(x, z)``) when ``callbackz=True`` show : :obj:`bool`, optional Display iterations log @@ -1752,17 +1791,18 @@ def setup( # type: ignore[override] """ self.proxf = proxf self.proxg = proxg - self.niter = niter self.gfirst = gfirst + self.niter = niter self.tol = tol + self.callbackz = callbackz self.ncp = get_array_module(x0) # check if tau is a vector - self.tau = self.ncp.asarray(tau, dtype=float) + self.tau = self.ncp.asarray(tau, dtype=np.float32) if self.tau.size == 1: tau_print = str(np.round(self.tau, 6)) - self.tau = self.tau * np.ones(niter) + self.tau = self.tau * self.ncp.ones(niter, dtype=np.float32) else: tau_print = "Variable" @@ -1784,7 +1824,10 @@ def setup( # type: ignore[override] return x, z def step( - self, x: NDArray, z: NDArray, show: bool = False + self, + x: NDArray, + z: NDArray, + show: bool = False, ) -> tuple[NDArray, NDArray]: r"""Run one step of solver @@ -1886,7 +1929,10 @@ def run( else False ) x, z = self.step(x, z, showstep) - self.callback(x) + if self.callbackz: + self.callback(x, z) + else: + self.callback(x) # check if any callback has raised a stop flag stop = _callback_stop(self.callbacks) if stop: @@ -1900,9 +1946,10 @@ def solve( # type: ignore[override] x0: NDArray, tau: float | NDArray, z0: NDArray | None = None, - niter: int = 10, gfirst: bool = True, + niter: int = 10, tol: float | None = None, + callbackz: bool = False, show: bool = False, itershow: tuple[int, int, int] = (10, 10, 10), ) -> tuple[NDArray, NDArray, int, NDArray]: @@ -1925,14 +1972,16 @@ def solve( # type: ignore[override] strategy) z0 : :obj:`numpy.ndarray`, optional Initial z vector (not required when ``gfirst=True``) - niter : :obj:`int` - Number of iterations of iterative scheme gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached + callbackz : :obj:`bool`, optional + Modify callback signature to (``callback(x, z)``) when ``callbackz=True`` show : :obj:`bool`, optional Display logs itershow : :obj:`tuple`, optional @@ -1944,7 +1993,7 @@ def solve( # type: ignore[override] ------- x : :obj:`numpy.ndarray` Estimated model - y : :obj:`numpy.ndarray` + z : :obj:`numpy.ndarray` Additional estimated model iiter : :obj:`int` Number of executed iterations @@ -1958,9 +2007,10 @@ def solve( # type: ignore[override] x0=x0, tau=tau, z0=z0, - niter=niter, gfirst=gfirst, + niter=niter, tol=tol, + callbackz=callbackz, show=show, ) @@ -2034,7 +2084,7 @@ def _print_setup(self, xcomplex: bool = False) -> None: f"Proximal operator (f): {type(self.proxf).__name__}\n" f"Proximal operator (g): {type(self.proxg).__name__}\n" ) - strpar1 = f"tau = {self.tau:6e}\tniter = {self.niter}" + strpar1 = f"tau = {float(self.tau):6e}\tniter = {self.niter}" strpar2 = f"gfirst = {self.gfirst}\t\ttol = {self.tol}" print(strpar) print(strpar1) @@ -2069,8 +2119,8 @@ def setup( # type: ignore[override] x0: NDArray, tau: float, z0: NDArray | None = None, - niter: int = 10, gfirst: bool = False, + niter: int = 10, tol: float | None = None, callbackz: bool = False, show: bool = False, @@ -2091,11 +2141,11 @@ def setup( # type: ignore[override] the Lipschitz constant of :math:`\nabla f`. z0 : :obj:`numpy.ndarray`, optional Initial z vector (not required when ``gfirst=True``) - niter : :obj:`int` - Number of iterations of iterative scheme gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -2143,7 +2193,10 @@ def setup( # type: ignore[override] return x, z def step( - self, x: NDArray, z: NDArray, show: bool = False + self, + x: NDArray, + z: NDArray, + show: bool = False, ) -> tuple[NDArray, NDArray]: r"""Run one step of solver @@ -2263,8 +2316,8 @@ def solve( # type: ignore[override] x0: NDArray, tau: float, z0: NDArray | None = None, - niter: int = 10, gfirst: bool = False, + niter: int = 10, tol: float | None = None, callbackz: bool = False, show: bool = False, @@ -2286,11 +2339,11 @@ def solve( # type: ignore[override] the Lipschitz constant of :math:`\nabla f`. z0 : :obj:`numpy.ndarray`, optional Initial z vector (not required when ``gfirst=True``) - niter : :obj:`int` - Number of iterations of iterative scheme gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -2321,8 +2374,8 @@ def solve( # type: ignore[override] x0=x0, tau=tau, z0=z0, - niter=niter, gfirst=gfirst, + niter=niter, tol=tol, callbackz=callbackz, show=show, @@ -2418,8 +2471,8 @@ def setup( # type: ignore[override] x0: NDArray, tau: float, z0: NDArray | None = None, - niter: int = 10, gfirst: bool = False, + niter: int = 10, tol: float | None = None, callbackz: bool = False, show: bool = False, @@ -2443,11 +2496,11 @@ def setup( # type: ignore[override] to guarantees convergence: :math:`\tau \in (0, 1/\lambda_{max}(\mathbf{A}^H\mathbf{A})]`. z0 : :obj:`numpy.ndarray` Initial auxiliary vector. If ``None``, initialized to ``A @ x0``. - niter : :obj:`int`, optional - Number of iterations of iterative scheme gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -2474,8 +2527,8 @@ def setup( # type: ignore[override] self.b = b self.A = A self.tau = tau - self.niter = niter self.gfirst = gfirst + self.niter = niter self.tol = tol self.callbackz = callbackz @@ -2665,8 +2718,8 @@ def solve( # type: ignore[override] x0: NDArray, tau: float, z0: NDArray | None = None, - niter: int = 10, gfirst: bool = False, + niter: int = 10, tol: float | None = None, callbackz: bool = False, show: bool = False, @@ -2693,11 +2746,11 @@ def solve( # type: ignore[override] the Lipschitz constant of :math:`\nabla f`. z0 : :obj:`numpy.ndarray`, optional Initial z vector (not required when ``gfirst=True``) - niter : :obj:`int` - Number of iterations of iterative scheme gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -2733,8 +2786,8 @@ def solve( # type: ignore[override] x0=x0, tau=tau, z0=z0, - niter=niter, gfirst=gfirst, + niter=niter, tol=tol, callbackz=callbackz, show=show, @@ -3314,7 +3367,7 @@ def step( if self.tol is not None: self.pfgold = self.pfg pf = self.proxf(x) - pg = self.proxg(self.Ax) + pg = self.proxg(x) self.pfg = pf + pg if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: self.tolbreak = True @@ -3553,8 +3606,8 @@ def setup( # type: ignore[override] x0: NDArray, tau: float, eta: float = 1.0, - niter: int = 10, gfirst: bool = True, + niter: int = 10, tol: float | None = None, callbacky: bool = False, show: bool = False, @@ -3573,11 +3626,11 @@ def setup( # type: ignore[override] Positive scalar weight eta : :obj:`float`, optional Relaxation parameter (must be between 0 and 2, 0 excluded). - niter : :obj:`int`, optional - Number of iterations of iterative scheme gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -3599,8 +3652,8 @@ def setup( # type: ignore[override] self.proxg = proxg self.tau = tau self.eta = eta - self.niter = niter self.gfirst = gfirst + self.niter = niter self.tol = tol self.callbacky = callbacky @@ -3622,7 +3675,10 @@ def setup( # type: ignore[override] return x, y def step( - self, x: NDArray, y: NDArray, show: bool = False + self, + x: NDArray, + y: NDArray, + show: bool = False, ) -> tuple[NDArray, NDArray]: r"""Run one step of solver @@ -3741,8 +3797,8 @@ def solve( # type: ignore[override] x0: NDArray, tau: float, eta: float = 1.0, - niter: int = 10, gfirst: bool = True, + niter: int = 10, tol: float | None = None, callbacky: bool = False, show: bool = False, @@ -3762,11 +3818,11 @@ def solve( # type: ignore[override] Positive scalar weight eta : :obj:`float`, optional Relaxation parameter (must be between 0 and 2, 0 excluded). - niter : :obj:`int`, optional - Number of iterations of iterative scheme gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + niter : :obj:`int`, optional + Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -3798,8 +3854,8 @@ def solve( # type: ignore[override] x0=x0, tau=tau, eta=eta, - niter=niter, gfirst=gfirst, + niter=niter, tol=tol, callbacky=callbacky, show=show, @@ -4009,6 +4065,10 @@ def step( Updated additional model vector """ + # store current solution prior to update + # for tolerance check + self.xold = x + # proximal steps p = self.ncp.stack( [self.proxfs[i].prox(y[i], self.tau / self.w[i]) for i in range(self.nprox)] @@ -4306,7 +4366,11 @@ def setup( # type: ignore[override] return x, x_bar, y def step( - self, x: NDArray, x_bar: NDArray, y: NDArray, show: bool = False + self, + x: NDArray, + x_bar: NDArray, + y: NDArray, + show: bool = False, ) -> tuple[NDArray, NDArray, NDArray]: r"""Run one step of solver @@ -4334,6 +4398,10 @@ def step( Updated additional model vector """ + # store current solution prior to update + # for tolerance check + self.xold = x + # proximal steps x = self.ncp.stack( [self.proxfs[i].prox(x_bar - y[i], self.tau) for i in range(self.nprox)] From 355d62686d84dff2093dbd42c5798537d426c58e Mon Sep 17 00:00:00 2001 From: mrava87 Date: Fri, 26 Jun 2026 10:52:31 +0100 Subject: [PATCH 21/24] feat: switched to calling class-based solvers in functional api --- pyproximal/optimization/primal.py | 1769 +++++++++-------------------- 1 file changed, 544 insertions(+), 1225 deletions(-) diff --git a/pyproximal/optimization/primal.py b/pyproximal/optimization/primal.py index 12ad1e4..342bfc7 100644 --- a/pyproximal/optimization/primal.py +++ b/pyproximal/optimization/primal.py @@ -1,116 +1,61 @@ -import time +__all__ = [ + "ProximalPoint", + "ProximalGradient", + "AcceleratedProximalGradient", + "AndersonProximalGradient", + "GeneralizedProximalGradient", + "HQS", + "ADMM", + "ADMML2", + "LinearizedADMM", + "TwIST", + "DouglasRachfordSplitting", + "PPXA", + "ConsensusADMM", +] + import warnings from collections.abc import Callable -from math import sqrt -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any -import numpy as np -from pylops.optimization.leastsquares import regularized_inversion -from pylops.utils.backend import get_array_module, to_numpy +from pylops.optimization.callback import CostNanInfCallback, CostToInitialCallback from pylops.utils.typing import NDArray -from pyproximal.proximal import L2 -from pyproximal.ProxOperator import ProxOperator -from pyproximal.utils.bilinear import BilinearOperator +from pyproximal.optimization.cls_primal import ADMM as cADMM +from pyproximal.optimization.cls_primal import ADMML2 as cADMML2 +from pyproximal.optimization.cls_primal import HQS as cHQS +from pyproximal.optimization.cls_primal import PPXA as cPPXA +from pyproximal.optimization.cls_primal import ( + AndersonProximalGradient as cAndersonProximalGradient, +) +from pyproximal.optimization.cls_primal import ConsensusADMM as cConsensusADMM +from pyproximal.optimization.cls_primal import ( + DouglasRachfordSplitting as cDouglasRachfordSplitting, +) +from pyproximal.optimization.cls_primal import ( + GeneralizedProximalGradient as cGeneralizedProximalGradient, +) +from pyproximal.optimization.cls_primal import LinearizedADMM as cLinearizedADMM +from pyproximal.optimization.cls_primal import ProximalGradient as cProximalGradient +from pyproximal.optimization.cls_primal import ProximalPoint as cProximalPoint +from pyproximal.optimization.cls_primal import TwIST as cTwIST if TYPE_CHECKING: from pylops.linearoperator import LinearOperator - -def _backtracking( - x: NDArray, - tau: float, - proxf: ProxOperator, - proxg: ProxOperator, - epsg: float, - beta: float = 0.5, - niterback: int = 10, -) -> tuple[NDArray, float]: - r"""Backtracking - - Line-search algorithm for finding step sizes in proximal algorithms when - the Lipschitz constant of the operator is unknown (or expensive to - estimate). - - """ - - def ftilde(x: NDArray, y: NDArray, f: ProxOperator, tau: float) -> float: - xy = x - y - return float( - f(y) + np.dot(f.grad(y), xy) + (1.0 / (2.0 * tau)) * np.linalg.norm(xy) ** 2 - ) - - iiterback = 0 - while iiterback < niterback: - z = proxg.prox(x - tau * proxf.grad(x), epsg * tau) - ft = ftilde(z, x, proxf, tau) - if proxf(z) <= ft: - break - tau *= beta - iiterback += 1 - return z, tau - - -def _x0z0_init( - x0: NDArray | None, - z0: NDArray | None, - Op: Optional["LinearOperator"] = None, - z0name: str | None = "z0", - Opname: str | None = "Op", -) -> tuple[NDArray, NDArray]: - r"""Initialize x0 and z0 - - Initialize x0 and z0 using the following convention. - - For ``Op=None``: - - if both are provided, they are simply returned; - - if only one is provided (the other is ``None``), the one provided - is copied to the other one. - - For ``Op!=None``, ``x0`` must be provided, and: - - if both are provided, they are simply returned; - - if ``z0`` is not provided, set to ``Op @ x0``. - - Parameters - ---------- - x0 : :obj:`numpy.ndarray` - Initial vector - z0 : :obj:`numpy.ndarray` - Initial auxiliary vector - Op : :obj:`pylops.LinearOperator`, optional - Linear Operator to apply to ``x0`` - z0name : :obj:`str`, optional - Name to display in error message instead of ``z0`` - Opname : :obj:`str`, optional - Name to display in error message instead of ``Op`` - - """ - if x0 is None and z0 is None: - msg = f"Both x0 or {z0name} are None, provide either of them or both" - raise ValueError(msg) - - if Op is None: - if x0 is None: - x0 = z0.copy() # type: ignore[union-attr] - elif z0 is None: - z0 = x0.copy() - else: - if x0 is None: - msg = f"x0 must be provided when {Opname} is also provided" - raise ValueError(msg) - elif z0 is None: - z0 = Op @ x0 - return x0, z0 + from pyproximal.ProxOperator import ProxOperator def ProximalPoint( - prox: ProxOperator, + prox: "ProxOperator", x0: NDArray, tau: float, niter: int = 10, tol: float | None = None, + rtol: float | None = None, callback: Callable[[NDArray], None] | None = None, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> NDArray: r"""Proximal point algorithm @@ -135,12 +80,23 @@ def ProximalPoint( Number of iterations of iterative scheme tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If - ``tol=None``, run until ``niter`` is reached + ``tol=None``, run until ``niter`` is reached or the other tolerance + criterion is met + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -149,66 +105,35 @@ def ProximalPoint( Notes ----- - The Proximal point algorithm can be expressed by the following recursion: - - .. math:: - - \mathbf{x}^{k+1} = \prox_{\tau f}(\mathbf{x}^k) + See :class:`pyproximal.optimization.cls_primal.ProximalPoint` """ - if show: - tstart = time.time() - print( - "Proximal point algorithm\n" - "---------------------------------------------------------\n" - "Proximal operator: %s\n" - "tau = %10e\tniter = %d\ttol = %s\n" % (type(prox), tau, niter, str(tol)) - ) - head = " Itn x[0] f" - print(head) - - # initialize model - x = x0.copy() - pf = np.inf - tolbreak = False - - # iterate - for iiter in range(niter): - x = prox.prox(x, tau) - - # run callback - if callback is not None: - callback(x) - - # tolerance check: break iterations if overall - # objective does not decrease below tolerance - if tol is not None: - pfold = pf - pf = prox(x) - if np.abs(1.0 - pf / pfold) < tol: - tolbreak = True - - # show iteration logger - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - if tol is None: - pf = prox(x) - msg = "%6g %12.5e %10.3e" % (iiter + 1, x[0], pf) - print(msg) - - # break if tolerance condition is met - if tolbreak: - break - - if show: - print("\nTotal time (s) = %.2f" % (time.time() - tstart)) - print("---------------------------------------------------------\n") + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + proxpsolve = cProximalPoint( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + proxpsolve.callback = callback + x, _, _ = proxpsolve.solve( + prox=prox, + x0=x0, + tau=tau, + niter=niter, + tol=tol or (0.0 if rtol else None), + show=show, + itershow=itershow, + ) return x def ProximalGradient( - proxf: ProxOperator, - proxg: ProxOperator, + proxf: "ProxOperator", + proxg: "ProxOperator", x0: NDArray, epsg: float | NDArray = 1.0, tau: float | None = None, @@ -219,8 +144,10 @@ def ProximalGradient( niterback: int = 100, acceleration: str | None = None, tol: float | None = None, + rtol: float | None = None, callback: Callable[[NDArray], None] | None = None, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> NDArray: r"""Proximal gradient (optionally accelerated) @@ -268,12 +195,23 @@ def ProximalGradient( Acceleration (``None``, ``vandenberghe`` or ``fista``) tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If - ``tol=None``, run until ``niter`` is reached + ``tol=None``, run until ``niter`` is reached or the other tolerance + criterion is met + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -282,177 +220,42 @@ def ProximalGradient( Notes ----- - The Proximal gradient algorithm can be expressed by the following recursion: - - .. math:: - - \mathbf{x}^{k+1} = \mathbf{y}^k + \eta (\prox_{\tau^k \epsilon g}(\mathbf{y}^k - - \tau^k \nabla f(\mathbf{y}^k)) - \mathbf{y}^k) \\ - \mathbf{y}^{k+1} = \mathbf{x}^k + \omega^k - (\mathbf{x}^k - \mathbf{x}^{k-1}) - - where at each iteration :math:`\tau^k` can be estimated by back-tracking - as follows: - - .. math:: - - \begin{aligned} - &\tau = \tau^{k-1} &\\ - &repeat \; \mathbf{z} = \prox_{\tau \epsilon g}(\mathbf{x}^k - - \tau \nabla f(\mathbf{x}^k)), \tau = \beta \tau \quad if \; - f(\mathbf{z}) \leq \tilde{f}_\tau(\mathbf{z}, \mathbf{x}^k) \\ - &\tau^k = \tau, \quad \mathbf{x}^{k+1} = \mathbf{z} &\\ - \end{aligned} - - where :math:`\tilde{f}_\tau(\mathbf{x}, \mathbf{y}) = f(\mathbf{y}) + - \nabla f(\mathbf{y})^T (\mathbf{x} - \mathbf{y}) + - 1/(2\tau)||\mathbf{x} - \mathbf{y}||_2^2`. - - Different accelerations are provided: - - - ``acceleration=None``: :math:`\omega^k = 0`; - - ``acceleration=vandenberghe`` [1]_: :math:`\omega^k = k / (k + 3)` for ` - - ``acceleration=fista``: :math:`\omega^k = (t_{k-1}-1)/t_k` where - :math:`t_k = (1 + \sqrt{1+4t_{k-1}^{2}}) / 2` [2]_ - - .. [1] Vandenberghe, L., "Fast proximal gradient methods", 2010. - .. [2] Beck, A., and Teboulle, M. "A Fast Iterative Shrinkage-Thresholding - Algorithm for Linear Inverse Problems", SIAM Journal on - Imaging Sciences, vol. 2, pp. 183-202. 2009. + See :class:`pyproximal.optimization.cls_primal.ProximalGradient` """ - # check if epgs is a vector - epsg = np.asarray(epsg, dtype=float) - if epsg.size == 1: - epsg = epsg * np.ones(niter) - epsg_print = str(epsg[0]) - else: - epsg_print = "Multi" - - if acceleration not in [None, "None", "vandenberghe", "fista"]: - msg = "Acceleration should be None, vandenberghe or fista" - raise NotImplementedError(msg) - - if show: - tstart = time.time() - print( - "Accelerated Proximal Gradient\n" - "---------------------------------------------------------\n" - "Proximal operator (f): %s\n" - "Proximal operator (g): %s\n" - "tau = %s\tbacktrack = %s\tbeta = %10e\n" - "epsg = %s\tniter = %d\ttol = %s\n" - "" - "niterback = %d\tacceleration = %s\n" - % ( - type(proxf), - type(proxg), - str(tau), - backtracking, - beta, - epsg_print, - niter, - str(tol), - niterback, - acceleration, - ) - ) - head = " Itn x[0] f g J=f+eps*g tau" - print(head) - - if tau is None: - backtracking = True - tau = 1.0 - - # initialize model - t = 1.0 - x = x0.copy() - y = x.copy() - pfg = np.inf - tolbreak = False - - # iterate - for iiter in range(niter): - xold = x.copy() - - # proximal step - if not backtracking: - if eta == 1.0: - x = proxg.prox(y - tau * proxf.grad(y), epsg[iiter] * tau) - else: - x = x + eta * ( - proxg.prox(x - tau * proxf.grad(x), epsg[iiter] * tau) - x - ) - else: - x, tau = _backtracking( - y, tau, proxf, proxg, epsg[iiter], beta=beta, niterback=niterback - ) - if eta != 1.0: - x = x + eta * ( - proxg.prox(x - tau * proxf.grad(x), epsg[iiter] * tau) - x - ) - - # update internal parameters for bilinear operator - if isinstance(proxf, BilinearOperator): - proxf.updatexy(x) - - # update y - if acceleration == "vandenberghe": - omega = iiter / (iiter + 3) - elif acceleration == "fista": - told = t - t = (1.0 + np.sqrt(1.0 + 4.0 * t**2)) / 2.0 - omega = (told - 1.0) / t - else: - omega = 0 - y = x + omega * (x - xold) - - # run callback - if callback is not None: - callback(x) - - # tolerance check: break iterations if overall - # objective does not decrease below tolerance - if tol is not None: - pfgold = pfg - pf, pg = proxf(x), proxg(x) - pfg = pf + np.sum(epsg[iiter] * pg) - if np.abs(1.0 - pfg / pfgold) < tol: - tolbreak = True - - # show iteration logger - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - if tol is None: - pf, pg = proxf(x), proxg(x) - pfg = pf + np.sum(epsg[iiter] * pg) - msg = "%6g %12.5e %10.3e %10.3e %10.3e %10.3e" % ( - iiter + 1, - ( - np.real(to_numpy(x[0])) - if x.ndim == 1 - else np.real(to_numpy(x[0, 0])) - ), - pf, - pg, - pfg, - tau, - ) - print(msg) - - # break if tolerance condition is met - if tolbreak: - break - - if show: - print("\nTotal time (s) = %.2f" % (time.time() - tstart)) - print("---------------------------------------------------------\n") + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + proxgsolve = cProximalGradient( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + proxgsolve.callback = callback + x, _, _, _ = proxgsolve.solve( + proxf=proxf, + proxg=proxg, + x0=x0, + epsg=epsg, + tau=tau, + backtracking=backtracking, + beta=beta, + eta=eta, + acceleration=acceleration, + niterback=niterback, + niter=niter, + tol=tol or (0.0 if rtol else None), + show=show, + itershow=itershow, + ) return x def AcceleratedProximalGradient( - proxf: ProxOperator, - proxg: ProxOperator, + proxf: "ProxOperator", + proxg: "ProxOperator", x0: NDArray, tau: float | None = None, beta: float = 0.5, @@ -461,8 +264,10 @@ def AcceleratedProximalGradient( niterback: int = 100, acceleration: str = "vandenberghe", tol: float | None = None, + rtol: float | None = None, callback: Callable[[NDArray], None] | None = None, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> NDArray: r"""Accelerated Proximal gradient @@ -490,14 +295,16 @@ def AcceleratedProximalGradient( niterback=niterback, acceleration=acceleration, tol=tol, + rtol=rtol, callback=callback, show=show, + itershow=itershow, ) def AndersonProximalGradient( - proxf: ProxOperator, - proxg: ProxOperator, + proxf: "ProxOperator", + proxg: "ProxOperator", x0: NDArray, epsg: float | NDArray = 1.0, tau: float | NDArray = 1.0, @@ -506,8 +313,10 @@ def AndersonProximalGradient( epsr: float = 1e-10, safeguard: bool = False, tol: float | None = None, + rtol: float | None = None, callback: Callable[[NDArray], None] | None = None, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> NDArray: r"""Proximal gradient with Anderson acceleration @@ -549,11 +358,21 @@ def AndersonProximalGradient( tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -562,170 +381,40 @@ def AndersonProximalGradient( Notes ----- - The Proximal gradient algorithm with Anderson acceleration can be expressed by the - following recursion [1]_: - - .. math:: - m_k = min(m, k)\\ - \mathbf{g}^{k} = \mathbf{x}^{k} - \tau^k \nabla f(\mathbf{x}^k)\\ - \mathbf{r}^{k} = \mathbf{g}^{k} - \mathbf{g}^{k}\\ - \mathbf{G}^{k} = [\mathbf{g}^{k},..., \mathbf{g}^{k-m_k}]\\ - \mathbf{R}^{k} = [\mathbf{r}^{k},..., \mathbf{r}^{k-m_k}]\\ - \alpha_k = (\mathbf{R}^{kT} \mathbf{R}^{k})^{-1} \mathbf{1} / \mathbf{1}^T - (\mathbf{R}^{kT} \mathbf{R}^{k})^{-1} \mathbf{1}\\ - \mathbf{y}^{k+1} = \mathbf{G}^{k} \alpha_k\\ - \mathbf{x}^{k+1} = \prox_{\tau^{k+1} g}(\mathbf{y}^{k+1}) - - where :math:`m` equals ``nhistory``, :math:`k=1,2,...,n_{iter}`, :math:`\mathbf{y}^{0}=\mathbf{x}^{0}`, - :math:`\mathbf{y}^{1}=\mathbf{x}^{0} - \tau^0 \nabla f(\mathbf{x}^0)`, - :math:`\mathbf{x}^{1}=\prox_{\tau^k g}(\mathbf{y}^{1})`, and - :math:`\mathbf{g}^{0}=\mathbf{y}^{1}`. - - Refer to [1]_ for the guarded version of the algorithm (when ``safeguard=True``). - - .. [1] Mai, V., and Johansson, M. "Anderson Acceleration of Proximal Gradient - Methods", 2020. + See :class:`pyproximal.optimization.cls_primal.AndersonProximalGradient` """ - # check if epgs is a vector - epsg = np.asarray(epsg, dtype=float) - if epsg.size == 1: - epsg = epsg * np.ones(niter) - epsg_print = str(epsg[0]) - else: - epsg_print = "Multi" - - if show: - tstart = time.time() - print( - "Proximal Gradient with Anderson Acceleration \n" - "---------------------------------------------------------\n" - "Proximal operator (f): %s\n" - "Proximal operator (g): %s\n" - "tau = %s\t\tepsg = %s\tniter = %d\n" - "nhist = %d\tepsr = %s\n" - "guard = %s\ttol = %s\n" - % ( - type(proxf), - type(proxg), - str(tau), - epsg_print, - niter, - nhistory, - str(epsr), - str(safeguard), - str(tol), - ) - ) - head = " Itn x[0] f g J=f+eps*g tau" - print(head) - - # initialize model - y = x0 - tau * proxf.grad(x0) - x = proxg.prox(y, epsg[0] * tau) - g = y.copy() - r = g - x0 - R, G = ( - [ - g, - ], - [ - r, - ], + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + aproxgsolve = cAndersonProximalGradient( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + aproxgsolve.callback = callback + x, _, _, _ = aproxgsolve.solve( + proxf=proxf, + proxg=proxg, + x0=x0, + epsg=epsg, + tau=tau, + niter=niter, + nhistory=nhistory, + epsr=epsr, + safeguard=safeguard, + tol=tol or (0.0 if rtol else None), + show=show, + itershow=itershow, ) - pf = proxf(x) - pfg = np.inf - tolbreak = False - - # iterate - for iiter in range(niter): - # update fix point - g = x - tau * proxf.grad(x) - r = g - y - - # update history vectors - R.insert(0, r) - G.insert(0, g) - if iiter >= nhistory - 1: - R.pop(-1) - G.pop(-1) - - # solve for alpha coefficients - Rstack = np.vstack(R) - Rinv = np.linalg.pinv(Rstack @ Rstack.T + epsr * np.linalg.norm(Rstack) ** 2) - ones = np.ones(min(nhistory, iiter + 2)) - Rinvones = Rinv @ ones - alpha = Rinvones / (ones[None] @ Rinvones) - - if not safeguard: - # update auxiliary variable - y = np.vstack(G).T @ alpha - - # update main variable - x = proxg.prox(y, epsg[iiter] * tau) - - else: - # update auxiliary variable - ytest = np.vstack(G).T @ alpha - - # update main variable - xtest = proxg.prox(ytest, epsg[iiter] * tau) - - # check if function is decreased, otherwise do basic PG step - pfold, pf = pf, proxf(xtest) - if pf <= pfold - tau * np.linalg.norm(proxf.grad(x)) ** 2 / 2: - y = ytest - x = xtest - else: - x = proxg.prox(g, epsg[iiter] * tau) - y = g - - # run callback - if callback is not None: - callback(x) - - # tolerance check: break iterations if overall - # objective does not decrease below tolerance - if tol is not None: - pfgold = pfg - pf, pg = proxf(x), proxg(x) - pfg = pf + np.sum(epsg[iiter] * pg) - if np.abs(1.0 - pfg / pfgold) < tol: - tolbreak = True - - # show iteration logger - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - if tol is None: - pf, pg = proxf(x), proxg(x) - pfg = pf + np.sum(epsg[iiter] * pg) - msg = "%6g %12.5e %10.3e %10.3e %10.3e %10.3e" % ( - iiter + 1, - ( - np.real(to_numpy(x[0])) - if x.ndim == 1 - else np.real(to_numpy(x[0, 0])) - ), - pf, - pg, - pfg, - tau, - ) - print(msg) - - # break if tolerance condition is met - if tolbreak: - break - - if show: - print("\nTotal time (s) = %.2f" % (time.time() - tstart)) - print("---------------------------------------------------------\n") return x def GeneralizedProximalGradient( - proxfs: list[ProxOperator], - proxgs: list[ProxOperator], + proxfs: list["ProxOperator"], + proxgs: list["ProxOperator"], x0: NDArray, tau: float | None, epsg: float | NDArray = 1.0, @@ -733,8 +422,11 @@ def GeneralizedProximalGradient( eta: float = 1.0, niter: int = 10, acceleration: str | None = None, + tol: float | None = None, + rtol: float | None = None, callback: Callable[[NDArray], None] | None = None, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> NDArray: r"""Generalized Proximal gradient @@ -773,11 +465,25 @@ def GeneralizedProximalGradient( Number of iterations of iterative scheme acceleration: :obj:`str`, optional Acceleration (``None``, ``vandenberghe`` or ``fista``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached or the other tolerance + criterion is met + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -786,134 +492,51 @@ def GeneralizedProximalGradient( Notes ----- - The Generalized Proximal gradient algorithm can be expressed by the - following recursion [1]_: - - .. math:: - \text{for } j=1,\cdots,n, \\ - ~~~~\mathbf z_j^{k+1} = \mathbf z_j^{k} + \eta - \left[prox_{\frac{\tau^k \epsilon_j}{w_j} g_j}\left(2 \mathbf{x}^{k} - \mathbf{z}_j^{k} - - \tau^k \sum_{i=1}^n \nabla f_i(\mathbf{x}^{k})\right) - \mathbf{x}^{k} \right] \\ - \mathbf{x}^{k+1} = \sum_{j=1}^n w_j \mathbf z_j^{k+1} \\ - - where :math:`\sum_{j=1}^n w_j=1`. In the current implementation, :math:`w_j=1/n` when - not provided. - - .. [1] Raguet, H., Fadili, J. and Peyré, G. "Generalized Forward-Backward Splitting", - arXiv, 2012. + See :class:`pyproximal.optimization.cls_primal.GeneralizedProximalGradient` """ - # check if weights sum to 1 - if weights is None: - weights = np.ones(len(proxgs)) / len(proxgs) - if len(weights) != len(proxgs) or np.sum(weights) != 1.0: - msg = f"omega={weights} must be an array of size {len(proxgs)} summing to 1" - raise ValueError(msg) - - # check if epgs is a vector - epsg = np.asarray(epsg, dtype=float) - if epsg.size == 1: - epsg_print = str(epsg) - epsg = epsg * np.ones(len(proxgs)) - else: - epsg_print = "Multi" - - if acceleration not in [None, "None", "vandenberghe", "fista"]: - msg = "Acceleration should be None, vandenberghe or fista" - raise NotImplementedError(msg) - if show: - tstart = time.time() - print( - "Generalized Proximal Gradient\n" - "---------------------------------------------------------\n" - "Proximal operators (f): %s\n" - "Proximal operators (g): %s\n" - "tau = %10e\nepsg = %s\tniter = %d\n" - % ( - [type(proxf) for proxf in proxfs], - [type(proxg) for proxg in proxgs], - 0 if tau is None else tau, - epsg_print, - niter, - ) - ) - head = " Itn x[0] f g J=f+g" - print(head) - - if tau is None: - tau = 1.0 - - # initialize model - t = 1.0 - x = x0.copy() - y = x.copy() - zs = [x.copy() for _ in range(len(proxgs))] - - # iterate - for iiter in range(niter): - xold = x.copy() - - # gradient - grad = np.zeros_like(x) - for _, proxf in enumerate(proxfs): - grad += proxf.grad(x) - - # proximal step - x = np.zeros_like(x) - for i, proxg in enumerate(proxgs): - ztmp = 2 * y - zs[i] - tau * grad - ztmp = proxg.prox(ztmp, tau * epsg[i] / weights[i]) - zs[i] += eta * (ztmp - y) - x += weights[i] * zs[i] - - # update y - if acceleration == "vandenberghe": - omega = iiter / (iiter + 3) - elif acceleration == "fista": - told = t - t = (1.0 + np.sqrt(1.0 + 4.0 * t**2)) / 2.0 - omega = (told - 1.0) / t - else: - omega = 0 - y = x + omega * (x - xold) - - # run callback - if callback is not None: - callback(x) - - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf: float = np.sum([proxf(x) for proxf in proxfs]) - pg: float = np.sum( - [eg * proxg(x) for proxg, eg in zip(proxgs, epsg, strict=True)] - ) - msg = "%6g %12.5e %10.3e %10.3e %10.3e" % ( - iiter + 1, - np.real(to_numpy(x[0])) - if x.ndim == 1 - else np.real(to_numpy(x[0, 0])), - pf, - pg, - pf + pg, - ) - print(msg) - if show: - print("\nTotal time (s) = %.2f" % (time.time() - tstart)) - print("---------------------------------------------------------\n") + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + proxgsolve = cGeneralizedProximalGradient( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + proxgsolve.callback = callback + x, _, _, _ = proxgsolve.solve( + proxfs=proxfs, + proxgs=proxgs, + x0=x0, + epsg=epsg, + weights=weights, + tau=1.0 if tau is None else tau, + eta=eta, + niter=niter, + acceleration=acceleration, + tol=tol or (0.0 if rtol else None), + show=show, + itershow=itershow, + ) return x def HQS( - proxf: ProxOperator, - proxg: ProxOperator, + proxf: "ProxOperator", + proxg: "ProxOperator", x0: NDArray, tau: float | NDArray, niter: int = 10, z0: NDArray | None = None, gfirst: bool = True, + tol: float | None = None, + rtol: float | None = None, callback: Callable[..., None] | None = None, callbackz: bool = False, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> tuple[NDArray, NDArray]: r"""Half Quadratic splitting @@ -951,6 +574,15 @@ def HQS( gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector @@ -958,6 +590,10 @@ def HQS( Modify callback signature to (``callback(x, z)``) when ``callbackz=True`` show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -973,98 +609,50 @@ def HQS( Notes ----- - The HQS algorithm can be expressed by the following recursion [1]_: - - .. math:: - - \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{x}^{k}) \\ - \mathbf{x}^{k+1} = \prox_{\tau f}(\mathbf{z}^{k+1}) - - for ``gfirst=False``, or - - .. math:: - - \mathbf{x}^{k+1} = \prox_{\tau f}(\mathbf{z}^{k}) \\ - \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{x}^{k+1}) - - for ``gfirst=False``. Note that ``x`` and ``z`` converge to each other, - however if iterations are stopped too early ``x`` is guaranteed to belong to - the domain of ``f`` while ``z`` is guaranteed to belong to the domain of ``g``. - Depending on the problem either of the two may be the best solution. - - .. [1] D., Geman, and C., Yang, "Nonlinear image recovery with halfquadratic - regularization", IEEE Transactions on Image Processing, - 4, 7, pp. 932-946, 1995. + See :class:`pyproximal.optimization.cls_primal.HQS` """ - # initialize variables - x, z = _x0z0_init(x0, z0) - ncp = get_array_module(x) - - # check if tau is a vector - tau = ncp.asarray(tau, dtype=float) - if tau.size == 1: - tau_print = str(tau) - tau = tau * np.ones(niter) - else: - tau_print = "Variable" - - if show: - tstart = time.time() - print( - "HQS\n" - "---------------------------------------------------------\n" - "Proximal operator (f): %s\n" - "Proximal operator (g): %s\n" - "tau = %s\tniter = %d\n" % (type(proxf), type(proxg), tau_print, niter) - ) - head = " Itn x[0] f g J = f + g" - print(head) - - # run iterations - for iiter in range(niter): - if gfirst: - z = proxg.prox(x, tau[iiter]) - x = proxf.prox(z, tau[iiter]) - else: - x = proxf.prox(z, tau[iiter]) - z = proxg.prox(x, tau[iiter]) - - # run callback - if callback is not None: - if callbackz: - callback(x, z) - else: - callback(x) - - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf, pg = proxf(x), proxg(x) - msg = "%6g %12.5e %10.3e %10.3e %10.3e" % ( - iiter + 1, - np.real(to_numpy(x[0])), - pf, - pg, - pf + pg, - ) - print(msg) - if show: - print("\nTotal time (s) = %.2f" % (time.time() - tstart)) - print("---------------------------------------------------------\n") + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + hqssolve = cHQS( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + hqssolve.callback = callback + x, z, _, _ = hqssolve.solve( + proxf=proxf, + proxg=proxg, + x0=x0, + tau=tau, + z0=z0, + niter=niter, + gfirst=gfirst, + tol=tol or (0.0 if rtol else None), + callbackz=callbackz, + show=show, + itershow=itershow, + ) return x, z def ADMM( - proxf: ProxOperator, - proxg: ProxOperator, + proxf: "ProxOperator", + proxg: "ProxOperator", x0: NDArray, tau: float, niter: int = 10, z0: NDArray | None = None, gfirst: bool = False, + tol: float | None = None, + rtol: float | None = None, callback: Callable[..., None] | None = None, callbackz: bool = False, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> tuple[NDArray, NDArray]: r"""Alternating Direction Method of Multipliers @@ -1115,6 +703,15 @@ def ADMM( gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector @@ -1122,6 +719,10 @@ def ADMM( Modify callback signature to (``callback(x, z)``) when ``callbackz=True`` show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -1142,77 +743,38 @@ def ADMM( Notes ----- - The ADMM algorithm can be expressed by the following recursion [1]_: - - .. math:: - - \mathbf{x}^{k+1} = \prox_{\tau f}(\mathbf{z}^{k} - \mathbf{u}^{k})\\ - \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{x}^{k+1} + \mathbf{u}^{k})\\ - \mathbf{u}^{k+1} = \mathbf{u}^{k} + \mathbf{x}^{k+1} - \mathbf{z}^{k+1} - - Note that ``x`` and ``z`` converge to each other, however if iterations are - stopped too early ``x`` is guaranteed to belong to the domain of ``f`` - while ``z`` is guaranteed to belong to the domain of ``g``. Depending on - the problem either of the two may be the best solution. - - .. [1] S. Boyd, N. Parikh, E. Chu, B. Peleato, and J. Eckstein. 2011. - Distributed optimization and statistical learning via the alternating - direction method of multipliers. Foundations and Trends in Machine - Learning, 3 (1), 1-122. https://doi.org/10.1561/2200000016. + See :class:`pyproximal.optimization.cls_primal.ADMM` """ - # initialize variables - x, z = _x0z0_init(x0, z0) - ncp = get_array_module(x) - u = ncp.zeros_like(x) - - if show: - tstart = time.time() - print( - "ADMM\n" - "---------------------------------------------------------\n" - "Proximal operator (f): %s\n" - "Proximal operator (g): %s\n" - "tau = %10e\tniter = %d\n" % (type(proxf), type(proxg), tau, niter) - ) - head = " Itn x[0] f g J = f + g" - print(head) - - # run iterations - for iiter in range(niter): - if gfirst: - z = proxg.prox(x + u, tau) - x = proxf.prox(z - u, tau) - else: - x = proxf.prox(z - u, tau) - z = proxg.prox(x + u, tau) - u = u + x - z - - # run callback - if callback is not None: - if callbackz: - callback(x, z) - else: - callback(x) - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf, pg = proxf(x), proxg(x) - msg = "%6g %12.5e %10.3e %10.3e %10.3e" % ( - iiter + 1, - np.real(to_numpy(x[0])), - pf, - pg, - pf + pg, - ) - print(msg) - if show: - print("\nTotal time (s) = %.2f" % (time.time() - tstart)) - print("---------------------------------------------------------\n") + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + admmsolve = cADMM( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + admmsolve.callback = callback + x, z, _, _ = admmsolve.solve( + proxf=proxf, + proxg=proxg, + x0=x0, + tau=tau, + z0=z0, + gfirst=gfirst, + niter=niter, + tol=tol or (0.0 if rtol else None), + callbackz=callbackz, + show=show, + itershow=itershow, + ) return x, z def ADMML2( - proxg: ProxOperator, + proxg: "ProxOperator", Op: "LinearOperator", b: NDArray, A: "LinearOperator", @@ -1221,8 +783,12 @@ def ADMML2( niter: int = 10, z0: NDArray | None = None, gfirst: bool = False, + tol: float | None = None, + rtol: float | None = None, callback: Callable[[NDArray], None] | None = None, + callbackz: bool = False, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), **kwargs_solver: dict[str, Any], ) -> tuple[NDArray, NDArray]: r"""Alternating Direction Method of Multipliers for L2 misfit term @@ -1260,11 +826,24 @@ def ADMML2( gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. **kwargs_solver Arbitrary keyword arguments for :py:func:`scipy.sparse.linalg.lsqr` used to solve the x-update @@ -1288,108 +867,53 @@ def ADMML2( Notes ----- - The ADMM algorithm can be expressed by the following recursion: - - .. math:: - - \mathbf{x}^{k+1} = \argmin_{\mathbf{x}} \frac{1}{2}||\mathbf{Op}\mathbf{x} - - \mathbf{b}||_2^2 + \frac{1}{2\tau} ||\mathbf{Ax} - \mathbf{z}^k + \mathbf{u}^k||_2^2\\ - \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{Ax}^{k+1} + \mathbf{u}^{k})\\ - \mathbf{u}^{k+1} = \mathbf{u}^{k} + \mathbf{Ax}^{k+1} - \mathbf{z}^{k+1} + See :class:`pyproximal.optimization.cls_primal.ADMML2` """ - # initialize variables - x, z = _x0z0_init(x0, z0, A, Opname="A") - ncp = get_array_module(x) - u = ncp.zeros_like(z) - - if show: - tstart = time.time() - print( - "ADMM\n" - "---------------------------------------------------------\n" - "Proximal operator (g): %s\n" - "tau = %10e\tniter = %d\n" % (type(proxg), tau, niter) - ) - head = " Itn x[0] f g J = f + g" - print(head) - - # run iterations - sqrttau = 1.0 / sqrt(tau) - for iiter in range(niter): - if gfirst: - Ax = A @ x - z = proxg.prox(Ax + u, tau) - - # solve augumented system - x = regularized_inversion( - Op, - b, - [ - A, - ], - x0=x, - dataregs=[ - z - u, - ], - epsRs=[ - sqrttau, - ], - **kwargs_solver, - )[0] - else: - # solve augumented system - x = regularized_inversion( - Op, - b, - [ - A, - ], - x0=x, - dataregs=[ - z - u, - ], - epsRs=[ - sqrttau, - ], - **kwargs_solver, - )[0] - Ax = A @ x - z = proxg.prox(Ax + u, tau) - u = u + Ax - z - - # run callback - if callback is not None: - callback(x) - - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf, pg = 0.5 * np.linalg.norm(Op @ x - b) ** 2, proxg(Ax) - msg = "%6g %12.5e %10.3e %10.3e %10.3e" % ( - iiter + 1, - np.real(to_numpy(x[0])), - pf, - pg, - pf + pg, - ) - print(msg) - if show: - print("\nTotal time (s) = %.2f" % (time.time() - tstart)) - print("---------------------------------------------------------\n") + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + admml2solve = cADMML2( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + admml2solve.callback = callback + x, z, _, _ = admml2solve.solve( + proxg=proxg, + Op=Op, + b=b, + A=A, + x0=x0, + tau=tau, + z0=z0, + gfirst=gfirst, + niter=niter, + tol=tol or (0.0 if rtol else None), + callbackz=callbackz, + show=show, + itershow=itershow, + **kwargs_solver, + ) return x, z def LinearizedADMM( - proxf: ProxOperator, - proxg: ProxOperator, + proxf: "ProxOperator", + proxg: "ProxOperator", A: "LinearOperator", x0: NDArray, tau: float, mu: float, niter: int = 10, z0: NDArray | None = None, + tol: float | None = None, + rtol: float | None = None, callback: Callable[[NDArray], None] | None = None, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> tuple[NDArray, NDArray]: r"""Linearized Alternating Direction Method of Multipliers @@ -1426,11 +950,24 @@ def LinearizedADMM( Number of iterations of iterative scheme z0 : :obj:`numpy.ndarray` Initial auxiliary vector. If ``None``, initialized to ``A @ x0``. + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -1451,71 +988,38 @@ def LinearizedADMM( Notes ----- - The Linearized-ADMM algorithm can be expressed by the following recursion [1]_: - - .. math:: - - \mathbf{x}^{k+1} = \prox_{\mu f}(\mathbf{x}^{k} - \frac{\mu}{\tau} - \mathbf{A}^H(\mathbf{A} \mathbf{x}^k - \mathbf{z}^k + \mathbf{u}^k))\\ - \mathbf{z}^{k+1} = \prox_{\tau g}(\mathbf{A} \mathbf{x}^{k+1} + - \mathbf{u}^k)\\ - \mathbf{u}^{k+1} = \mathbf{u}^{k} + \mathbf{A}\mathbf{x}^{k+1} - - \mathbf{z}^{k+1} - - .. [1] N., Parikh, "Proximal Algorithms", Foundations and Trends - in Optimization. 2013. + See :class:`pyproximal.optimization.cls_primal.LinearizedADMM` """ - # initialize variables - x, z = _x0z0_init(x0, z0, A, Opname="A") - Ax = A.matvec(x) if z0 is None else z - ncp = get_array_module(x) - u = ncp.zeros_like(z) - - if show: - tstart = time.time() - print( - "Linearized-ADMM\n" - "---------------------------------------------------------\n" - "Proximal operator (f): %s\n" - "Proximal operator (g): %s\n" - "Linear operator (A): %s\n" - "tau = %10e\tmu = %10e\tniter = %d\n" - % (type(proxf), type(proxg), type(A), tau, mu, niter) - ) - head = " Itn x[0] f g J = f + g" - print(head) - - # run iterations - for iiter in range(niter): - x = proxf.prox(x - mu / tau * A.rmatvec(Ax - z + u), mu) - Ax = A.matvec(x) - z = proxg.prox(Ax + u, tau) - u = u + Ax - z - - # run callback - if callback is not None: - callback(x) - - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf, pg = proxf(x), proxg(Ax) - msg = "%6g %12.5e %10.3e %10.3e %10.3e" % ( - iiter + 1, - np.real(to_numpy(x[0])), - pf, - pg, - pf + pg, - ) - print(msg) - if show: - print("\nTotal time (s) = %.2f" % (time.time() - tstart)) - print("---------------------------------------------------------\n") + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + ladmmsolve = cLinearizedADMM( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + ladmmsolve.callback = callback + x, z, _, _ = ladmmsolve.solve( + proxf=proxf, + proxg=proxg, + A=A, + x0=x0, + tau=tau, + mu=mu, + z0=z0, + niter=niter, + tol=tol or (0.0 if rtol else None), + show=show, + itershow=itershow, + ) return x, z def TwIST( - proxg: ProxOperator, + proxg: "ProxOperator", A: "LinearOperator", b: NDArray, x0: NDArray, @@ -1523,8 +1027,11 @@ def TwIST( beta: float | None = None, eigs: tuple[float, float] | None = None, niter: int = 10, + tol: float | None = None, + rtol: float | None = None, callback: Callable[[NDArray], None] | None = None, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), returncost: bool = False, ) -> NDArray | tuple[NDArray, NDArray]: r"""Two-step Iterative Shrinkage/Threshold @@ -1561,11 +1068,24 @@ def TwIST( If passed, computes `alpha` and `beta` based on them. niter : :obj:`int`, optional Number of iterations of iterative scheme + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. returncost : :obj:`bool`, optional Return cost function @@ -1573,110 +1093,38 @@ def TwIST( ------- x : :obj:`numpy.ndarray` Inverted model - j : :obj:`numpy.ndarray` + j : :obj:`numpy.ndarray`, optional Cost function Notes ----- - The TwIST algorithm can be expressed by the following recursion: - - .. math:: - - \mathbf{x}^{k+1} = (1-\alpha) \mathbf{x}^{k-1} + - (\alpha-\beta) \mathbf{x}^k + - \beta \prox_{g} (\mathbf{x}^k + \mathbf{A}^H - (\mathbf{b} - \mathbf{A}\mathbf{x}^k)). - - where :math:`\mathbf{x}^{1} = \prox_{g} (\mathbf{x}^0 + \mathbf{A}^T - (\mathbf{b} - \mathbf{A}\mathbf{x}^0))`. - - The optimal weighting parameters :math:`\alpha` and :math:`\beta` are - linked to the smallest and largest eigenvalues of - :math:`\mathbf{A}^H\mathbf{A}` as follows: - - .. math:: - - \alpha = 1 + \rho^2 \\ - \beta =\frac{2 \alpha}{\Lambda_{max} + \lambda_{min}} - - where :math:`\rho=\frac{1-\sqrt{k}}{1+\sqrt{k}}` with - :math:`k=\frac{\lambda_{min}}{\Lambda_{max}}` and - :math:`\Lambda_{max}=max(1, \lambda_{max})`. - - Experimentally, it has been observed that TwIST is robust to the - choice of such parameters. Finally, note that in the case of - :math:`\alpha=1` and :math:`\beta=1`, TwIST is identical to IST. + See :class:`pyproximal.optimization.cls_primal.TwIST` """ - # define proxf as L2 proximal - proxf = L2(Op=A, b=b) - - # find alpha and beta - if alpha is None or beta is None: - if eigs is None: - emin = A.eigs(neigs=1, which="SM") - emax = max([1, A.eigs(neigs=1, which="LM")]) - else: - emax, emin = eigs - k = emin / emax - rho = (1 - sqrt(k)) / (1 + sqrt(k)) - alpha = 1 + rho**2 - beta = 2 * alpha / (emax + emin) - - # compute proximal of g on initial guess (x_1) - xold = x0.copy() - x = proxg.prox(xold - proxf.grad(xold), 1.0) - - if show: - tstart = time.time() - print( - "TwIST\n" - "---------------------------------------------------------\n" - "Proximal operator (g): %s\n" - "Linear operator (A): %s\n" - "alpha = %10e\tbeta = %10e\tniter = %d\n" - % (type(proxg), type(A), alpha, beta, niter) - ) - head = " Itn x[0] f g J = f + g" - print(head) - - # iterate - if returncost: - j = np.zeros(niter) - for iiter in range(niter): - # compute new x - xnew = ( - (1 - alpha) * xold - + (alpha - beta) * x - + beta * proxg.prox(x - proxf.grad(x), 1.0) - ) - # save current x as old (x_i -> x_i-1) - xold = x.copy() - # save new x as current (x_i+1 -> x_i) - x = xnew.copy() - - # compute cost function - if returncost: - j[iiter] = proxf(x) + proxg(x) - - # run callback - if callback is not None: - callback(x) - - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf, pg = proxf(x), proxg(x) - msg = "%6g %12.5e %10.3e %10.3e %10.3e" % ( - iiter + 1, - np.real(to_numpy(x[0])), - pf, - pg, - pf + pg, - ) - print(msg) - if show: - print("\nTotal time (s) = %.2f" % (time.time() - tstart)) - print("---------------------------------------------------------\n") + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + twistsolve = cTwIST( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + twistsolve.callback = callback + x, _, j = twistsolve.solve( + proxg=proxg, + A=A, + b=b, + x0=x0, + alpha=alpha, + beta=beta, + eigs=eigs, + niter=niter, + tol=tol or (0.0 if rtol or returncost else None), + show=show, + itershow=itershow, + ) if returncost: return x, j else: @@ -1684,16 +1132,19 @@ def TwIST( def DouglasRachfordSplitting( - proxf: ProxOperator, - proxg: ProxOperator, + proxf: "ProxOperator", + proxg: "ProxOperator", x0: NDArray, tau: float, eta: float = 1.0, niter: int = 10, gfirst: bool = True, + tol: float | None = None, + rtol: float | None = None, callback: Callable[..., None] | None = None, callbacky: bool = False, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> tuple[NDArray, NDArray]: r"""Douglas-Rachford Splitting @@ -1724,6 +1175,15 @@ def DouglasRachfordSplitting( gfirst : :obj:`bool`, optional Apply Proximal of operator ``g`` first (``True``) or Proximal of operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector @@ -1732,6 +1192,10 @@ def DouglasRachfordSplitting( when ``callbacky=True`` show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -1742,86 +1206,47 @@ def DouglasRachfordSplitting( Notes ----- - The Douglas-Rachford Splitting algorithm can be expressed by the following - recursion [1]_, [2]_, [3]_, [4]_: - - .. math:: - - \mathbf{x}^{k} &= \prox_{\tau g}(\mathbf{y}^k) \\ - \mathbf{y}^{k+1} &= \mathbf{y}^{k} + - \eta (\prox_{\tau f}(2 \mathbf{x}^{k} - \mathbf{y}^{k}) - - \mathbf{x}^{k}) - - .. [1] Patrick L. Combettes and Jean-Christophe Pesquet. 2011. Proximal - Splitting Methods in Signal Processing. In Fixed-Point Algorithms for - Inverse Problems in Science and Engineering, Springer, pp. 185-212. - Algorithm 10.15. - https://doi.org/10.1007/978-1-4419-9569-8_10 - .. [2] Scott B. Lindstrom and Brailey Sims. 2021. Survey: Sixty Years of - Douglas-Rachford. Journal of the Australian Mathematical Society, 110, - 3, 333-370. Eq.(15). https://doi.org/10.1017/S1446788719000570 - https://arxiv.org/abs/1809.07181 - .. [3] Ryu, E.K., Yin, W., 2022. Large-Scale Convex Optimization: Algorithms - & Analyses via Monotone Operators. Cambridge University Press, - Cambridge. Eq.(2.18). https://doi.org/10.1017/9781009160865 - https://large-scale-book.mathopt.com/ - .. [4] Combettes, P.L., Pesquet, J.-C., 2008. A proximal decomposition - method for solving convex variational inverse problems. Inverse Problems - 24, 065014. Proposition 3.2. https://doi.org/10.1088/0266-5611/24/6/065014 - https://arxiv.org/abs/0807.2617 + See :class:`pyproximal.optimization.cls_primal.DouglasRachfordSplitting` """ - if show: - tstart = time.time() - print( - "Douglas-Rachford Splitting\n" - "---------------------------------------------------------\n" - f"Proximal operator (f): {type(proxf)}\n" - f"Proximal operator (g): {type(proxg)}\n" - f"tau = {tau:10e}\tniter = {niter:d}\n" - ) - head = " Itn x[0] f g J = f + g" - print(head) - - y = x0.copy() - for iiter in range(niter): - if gfirst: - x = proxg.prox(y, tau) - y = y + eta * (proxf.prox(2 * x - y, tau) - x) - else: - x = proxf.prox(y, tau) - y = y + eta * (proxg.prox(2 * x - y, tau) - x) - - # run callback - if callback is not None: - if callbacky: - callback(x, y) - else: - callback(x) - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf, pg = proxf(x), proxg(x) - print( - f"{iiter + 1:6d} {np.real(to_numpy(x[0])):12.5e} " - f"{pf:10.3e} {pg:10.3e} {pf + pg:10.3e}" - ) - - if show: - print(f"\nTotal time (s) = {time.time() - tstart:.2f}") - print("---------------------------------------------------------\n") + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + drssolve = cDouglasRachfordSplitting( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + drssolve.callback = callback + x, y, _, _ = drssolve.solve( + proxf=proxf, + proxg=proxg, + x0=x0, + tau=tau, + eta=eta, + gfirst=gfirst, + niter=niter, + tol=tol or (0.0 if rtol else None), + show=show, + itershow=itershow, + ) return x, y def PPXA( # pylint: disable=invalid-name - proxfs: list[ProxOperator], + proxfs: list["ProxOperator"], x0: NDArray | list[NDArray], tau: float, eta: float = 1.0, weights: NDArray | list[float] | None = None, niter: int = 1000, - tol: float | None = 1e-7, + tol: float | None = None, + rtol: float | None = None, callback: Callable[..., None] | None = None, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> NDArray: r"""Parallel Proximal Algorithm (PPXA) @@ -1854,13 +1279,23 @@ def PPXA( # pylint: disable=invalid-name niter : :obj:`int`, optional Number of iterations of iterative scheme. tol : :obj:`float`, optional - Tolerance on change of the solution (used as stopping criterion). - If ``tol=0``, run until ``niter`` is reached. + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -1873,115 +1308,44 @@ def PPXA( # pylint: disable=invalid-name Notes ----- - The Parallel Proximal Algorithm (PPXA) can be expressed by the following - recursion [1]_, [2]_, [3]_, [4]_: - - * :math:`\mathbf{y}_{i}^{0} = \mathbf{x}` or :math:`\mathbf{y}_{i}^{0} = \mathbf{x}_{i}` for :math:`i=1,\ldots,m` - * :math:`\mathbf{x}^{0} = \sum_{i=1}^m w_i \mathbf{y}_{i}^{0}` - * for :math:`k = 1, \ldots` - - * for :math:`i = 1, \ldots, m` - - * :math:`\mathbf{p}_{i}^{k} = \prox_{\frac{\tau}{w_i} f_i} (\mathbf{y}_{i}^{k})` - - * :math:`\mathbf{p}^{k} = \sum_{i=1}^{m} w_i \mathbf{p}_{i}^{k}` - * for :math:`i = 1, \ldots, m` - - * :math:`\mathbf{y}_{i}^{k+1} = \mathbf{y}_{i}^{k} + \eta (2 \mathbf{p}^{k} - \mathbf{x}^{k} - \mathbf{p}_i^{k})` - - * :math:`\mathbf{x}^{k+1} = \mathbf{x}^{k} + \eta (\mathbf{p}^{k} - \mathbf{x}^{k})` - - where :math:`0 < \eta < 2` and - :math:`\sum_{i=1}^m w_i = 1, \ 0 < w_i < 1`. - In the current implementation, :math:`w_i = 1 / m` when not provided. - - References - ---------- - .. [1] Combettes, P.L., Pesquet, J.-C., 2008. A proximal decomposition - method for solving convex variational inverse problems. Inverse Problems - 24, 065014. Algorithm 3.1. https://doi.org/10.1088/0266-5611/24/6/065014 - https://arxiv.org/abs/0807.2617 - .. [2] Combettes, P.L., Pesquet, J.-C., 2011. Proximal Splitting Methods in - Signal Processing, in Fixed-Point Algorithms for Inverse Problems in - Science and Engineering, Springer, pp. 185-212. Algorithm 10.27. - https://doi.org/10.1007/978-1-4419-9569-8_10 - .. [3] Bauschke, H.H., Combettes, P.L., 2011. Convex Analysis and Monotone - Operator Theory in Hilbert Spaces, 1st ed, CMS Books in Mathematics. - Springer, New York, NY. Proposition 27.8. - https://doi.org/10.1007/978-1-4419-9467-7 - .. [4] Ryu, E.K., Yin, W., 2022. Large-Scale Convex Optimization: Algorithms - & Analyses via Monotone Operators. Cambridge University Press, - Cambridge. Exercise 2.38 https://doi.org/10.1017/9781009160865 - https://large-scale-book.mathopt.com/ + See :class:`pyproximal.optimization.cls_primal.PPXA` """ - if show: - tstart = time.time() - print( - "Parallel Proximal Algorithm\n" - "---------------------------------------------------------" - ) - for i, proxf in enumerate(proxfs): - print(f"Proximal operator (f{i}): {type(proxf)}") - print(f"tau = {tau:10e}\tniter = {niter:d}\n") - head = " Itn x[0] J=sum_i f_i" - print(head) - - ncp = get_array_module(x0) - - # initialize model - m = len(proxfs) - if weights is None: - w = ncp.full(m, 1.0 / m) - else: - w = ncp.asarray(weights) - - if isinstance(x0, list) or x0.ndim == 2: - y = ncp.asarray(x0) # yi_0 = xi_0, for i = 1, ..., m - else: - y = ncp.full((m, x0.size), x0) # y1_0 = y2_0 = ... = ym_0 = x0 - - x = ncp.mean(y, axis=0) - x_old = x.copy() - - # iterate - for iiter in range(niter): - p = ncp.stack([proxfs[i].prox(y[i], tau / w[i]) for i in range(m)]) - pn = ncp.sum(w[:, None] * p, axis=0) - y = y + eta * (2 * pn - x - p) - x = x + eta * (pn - x) - - # run callback - if callback is not None: - callback(x) - - # show iteration logger - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf = ncp.sum([proxfs[i](x) for i in range(m)]) - print(f"{iiter + 1:6d} {ncp.real(to_numpy(x[0])):12.5e} {pf:10.3e}") - - # break if tolerance condition is met - if ncp.abs(x - x_old).max() < tol: - break - - x_old = x - - if show: - print(f"\nTotal time (s) = {time.time() - tstart:.2f}") - print("---------------------------------------------------------\n") - + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + ppxasolve = cPPXA( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + ppxasolve.callback = callback + x, _, _, _ = ppxasolve.solve( + proxfs=proxfs, + x0=x0, + tau=tau, + eta=eta, + weights=weights, + niter=niter, + tol=tol or (0.0 if rtol else None), + show=show, + itershow=itershow, + ) return x def ConsensusADMM( # pylint: disable=invalid-name - proxfs: list[ProxOperator], + proxfs: list["ProxOperator"], x0: NDArray, tau: float, niter: int = 1000, - tol: float | None = 1e-7, + tol: float | None = None, + rtol: float | None = None, callback: Callable[..., None] | None = None, show: bool = False, + itershow: tuple[int, int, int] = (10, 10, 10), ) -> NDArray: r"""Consensus ADMM @@ -2007,13 +1371,23 @@ def ConsensusADMM( # pylint: disable=invalid-name niter : :obj:`int`, optional Number of iterations of iterative scheme. tol : :obj:`float`, optional - Tolerance on change of the solution (used as stopping criterion). - If ``tol=0``, run until ``niter`` is reached. + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + rtol : :obj:`float`, optional + Relative tolerance on objective function wrt initial value. Stops + the solver when the ratio of the current objective function to the + initial objective function is below this value. If ``rtol=None``, + run until ``niter`` is reached or the other tolerance criterion is + met callback : :obj:`callable`, optional Function with signature (``callback(x)``) to call after each iteration where ``x`` is the current model vector show : :obj:`bool`, optional Display iterations log + itershow : :obj:`tuple`, optional + Display set log for the first N1 steps, last N2 steps, + and every N3 steps in between where N1, N2, N3 are the + three element of the list. Returns ------- @@ -2027,82 +1401,27 @@ def ConsensusADMM( # pylint: disable=invalid-name Notes ----- - The ADMM for the consensus problem can be expressed by the following - recursion [1]_, [2]_: - - * :math:`\bar{\mathbf{x}}^{0} = \mathbf{x}` - * for :math:`k = 1, \ldots` - - * for :math:`i = 1, \ldots, m` - - * :math:`\mathbf{x}_i^{k+1} = \mathrm{prox}_{\tau f_i} \left(\bar{\mathbf{x}}^{k} - \mathbf{y}_i^{k}\right)` - - * :math:`\bar{\mathbf{x}}^{k+1} = \frac{1}{m} \sum_{i=1}^m \mathbf{x}_i^{k}` - - * for :math:`i = 1, \ldots, m` - - * :math:`\mathbf{y}_i^{k+1} = \mathbf{y}_i^{k} + \mathbf{x}_i^{k+1} - \bar{\mathbf{x}}^{k+1}` - - The current implementation returns :math:`\bar{\mathbf{x}}`. - - References - ---------- - .. [1] Boyd, S., Parikh, N., Chu, E., Peleato, B., Eckstein, J., 2011. - Distributed Optimization and Statistical Learning via the Alternating - Direction Method of Multipliers. Foundations and Trends in Machine Learning, - Vol. 3, No. 1, pp 1-122. Section 7.1. https://doi.org/10.1561/2200000016 - https://stanford.edu/~boyd/papers/pdf/admm_distr_stats.pdf - .. [2] Parikh, N., Boyd, S., 2014. Proximal Algorithms. Foundations and - Trends in Optimization, Vol. 1, No. 3, pp 127-239. - Section 5.2.1. https://doi.org/10.1561/2400000003 - https://web.stanford.edu/~boyd/papers/pdf/prox_algs.pdf + See :class:`pyproximal.optimization.cls_primal.ConsensusADMM` """ - if show: - tstart = time.time() - print( - "Consensus ADMM\n---------------------------------------------------------" - ) - for i, proxf in enumerate(proxfs): - print(f"Proximal operator (f{i}): {type(proxf)}") - print(f"tau = {tau:10e}\tniter = {niter:d}\n") - head = " Itn x[0] J=sum_i f_i" - print(head) - - ncp = get_array_module(x0) - - # initialize model - m = len(proxfs) - x_bar = x0.copy() - x_bar_old = x0.copy() - y = ncp.zeros((m, x0.size), dtype=x0.dtype) - - # iterate - for iiter in range(niter): - x = ncp.stack([proxfs[i].prox(x_bar - y[i], tau) for i in range(m)]) - x_bar = ncp.mean(x, axis=0) - y = y + x - x_bar - - # run callback - if callback is not None: - callback(x_bar) - - # show iteration logger - if show: - if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf = ncp.sum([proxfs[i](x_bar) for i in range(m)]) - print( - f"{iiter + 1:6d} {ncp.real(to_numpy(x_bar[0])):12.5e} {pf:10.3e}" - ) - - # break if tolerance condition is met - if ncp.abs(x_bar - x_bar_old).max() < tol: - break - - x_bar_old = x_bar - - if show: - print(f"\nTotal time (s) = {time.time() - tstart:.2f}") - print("---------------------------------------------------------\n") - - return x_bar + callbacks = [] + if tol is not None or rtol is not None: + callbacks.append(CostNanInfCallback()) + if rtol is not None: + callbacks.append(CostToInitialCallback(rtol)) + + ccadmmsolve = cConsensusADMM( + callbacks=callbacks if len(callbacks) > 0 else None, + ) + if callback is not None: + ccadmmsolve.callback = callback + _, x, _, _, _ = ccadmmsolve.solve( + proxfs=proxfs, + x0=x0, + tau=tau, + niter=niter, + tol=tol or (0.0 if rtol else None), + show=show, + itershow=itershow, + ) + return x From b819bea1e4f5689a53ef81692e1983a4650d49cc Mon Sep 17 00:00:00 2001 From: mrava87 Date: Fri, 26 Jun 2026 10:52:49 +0100 Subject: [PATCH 22/24] feat: added custom callbacks --- pyproximal/optimization/callback.py | 39 +++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 pyproximal/optimization/callback.py diff --git a/pyproximal/optimization/callback.py b/pyproximal/optimization/callback.py new file mode 100644 index 0000000..91fdc50 --- /dev/null +++ b/pyproximal/optimization/callback.py @@ -0,0 +1,39 @@ +__all__ = [ + "ModuleUpdateCallback", +] + +from typing import TYPE_CHECKING + +from pylops.optimization.callback import Callbacks +from pylops.utils.typing import NDArray + +if TYPE_CHECKING: + from pyproximal.optimization.basesolver import Solver + + +class ModuleUpdateCallback(Callbacks): # type: ignore[misc] + """Module update callback + + This callback can be used to stop the solver when each element of + the model (i.e, solution) is updated below a certain threshold. + + It requires the solver to store the model of the previous iteration + in a variable named ``xold``. + + Parameters + ---------- + tol : :obj:`float` + Absolute value of model update below which the solver + will stop iterating. For example, if ``tol`` is 0.1, the solver + will stop when the absolute value of each element of the difference + between the current model is below below 0.1. + + """ + + def __init__(self, tol: float) -> None: + self.tol = tol + self.stop = False + + def on_step_end(self, solver: "Solver", x: NDArray) -> None: + if solver.ncp.abs(x - solver.xold).max() < self.tol: + self.stop = True From b76dbedc7eeac717d4fc2359bf0a577ef9400596 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Fri, 26 Jun 2026 11:15:16 +0100 Subject: [PATCH 23/24] feat: allow niter=None in ProximalGradient setup --- pyproximal/optimization/cls_primal.py | 45 +++++++++++++++++---------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index 0a39745..21b303d 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -454,10 +454,12 @@ def _print_setup(self, epsg_print: str, xcomplex: bool = False) -> None: head1 = " Itn x[0] f g J=f+eps*g tau" print(head1) - def _print_step(self, x: NDArray, pf: float | None, pg: float | None) -> None: + def _print_step( + self, x: NDArray, pf: float | None, pg: float | None, epsg_prev: int + ) -> None: if self.tol is None: pf, pg = self.proxf(x), self.proxg(x) - self.pfg = pf + np.sum(self.epsg[self.iiter - 1] * pg) + self.pfg = pf + np.sum(epsg_prev * pg) x0 = to_numpy(x[0]) if x.ndim == 1 else to_numpy(x[0, 0]) strx = f"{x0:1.2e} " if np.iscomplexobj(x) else f"{x0:11.4e} " tau_str = ( @@ -484,7 +486,7 @@ def setup( # type: ignore[override] eta: float = 1.0, acceleration: str | None = None, niterback: int = 100, - niter: int = 10, + niter: int | None = None, tol: float | None = None, show: bool = False, ) -> tuple[NDArray, NDArray]: @@ -520,7 +522,8 @@ def setup( # type: ignore[override] niterback : :obj:`int`, optional Max number of iterations of backtracking niter : :obj:`int`, optional - Number of iterations of iterative scheme + Number of iterations (default to ``None`` in case a user wants to + manually step over the solver) tol : :obj:`float`, optional Tolerance on change of objective function (used as stopping criterion). If ``tol=None``, run until ``niter`` is reached @@ -550,8 +553,11 @@ def setup( # type: ignore[override] # check if epgs is a vector self.epsg = self.ncp.asarray(epsg, dtype=np.float32) if self.epsg.size == 1: - self.epsg = epsg * self.ncp.ones(niter, dtype=np.float32) - epsg_print = str(self.epsg[0]) + if niter is None: + epsg_print = str(self.epsg) + else: + self.epsg = epsg * self.ncp.ones(niter, dtype=np.float32) + epsg_print = str(self.epsg[0]) else: epsg_print = "Multi" @@ -576,8 +582,9 @@ def setup( # type: ignore[override] self.t = 1.0 # create variables to track the objective function and iterations + epsg_ = self.epsg if niter is None else self.epsg[0] pf, pg = self.proxf(x), self.proxg(x) - pfg = pf + np.sum(self.epsg[self.iiter] * pg) + pfg = pf + np.sum(epsg_ * pg) self.pfg, self.pfgold = pfg, pfg self.cost: list[float] = [] @@ -619,17 +626,23 @@ def step( """ xold = x.copy() + # define epsg for current iteration + if self.epsg.ndim == 0: + epsg = self.epsg + epsg_prev = self.epsg + else: + epsg = self.epsg[self.iiter] + epsg_prev = self.epsg[self.iiter - 1] + # proximal step if not self.backtracking: if self.eta == 1.0: - x = self.proxg.prox( - y - self.tau * self.proxf.grad(y), self.epsg[self.iiter] * self.tau - ) + x = self.proxg.prox(y - self.tau * self.proxf.grad(y), epsg * self.tau) else: x = x + self.eta * ( self.proxg.prox( x - self.tau * self.proxf.grad(x), - self.epsg[self.iiter] * self.tau, + epsg * self.tau, ) - x ) @@ -639,7 +652,7 @@ def step( cast(float, self.tau), self.proxf, self.proxg, - self.epsg[self.iiter], + epsg, beta=self.beta, niterback=self.niterback, ) @@ -647,7 +660,7 @@ def step( x = x + self.eta * ( self.proxg.prox( x - self.tau * self.proxf.grad(x), - self.epsg[self.iiter] * self.tau, + epsg * self.tau, ) - x ) @@ -672,7 +685,7 @@ def step( if self.tol is not None: self.pfgold = self.pfg pf, pg = self.proxf(x), self.proxg(x) - self.pfg = pf + np.sum(self.epsg[self.iiter] * pg) + self.pfg = pf + np.sum(epsg * pg) if np.abs(1.0 - self.pfg / self.pfgold) < self.tol: self.tolbreak = True else: @@ -680,7 +693,7 @@ def step( self.iiter += 1 if show: - self._print_step(x, pf, pg) + self._print_step(x, pf, pg, epsg_prev) if self.tol is not None or show: self.cost.append(float(self.pfg)) return x, y @@ -4554,5 +4567,3 @@ def solve( # type: ignore[override] # TOD0: # - make niter optional in setup and fix and prints # in _print_setup to not include if not provided -# - add the tolerance check on the solution for PPXA and -# ConsensunADMM for the function based solvers via callbacks From 0a46108ae8bcec0e819ce7be44958d5670480975 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Fri, 26 Jun 2026 11:17:24 +0100 Subject: [PATCH 24/24] fix: fixed initialization of u in ADMM solvers --- pyproximal/optimization/cls_primal.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproximal/optimization/cls_primal.py b/pyproximal/optimization/cls_primal.py index 21b303d..a7453f2 100644 --- a/pyproximal/optimization/cls_primal.py +++ b/pyproximal/optimization/cls_primal.py @@ -2549,7 +2549,7 @@ def setup( # type: ignore[override] # initialize solver x, z = _x0z0_init(x0, z0, A, Opname="A") - self.u = self.ncp.zeros_like(x) + self.u = self.ncp.zeros_like(z) # other parameters self.sqrttau = 1.0 / sqrt(self.tau) @@ -2961,7 +2961,7 @@ def setup( # type: ignore[override] # initialize solver x, z = _x0z0_init(x0, z0, A, Opname="A") self.Ax = A.matvec(x) if z0 is None else z - self.u = self.ncp.zeros_like(x) + self.u = self.ncp.zeros_like(z) # create variables to track the objective function and iterations self.pfg, self.pfgold = np.inf, np.inf