__all__ = ["Solver"]
import functools
import time
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, Any
from pylops.optimization.callback import Callbacks
from pylops.utils.typing import NDArray
if TYPE_CHECKING:
from pylops.linearoperator import LinearOperator
[docs]class Solver(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 to which the operator `Op` must be passed
- ``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
----------
Op : :obj:`pylops.LinearOperator`
Operator to invert of
callbacks : :obj:`pylops.optimization.callback.Callbacks`
Callbacks object used to implement custom callbacks
"""
def __init__(
self,
Op: "LinearOperator",
callbacks: Callbacks = None,
) -> None:
self.Op = Op
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 + "\n"
f"The Operator Op has {self.Op.shape[0]} rows and {self.Op.shape[1]} cols"
)
def _print_setup(self, *args: Any, **kwargs: Any) -> None:
pass
def _print_step(self, *args: Any, **kwargs: Any) -> None:
pass
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")
def _registercallbacks(self) -> None:
# We want to make sure that the appropriate callbacks are called
# for each method. Instead of just calling self.step, we want
# to call self.callbacks[:].on_step_begin, self.step and finally
# self.callbacks[::-1].on_step_end, for all callbacks in the list
# We can do this in an automated way by decorating all methods
def cbdecorator(func, setup=False): # func will be self.setup, self.step, etc.
@functools.wraps(func)
def wrapper(*args, **kwargs):
if self.callbacks:
for cb in self.callbacks:
# Call all on_*_begin callbacks
if setup:
getattr(cb, f"on_{func.__name__}_begin")(
self, kwargs.get("x0", None)
) # self is solver, args[0] is x
else:
getattr(cb, f"on_{func.__name__}_begin")(
self, args[0]
) # self is solver, args[0] is x
ret = func(*args, **kwargs)
if self.callbacks:
for cb in self.callbacks[::-1]:
# Call all on_*_end callbacks in reverse order
if setup:
getattr(cb, f"on_{func.__name__}_end")(
self, kwargs.get("x0", None)
)
else:
getattr(cb, f"on_{func.__name__}_end")(self, args[0])
return ret
return wrapper
for method in ["setup", "step", "run"]:
# Replace each method by its decorator
setattr(
self,
method,
cbdecorator(
getattr(self, method), True if method == "setup" else False
),
)
@abstractmethod
def setup(
self,
y: NDArray,
*args,
show: bool = False,
**kwargs,
) -> 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
----------
y : :obj:`np.ndarray`
Data of size :math:`[N \times 1]`
show : :obj:`bool`, optional
Display setup log
"""
pass
@abstractmethod
def step(
self,
x: NDArray,
*args,
show: bool = False,
**kwargs,
) -> Any:
"""Run one step of solver
This method is used to run one step of the solver. Users can change the
function signature by including any other input parameter required when applying
one step of the solver
Parameters
----------
x : :obj:`np.ndarray`
Current model vector to be updated by a step of the solver
show : :obj:`bool`, optional
Display step log
"""
pass
@abstractmethod
def run(
self,
x: NDArray,
*args,
show: bool = False,
**kwargs,
) -> Any:
"""Run multiple steps of solver
This method is used to run multiple step of the solver. Users can change the
function signature by including any other input parameter required when applying
multiple steps of the solver
Parameters
----------
x : :obj:`np.ndarray`
Current model vector to be updated by multiple steps of the solver
show : :obj:`bool`, optional
Display step log
"""
pass
def finalize(
self,
*args,
show: bool = False,
**kwargs,
) -> Any:
"""Finalize solver
This method is used to finalize the solver. Users can change the
function signature by including any other input parameter required when
finalizing the solver
Parameters
----------
show : :obj:`bool`, optional
Display finalize log
"""
self.tend = time.time()
self.telapsed = self.tend - self.tstart
if show:
self._print_finalize()
@abstractmethod
def solve(
self,
y: NDArray,
*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
----------
y : :obj:`np.ndarray`
Data
show : :obj:`bool`, optional
Display finalize log
"""
pass
def callback(
self,
x: NDArray,
*args,
**kwargs,
) -> None:
"""Callback routine
This routine must be passed by the user. Its function signature must contain
a single input that contains the current solution (when using the `solve`
method it will be automatically invoked after each step of the solve)
Parameters
----------
x : :obj:`np.ndarray`
Current solution
Examples
--------
>>> import numpy as np
>>> from pylops.basicoperators import Identity
>>> from pylops.optimization.solver import CG
>>> def callback(x):
... print(f"Running callback, current solution {x}")
...
>>> I = Identity(10)
>>> I
<10x10 Identity with dtype=float64>
>>> cgsolve = CG(I, np.arange(10))
>>> cgsolve.callback = callback
>>> x = np.ones(10)
>>> cgsolve.callback(x)
Running callback, current solution [1,1,1...]
"""
pass
