Added BB step size rule to our step size methods

Wrappers/Python/cil/optimisation/utilities/StepSizeMethods.py

@@ -18,7 +18,7 @@
 
 from abc import ABC, abstractmethod
 import numpy
-
+from numbers import Number
 
 class StepSizeRule(ABC):
     """
@@ -138,3 +138,88 @@ def get_step_size(self, algorithm):
             raise ValueError(
                 'Could not find a proper step_size in {} loops. Consider increasing alpha or max_iterations.'.format(self.max_iterations))
         return self.alpha
+
+
+class BarzilaiBorweinStepSizeRule(StepSizeRule):
+
+    """ 
+    Applies the Barzilai- Borwein rule to calculate the step size (step_size).
+
+    Let :math:`\Delta x=x_k-x_{k-1}` and :math:`\Delta g=g_k-g_{k-1}`. A Barzilai-Borwein (BB) iteration is :math:`x_{k+1}=x_k-\alpha _kg_k` where the step size :math:`\alpha _k` is either
+
+    - :math:`\alpha _k^{LONG}=\frac{\Delta x\cdot\Delta x}{\Delta x\cdot\Delta g}`, or
+
+    - :math:`\alpha _k^{SHORT}=\frac{\Delta x \cdot\Delta g}{\Delta g \cdot\Delta g}`.
+
+    Parameters
+    ----------
+    initial: float,
+        The step-size for the first iteration
+    mode: One of 'long', 'short' or 'alternate', default is 'short'. 
+        This calculates the step-size based on the LONG, SHORT or alternating between the two, starting with short. 
+    stabilisation_param: float or None, default is None 
+        In order to add stability the step-size has an upper limit of :math:`\Delta/\|g_k\|` were :math:`\Delta` is either a given positive constant, or determined to be the minimium of :math`:\Delta x: from the first 3 iterations. To "turn off" the stabilisation use `np.inf`.
+
+
+
+    Reference
+    ---------
+    - Barzilai, Jonathan; Borwein, Jonathan M. (1988). "Two-Point Step Size Gradient Methods". IMA Journal of Numerical Analysis. 8: 141–148
+    https://en.wikipedia.org/wiki/Barzilai-Borwein_method
+
+    - Burdakov, O., Dai, Y.H. and Huang, N., 2019. Stabilized barzilai-borwein method. arXiv preprint arXiv:1907.06409.
+
+    - https://en.wikipedia.org/wiki/Barzilai-Borwein_method
+    """
+
+    def __init__(self, initial, mode='short', stabilisation_param=None):
+        '''Initialises the step size rule 
+        '''
+
+        self.mode=mode
+        self.store_grad=None 
+        self.store_x=None
+        self.initial=initial
+        if stabilisation_param is None:
+            self.adaptive = True
+            stabilisation_param =numpy.inf
+        else:
+            self.adaptive = False 
+            if not ( isinstance(stabilisation_param, Number) and stabilisation_param >=0):
+                raise TypeError(" The stabilisation_param should be None, a positive number or np.inf")
+        self.stabilisation_param=stabilisation_param
+
+
+
+    def get_step_size(self, algorithm):
+        """
+        Applies the B-B rule to calculate the step size (`step_size`)
+
+        Returns
+        --------
+        the calculated step size:float
+
+        """
+        if self.store_x is None:
+            self.store_x=algorithm.x.copy()
+            self.store_grad=algorithm.gradient_update.copy()
+            return self.initial
+
+        if algorithm.gradient_update.norm()<1e-6:
+            raise StopIteration
+
+        if self.mode=='long' or (self.mode =='alternate' and algorithm.iteration%2 ==0):
+            ret = (( algorithm.x-self.store_x).dot(algorithm.x-self.store_x))/ (( algorithm.x-self.store_x).dot(algorithm.gradient_update-self.store_grad))
+        elif self.mode=='short' or (self.mode =='alternate' and algorithm.iteration%2 ==1):
+            ret = (( algorithm.x-self.store_x).dot(algorithm.gradient_update-self.store_grad))/ (( algorithm.gradient_update-self.store_grad).dot(algorithm.gradient_update-self.store_grad))
+        else:
+            raise ValueError('Mode should be chosen from "long", "short" or "alternate". ')
+
+        if (self.adaptive and algorithm.iteration <=3):
+            self.stabilisation_param = min(self.stabilisation_param,(algorithm.x-self.store_x).norm() )
+
+        ret = numpy.nanmin( numpy.array([ret, self.stabilisation_param/algorithm.gradient_update.norm()]))
+        self.store_x.fill(algorithm.x)
+        self.store_grad.fill(algorithm.gradient_update)
+
+        return ret

Wrappers/Python/cil/optimisation/utilities/__init__.py

@@ -19,5 +19,5 @@
 
 from .sampler import Sampler
 from .sampler import SamplerRandom
-from .StepSizeMethods import ConstantStepSize, ArmijoStepSizeRule, StepSizeRule
+from .StepSizeMethods import ConstantStepSize, ArmijoStepSizeRule, StepSizeRule, BarzilaiBorweinStepSizeRule
 from .preconditioner import  Preconditioner, AdaptiveSensitivity, Sensitivity

Wrappers/Python/test/test_stepsizes.py

@@ -1,9 +1,9 @@
 from cil.optimisation.algorithms import SIRT, GD, ISTA, FISTA
 from cil.optimisation.functions import LeastSquares, IndicatorBox
-from cil.framework import ImageGeometry, VectorGeometry
+from cil.framework import ImageGeometry, VectorGeometry, VectorData
 from cil.optimisation.operators import IdentityOperator, MatrixOperator
 
-from cil.optimisation.utilities import Sensitivity, AdaptiveSensitivity, Preconditioner, ConstantStepSize, ArmijoStepSizeRule
+from cil.optimisation.utilities import Sensitivity, AdaptiveSensitivity, Preconditioner, ConstantStepSize, ArmijoStepSizeRule, BarzilaiBorweinStepSizeRule
 import numpy as np
 
 from testclass import CCPiTestClass
@@ -77,3 +77,152 @@ def test_armijo_calculation(self):
         alg.gradient_update = ig.allocate(-2)
         step_size = test_stepsize.get_step_size(alg)
         self.assertAlmostEqual(step_size, 2)
+
+    def test_bb(self):
+        n = 10
+        m = 5
+
+        A = np.random.uniform(0, 1, (m, n)).astype('float32')
+        b = (A.dot(np.random.randn(n)) + 0.1 *
+             np.random.randn(m)).astype('float32')
+
+        Aop = MatrixOperator(A)
+        bop = VectorData(b)
+        ig=Aop.domain
+        initial = ig.allocate()
+        f = LeastSquares(Aop, b=bop, c=0.5)
+
+        ss_rule=BarzilaiBorweinStepSizeRule(2 )
+        self.assertEqual(ss_rule.mode, 'short')
+        self.assertEqual(ss_rule.initial, 2)
+        self.assertEqual(ss_rule.adaptive, True)
+        self.assertEqual(ss_rule.stabilisation_param, np.inf)
+
+        #Check the right errors are raised for incorrect parameters 
+
+        with self.assertRaises(TypeError):
+            ss_rule=BarzilaiBorweinStepSizeRule(2,'short',-4, )
+        with self.assertRaises(TypeError):
+            ss_rule=BarzilaiBorweinStepSizeRule(2,'long', 'banana', )
+        with self.assertRaises(ValueError):
+            ss_rule=BarzilaiBorweinStepSizeRule(2, 'banana',3 )
+            alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+            alg.run(2)
+
+        #Check stabilisation parameter unchanges if fixed 
+        ss_rule=BarzilaiBorweinStepSizeRule(2, 'long',3 )
+        self.assertEqual(ss_rule.mode, 'long')
+        self.assertFalse(ss_rule.adaptive)
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        self.assertEqual(ss_rule.stabilisation_param,3)
+        alg.run(2)
+        self.assertEqual(ss_rule.stabilisation_param,3)
+
+        #Check infinity can be passed 
+        ss_rule=BarzilaiBorweinStepSizeRule(2, 'short',np.inf )
+        self.assertEqual(ss_rule.mode, 'short')
+        self.assertFalse(ss_rule.adaptive)
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        alg.run(2)
+
+        n = 5
+        m = 5
+
+        A = np.eye(5).astype('float32')
+        b = (np.array([1,1,1,1,1])).astype('float32')
+
+        Aop = MatrixOperator(A)
+        bop = VectorData(b)
+        ig=Aop.domain
+        initial = ig.allocate(0)
+        f = LeastSquares(Aop, b=bop, c=0.5)
+        ss_rule=BarzilaiBorweinStepSizeRule(1/4, 'long',np.inf )
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        #Check the initial step size was used
+        alg.run(1)
+        self.assertNumpyArrayEqual( np.array([.25,.25,.25,.25,.25]), alg.x.as_array() )
+        #check long 
+        alg.run(1)
+        x_change= np.array([.25,.25,.25,.25,.25])-np.array([0,0,0,0,0])
+        grad_change = -np.array([.75,.75,.75,.75,.75])+np.array([1,1,1,1,1])
+        step= x_change.dot(x_change)/x_change.dot(grad_change)
+        self.assertNumpyArrayEqual( np.array([.25,.25,.25,.25,.25])+step*np.array([.75,.75,.75,.75,.75]), alg.x.as_array() )
+
+        ss_rule=BarzilaiBorweinStepSizeRule(1/4, 'long',np.inf )
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        #Check the initial step size was used
+        alg.run(1)
+        self.assertNumpyArrayEqual( np.array([.25,.25,.25,.25,.25]), alg.x.as_array() )
+        #check short
+        alg.run(1)
+        x_change= np.array([.25,.25,.25,.25,.25])-np.array([0,0,0,0,0])
+        grad_change = -np.array([.75,.75,.75,.75,.75])+np.array([1,1,1,1,1])
+        step= x_change.dot(grad_change)/grad_change.dot(grad_change)
+        self.assertNumpyArrayEqual( np.array([.25,.25,.25,.25,.25])+step*np.array([.75,.75,.75,.75,.75]), alg.x.as_array() )
+
+        #check stop iteration 
+        ss_rule=BarzilaiBorweinStepSizeRule(1, 'long',np.inf )
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        alg.run(500)
+        self.assertEqual(alg.iteration, 1)
+
+        #check adaptive
+        ss_rule=BarzilaiBorweinStepSizeRule(0.001, 'long',None )
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        self.assertEqual(ss_rule.stabilisation_param, np.inf)
+        alg.run(2)
+        self.assertNotEqual(ss_rule.stabilisation_param, np.inf)
+
+        #check stops being adaptive 
+
+        ss_rule=BarzilaiBorweinStepSizeRule(0.0000001, 'long',None )
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        self.assertEqual(ss_rule.stabilisation_param, np.inf)
+        alg.run(4)
+        self.assertNotEqual(ss_rule.stabilisation_param, np.inf)
+        a=ss_rule.stabilisation_param
+        alg.run(1)
+        self.assertEqual(ss_rule.stabilisation_param, a)
+
+
+    def test_bb_converge(self):
+        n = 10
+        m = 5
+        np.random.seed(4)
+        A = np.random.uniform(0, 1, (m, n)).astype('float32')
+        b = (A.dot(np.random.randn(n)) + 0.1 *
+             np.random.randn(m)).astype('float32')
+
+        Aop = MatrixOperator(A)
+        bop = VectorData(b)
+        ig=Aop.domain
+        initial = ig.allocate()
+        f = LeastSquares(Aop, b=bop, c=0.5)
+
+        ss_rule=ArmijoStepSizeRule(max_iterations=40)
+        alg_true = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        alg_true .run(300, verbose=0)
+
+
+
+
+
+        ss_rule=BarzilaiBorweinStepSizeRule(1/f.L, 'short')
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        alg.run(80, verbose=0)
+        self.assertNumpyArrayAlmostEqual(alg.x.as_array(), alg_true.x.as_array(), decimal=3)
+
+
+        ss_rule=BarzilaiBorweinStepSizeRule(1/f.L, 'long')
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+        alg.run(80, verbose=0)
+        self.assertNumpyArrayAlmostEqual(alg.x.as_array(), alg_true.x.as_array(), decimal=3)
+
+
+        ss_rule=BarzilaiBorweinStepSizeRule(1/f.L, 'alternate')
+        alg = GD(initial=initial, objective_function=f, step_size=ss_rule)
+
+        alg.run(80, verbose=0)
+        self.assertNumpyArrayAlmostEqual(alg.x.as_array(), alg_true.x.as_array(), decimal=3)
+
+

docs/source/optimisation.rst

@@ -621,6 +621,10 @@ We also have a number of example classes:
 .. autoclass:: cil.optimisation.utilities.StepSizeMethods.ArmijoStepSizeRule
    :members:
 
+.. autoclass:: cil.optimisation.utilities.StepSizeMethods.BarzilaiBorweinStepSizeRule
+   :members:
+
+
 
 Preconditioners
 ----------------