TomographicImaging · MargaretDuff · May 24, 2024 · Mar 28, 2024 · Mar 28, 2024 · Mar 28, 2024
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,16 +3,20 @@
 
   - New Features:
     - Added wavelet operator, wrapping PyWavelets operator as a CIL operator (#1618)
-    - Added L1Sparsity function, allowing calculations of `|Ax-b|_1` and it's proximal, in the case of orthogonal operators, `A` (#1618)
+    - Added L1Sparsity function, allowing calculations of `|Ax-b|_1` and its proximal, in the case of orthogonal operators, `A` (#1618)
+    - Options in algorithms GD, ISTA and FISTA to pass a `cil.optimisation.utilities.StepSizeRule` or a `cil.optimisation.utilities.Preconditioner`(#1768)
+    - an implementation of the Armijo Rule as a child class of  `cil.optimisation.utilities.StepSizeRule` (#1768)
+    - Sensitivity preconditioners added as child classes of `cil.optimisation.utilities.Preconditioner`(#1768)
   - Enhancements:
     - Added `geometry` property to `BlockDataContainer`. Adds `__eq__` to `BlockGeometry` (#1799)
     - Raises error in `BlockDataContainer.pnorm` if the shape of the containers is not the same (#1799)
     - Operators and functions now also return when out is specified (#1742)
   - Bug fixes:
-    - gradient descent `update_objective` called twice on the initial point.
+    - gradient descent `update_objective` called twice on the initial point.(#1789)
     - ProjectionMap operator bug fix in adjoint and added documentation (#1743)
     - BlockOperator that would return a BlockDataContainer of shape (1,1) now returns the appropriate DataContainer. BlockDataContainer direct and adjoint methods accept DataContainer as parameter (#1802).
-    - BlurringOperator: remove check for geometry class (old SIRF integration bug)
+    - BlurringOperator: remove check for geometry class (old SIRF integration bug) (#1807)
+    - The `ZeroFunction` and `ConstantFunction` now have a Lipschitz constant of 1. (#1768)
   - Changes that break backwards compatibility:
     - Merged the files `BlockGeometry.py` and `BlockDataContainer.py` in `framework` to one file `block.py`. Please use `from cil.framework import BlockGeometry, BlockDataContainer` as before (#1799)
 

diff --git a/Wrappers/Python/cil/optimisation/algorithms/FISTA.py b/Wrappers/Python/cil/optimisation/algorithms/FISTA.py
@@ -18,9 +18,11 @@
 
 from cil.optimisation.algorithms import Algorithm
 from cil.optimisation.functions import ZeroFunction
+from cil.optimisation.utilities import ConstantStepSize, StepSizeRule
 import numpy
 import logging
-from numbers import Number
+from numbers import Real, Number
+import warnings
 
 log = logging.getLogger(__name__)
 
@@ -59,9 +61,13 @@ class ISTA(Algorithm):
         Differentiable function. If `None` is passed, the algorithm will use the ZeroFunction.
     g : Function or `None`
         Convex function with *simple* proximal operator. If `None` is passed, the algorithm will use the ZeroFunction.
-    step_size : positive :obj:`float`, default = None
-                Step size for the gradient step of ISTA.
-                The default :code:`step_size` is :math:`\frac{1}{L}` or 1 if `f=None`.
+    step_size : positive :obj:`float` or child class of :meth:`cil.optimisation.utilities.StepSizeRule`',  default = None
+                Step size for the gradient step of ISTA. If a float is passed, this is used as a constant step size.  If a child class of :meth:`cil.optimisation.utilities.StepSizeRule`' is passed then it's method `get_step_size` is called for each update. 
+                The default :code:`step_size` is a constant :math:`\frac{0.99*2}{L}` or 1 if `f=None`.
+     preconditioner: class with a `apply` method or a function that takes an initialised CIL function as an argument and modifies a provided `gradient`.
+            This could be a custom `preconditioner` or one provided in :meth:`~cil.optimisation.utilities.preconditoner`. If None is passed  then `self.gradient_update` will remain unmodified. 
+
+
     kwargs: Keyword arguments
         Arguments from the base class :class:`.Algorithm`.
 
@@ -93,43 +99,51 @@ class ISTA(Algorithm):
     """
 
     def _provable_convergence_condition(self):
-        return self.step_size <= 0.99*2.0/self.f.L
+        if self.preconditioner is not None:
+            raise NotImplementedError(
+                "Can't check convergence criterion if a preconditioner is used ")
+
+        if isinstance(self.step_size_rule, ConstantStepSize):
+            return self.step_size_rule.step_size <= 0.99*2.0/self.f.L
+        else:
+            raise TypeError(
+                "Can't check convergence criterion for non-constant step size")
 
     @property
     def step_size(self):
-       return self._step_size
+        if isinstance(self.step_size_rule, ConstantStepSize):
+            return self.step_size_rule.step_size
+        else:
+            warnings.warn(
+                "Note the step-size is set by a step-size rule and could change wit each iteration")
+            return self.step_size_rule.get_step_size()
 
     # Set default step size
-    def set_step_size(self, step_size):
-        """ Set default step size.
-        """
 
-        if step_size is None:
-            if isinstance(self.f, ZeroFunction):
-                self._step_size = 1
-
-            elif isinstance(self.f.L, Number):
-                self._step_size = 0.99*2.0/self.f.L
+    def _calculate_default_step_size(self):
+        """ Calculates the default step size if a step size rule or a step size is not provided. 
+        """
 
-            else:
-                raise ValueError("Function f is not differentiable")
+
+        return 0.99*2.0/self.f.L
 
-        else:
-            self._step_size = step_size
+
 
-    def __init__(self, initial, f, g, step_size = None, **kwargs):
+    def __init__(self, initial, f, g, step_size=None, preconditioner=None, **kwargs):
 
         super(ISTA, self).__init__(**kwargs)
-        self._step_size = None
-        self.set_up(initial=initial, f=f, g=g, step_size=step_size, **kwargs)
+        self._step_size = step_size
+        self.set_up(initial=initial, f=f, g=g, step_size=step_size,
+                    preconditioner=preconditioner, **kwargs)
 
-    def set_up(self, initial, f, g, step_size, **kwargs):
+    def set_up(self, initial, f, g, step_size, preconditioner, **kwargs):
         """Set up of the algorithm"""
         log.info("%s setting up", self.__class__.__name__)
         # set up ISTA
         self.initial = initial
         self.x_old = initial.copy()
         self.x = initial.copy()
+        self.gradient_update = initial.copy()
 
         if f is None:
             f = ZeroFunction()
@@ -142,28 +156,45 @@ def set_up(self, initial, f, g, step_size, **kwargs):
         self.g = g
 
         if isinstance(f, ZeroFunction) and isinstance(g, ZeroFunction):
-            raise ValueError('You set both f and g to be the ZeroFunction and thus the iterative method will not update and will remain fixed at the initial value.')
+            raise ValueError(
+                'You set both f and g to be the ZeroFunction and thus the iterative method will not update and will remain fixed at the initial value.')
 
         # set step_size
-        self.set_step_size(step_size=step_size)
+        if step_size is None:
+            self.step_size_rule = ConstantStepSize(
+                self._calculate_default_step_size())
+        elif isinstance(step_size, Real):
+            self.step_size_rule = ConstantStepSize(step_size)
+        elif isinstance(step_size, StepSizeRule):
+            self.step_size_rule = step_size
+
+        self.preconditioner = preconditioner
+
         self.configured = True
         log.info("%s configured", self.__class__.__name__)
 
-
     def update(self):
-
         r"""Performs a single iteration of ISTA
 
         .. math:: x_{k+1} = \mathrm{prox}_{\alpha g}(x_{k} - \alpha\nabla f(x_{k}))
 
         """
 
         # gradient step
-        self.f.gradient(self.x_old, out=self.x)
-        self.x_old.sapyb(1., self.x, -self.step_size, out=self.x_old)
+        self.f.gradient(self.x_old, out=self.gradient_update)
+        if self.preconditioner is not None:
+            self.preconditioner.apply(
+                self, self.gradient_update, out=self.gradient_update)
+
+        try:
+            step_size = self.step_size_rule.get_step_size(self)
+        except NameError:
+            raise NameError(msg='`step_size` must be `None`, a real float or a child class of :meth:`cil.optimisation.utilities.StepSizeRule`')
+
+        self.x_old.sapyb(1., self.gradient_update, -step_size, out=self.x_old)
 
         # proximal step
-        self.g.proximal(self.x_old, self.step_size, out=self.x)
+        self.g.proximal(self.x_old, step_size, out=self.x)
 
     def _update_previous_solution(self):
         """ Swaps the references to current and previous solution based on the :func:`~Algorithm.update_previous_solution` of the base class :class:`Algorithm`.
@@ -182,7 +213,7 @@ def update_objective(self):
         .. math:: f(x) + g(x)
 
         """
-        self.loss.append( self.f(self.x_old) + self.g(self.x_old) )
+        self.loss.append(self.f(self.x_old) + self.g(self.x_old))
 
 
 class FISTA(ISTA):
@@ -217,9 +248,12 @@ class FISTA(ISTA):
         Differentiable function.  If `None` is passed, the algorithm will use the ZeroFunction.
     g : Function or `None`
         Convex function with *simple* proximal operator. If `None` is passed, the algorithm will use the ZeroFunction.
-    step_size : positive :obj:`float`, default = None
-                Step size for the gradient step of FISTA.
-                The default :code:`step_size` is :math:`\frac{1}{L}` or 1 if `f=None`.
+    step_size : positive :obj:`float` or child class of :meth:`cil.optimisation.utilities.StepSizeRule`',  default = None
+                Step size for the gradient step of ISTA. If a float is passed, this is used as a constant step size. If a child class of :meth:`cil.optimisation.utilities.StepSizeRule`' is passed then it's method `get_step_size` is called for each update. 
+                The default :code:`step_size` is a constant :math:`\frac{1}{L}` or 1 if `f=None`.
+    preconditioner: class with a `apply` method or a function that takes an initialised CIL function as an argument and modifies a provided `gradient`.
+            This could be a custom `preconditioner` or one provided in :meth:`~cil.optimisation.utilities.preconditoner`. If None is passed  then `self.gradient_update` will remain unmodified. 
+
     kwargs: Keyword arguments
         Arguments from the base class :class:`.Algorithm`.
 
@@ -248,36 +282,32 @@ class FISTA(ISTA):
 
     """
 
-    def set_step_size(self, step_size):
-
-        """Set the default step size
+    def _calculate_default_step_size(self):
+        """Calculate the default step size if a step size rule or step size is not provided 
         """
+        return 1./self.f.L
 
-        if step_size is None:
-
-            if isinstance(self.f, ZeroFunction):
-                self._step_size = 1
 
-            elif isinstance(self.f.L, Number):
-                self._step_size = 1./self.f.L
 
-            else:
-                raise ValueError("Function f is not differentiable")
+    def _provable_convergence_condition(self):
+        if self.preconditioner is not None:
+            raise NotImplementedError(
+                "Can't check convergence criterion if a preconditioner is used ")
 
+        if isinstance(self.step_size_rule, ConstantStepSize):
+            return self.step_size_rule.step_size <= 1./self.f.L
         else:
-            self._step_size = step_size
-
-    def _provable_convergence_condition(self):
-        return self.step_size <= 1./self.f.L
+            raise TypeError(
+                "Can't check convergence criterion for non-constant step size")
 
-    def __init__(self, initial, f, g, step_size = None, **kwargs):
+    def __init__(self, initial, f, g, step_size=None,  preconditioner=None, **kwargs):
 
         self.y = initial.copy()
         self.t = 1
-        super(FISTA, self).__init__(initial=initial, f=f, g=g, step_size=step_size, **kwargs)
+        super(FISTA, self).__init__(initial=initial, f=f, g=g,
+                                    step_size=step_size,  preconditioner=preconditioner, **kwargs)
 
     def update(self):
-
         r"""Performs a single iteration of FISTA
 
         .. math::
@@ -292,11 +322,17 @@ def update(self):
 
         self.t_old = self.t
 
-        self.f.gradient(self.y, out=self.x)
+        self.f.gradient(self.y, out=self.gradient_update)
+
+        if self.preconditioner is not None:
+            self.preconditioner.apply(
+                self, self.gradient_update, out=self.gradient_update)
+
+        step_size = self.step_size_rule.get_step_size(self)
 
-        self.y.sapyb(1., self.x, -self.step_size, out=self.y)
+        self.y.sapyb(1., self.gradient_update, -step_size, out=self.y)
 
-        self.g.proximal(self.y, self.step_size, out=self.x)
+        self.g.proximal(self.y, step_size, out=self.x)
 
         self.t = 0.5*(1 + numpy.sqrt(1 + 4*(self.t_old**2)))
 
@@ -311,10 +347,10 @@ def update(self):
     from cil.framework import ImageGeometry
     f = L2NormSquared()
     g = L2NormSquared()
-    ig = ImageGeometry(3,4,4)
+    ig = ImageGeometry(3, 4, 4)
     initial = ig.allocate()
-    fista = FISTA(initial, f, g, step_size = 1443432)
+    fista = FISTA(initial, f, g, step_size=1443432)
     print(fista.is_provably_convergent())
 
-    gd = GD(initial=initial, objective = f, step_size = 1023123)
+    gd = GD(initial=initial, objective=f, step_size=1023123)
     print(gd.is_provably_convergent())