Merge pull request #332 from vais-ral/cgls_fix

adds tolerance check to CGLS iterations as it may be unstable and diverge

8 files changed, 559 insertions, 227 deletions

diff --git a/Wrappers/Python/ccpi/optimisation/algorithms/CGLS.py b/Wrappers/Python/ccpi/optimisation/algorithms/CGLS.py
index 15acc31..d3aea8e 100755
--- a/Wrappers/Python/ccpi/optimisation/algorithms/CGLS.py
+++ b/Wrappers/Python/ccpi/optimisation/algorithms/CGLS.py
@@ -34,90 +34,72 @@ class CGLS(Algorithm):
       x_init: initial guess
       operator: operator for forward/backward projections
       data: data to operate on
+      tolerance: tolerance to stop algorithm
+      
+    Reference:
+        https://web.stanford.edu/group/SOL/software/cgls/
+      
     '''
+    
+    
+    
     def __init__(self, **kwargs):
+        
         super(CGLS, self).__init__()
         self.x        = kwargs.get('x_init', None)
         self.operator = kwargs.get('operator', None)
         self.data     = kwargs.get('data', None)
+        self.tolerance     = kwargs.get('tolerance', 1e-6)
         if self.x is not None and self.operator is not None and \
            self.data is not None:
-            print ("Calling from creator")
+            print (self.__class__.__name__ , "set_up called from creator")
             self.set_up(x_init  =kwargs['x_init'],
                                operator=kwargs['operator'],
                                data    =kwargs['data'])
-
+            
+                                    
     def set_up(self, x_init, operator , data ):
 
-        self.r = data.copy()
-        self.x = x_init * 0
-
-        self.operator = operator
-        self.d = operator.adjoint(self.r)
-
+        self.x = x_init * 0.
+        self.r = data - self.operator.direct(self.x)
+        self.s = self.operator.adjoint(self.r)
         
-        self.normr2 = self.d.squared_norm()
+        self.p = self.s
+        self.norms0 = self.s.norm()
         
-        self.s = self.operator.domain_geometry().allocate()
-        #if isinstance(self.normr2, Iterable):
-        #    self.normr2 = sum(self.normr2)
-        #self.normr2 = numpy.sqrt(self.normr2)
-        #print ("set_up" , self.normr2)
-        n = Norm2Sq(operator, self.data)
-        self.loss.append(n(x_init))
-        self.configured = True
-
-    def update(self):
-        self.update_new()
-    def update_old(self):
-        Ad = self.operator.direct(self.d)
-        #norm = (Ad*Ad).sum()
-        #if isinstance(norm, Iterable):
-        #    norm = sum(norm)
-        norm = Ad.squared_norm()
+        ##
+        self.norms = self.s.norm()
+        ##
         
-        alpha = self.normr2/norm
-        self.x += (self.d * alpha)
-        self.r -= (Ad * alpha)
-        s  = self.operator.adjoint(self.r)
-
-        normr2_new = s.squared_norm()
-        #if isinstance(normr2_new, Iterable):
-        #    normr2_new = sum(normr2_new)
-        #normr2_new = numpy.sqrt(normr2_new)
-        #print (normr2_new)
         
-        beta = normr2_new/self.normr2
-        self.normr2 = normr2_new
-        self.d = s + beta*self.d
-
-    def update_new(self):
+        self.gamma = self.norms0**2
+        self.normx = self.x.norm()
+        self.xmax = self.normx   
+        
+        self.loss.append(self.r.squared_norm())
+        self.configured = True         
 
-        Ad = self.operator.direct(self.d)
-        norm = Ad.squared_norm()
-        if norm == 0.:
-            print ('norm = 0, cannot update solution')
-            print ("self.d norm", self.d.squared_norm(), self.d.as_array())
-            raise StopIteration()
-        alpha = self.normr2/norm
-        if alpha == 0.:
-            print ('alpha = 0, cannot update solution')
-            raise StopIteration()
-        self.d *= alpha
-        Ad *= alpha
-        self.r -= Ad
         
-        self.x += self.d
+    def update(self):
         
-        self.operator.adjoint(self.r, out=self.s)
-        s = self.s
-
-        normr2_new = s.squared_norm()
+        self.q = self.operator.direct(self.p)
+        delta = self.q.squared_norm()
+        alpha = self.gamma/delta
+                        
+        self.x += alpha * self.p
+        self.r -= alpha * self.q
+        
+        self.s = self.operator.adjoint(self.r)
+        
+        self.norms = self.s.norm()
+        self.gamma1 = self.gamma
+        self.gamma = self.norms**2
+        self.beta = self.gamma/self.gamma1
+        self.p = self.s + self.beta * self.p   
         
-        beta = normr2_new/self.normr2
-        self.normr2 = normr2_new
-        self.d *= (beta/alpha) 
-        self.d += s
+        self.normx = self.x.norm()
+        self.xmax = numpy.maximum(self.xmax, self.normx)
+                    
 
     def update_objective(self):
         a = self.r.squared_norm()
@@ -125,10 +107,17 @@ class CGLS(Algorithm):
             raise StopIteration()
         self.loss.append(a)
         
-#    def should_stop(self):
-#        if self.iteration > 0:
-#            x = self.get_last_objective()
-#            a = x > 0
-#            return self.max_iteration_stop_cryterion() or (not a)
-#        else:
-#            return False
+    def should_stop(self):
+        return self.flag() or self.max_iteration_stop_cryterion()
+ 
+    def flag(self):
+        flag  = (self.norms <= self.norms0 * self.tolerance) or (self.normx * self.tolerance >= 1)
+
+        if flag:
+            self.update_objective()
+            if self.iteration > self._iteration[-1]:
+                print (self.verbose_output())
+            print('Tolerance is reached: {}'.format(self.tolerance))
+
+        return flag
+ 
diff --git a/Wrappers/Python/demos/CGLS_examples/CGLS_Tikhonov.py b/Wrappers/Python/demos/CGLS_examples/CGLS_Tikhonov.py
deleted file mode 100644
index 653e191..0000000
--- a/Wrappers/Python/demos/CGLS_examples/CGLS_Tikhonov.py
+++ /dev/null
@@ -1,106 +0,0 @@
-#========================================================================
-# Copyright 2019 Science Technology Facilities Council
-# Copyright 2019 University of Manchester
-#
-# This work is part of the Core Imaging Library developed by Science Technology
-# Facilities Council and University of Manchester
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#         http://www.apache.org/licenses/LICENSE-2.0.txt
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-#=========================================================================
-
-""" 
-Compare solutions of PDHG & "Block CGLS" algorithms for 
-
-
-Problem:     min_x alpha * ||\grad x ||^{2}_{2} + || A x - g ||_{2}^{2}
-
-
-             A: Projection operator
-             g: Sinogram
-
-"""
-
-
-from ccpi.framework import AcquisitionGeometry, BlockDataContainer, AcquisitionData
-
-import numpy as np 
-import numpy                          
-import matplotlib.pyplot as plt
-
-from ccpi.optimisation.algorithms import CGLS
-from ccpi.optimisation.operators import BlockOperator, Gradient
-       
-from ccpi.framework import TestData
-import os, sys
-from ccpi.astra.ops import AstraProjectorSimple 
-
-# Load Data  
-loader = TestData(data_dir=os.path.join(sys.prefix, 'share','ccpi'))                 
-N = 150
-M = 150
-data = loader.load(TestData.SIMPLE_PHANTOM_2D, size=(N,M), scale=(0,1))
-
-ig = data.geometry
-
-detectors = N
-angles = np.linspace(0, np.pi, N, dtype=np.float32)
-
-ag = AcquisitionGeometry('parallel','2D', angles, detectors)
-
-device = input('Available device: GPU==1 / CPU==0 ')
-if device=='1':
-    dev = 'gpu'
-else:
-    dev = 'cpu'
-
-Aop = AstraProjectorSimple(ig, ag, dev)    
-sin = Aop.direct(data)
-
-noisy_data = AcquisitionData( sin.as_array() + np.random.normal(0,3,ig.shape))
-
-# Show Ground Truth and Noisy Data
-plt.figure(figsize=(10,10))
-plt.subplot(2,1,1)
-plt.imshow(data.as_array())
-plt.title('Ground Truth')
-plt.colorbar()
-plt.subplot(2,1,2)
-plt.imshow(noisy_data.as_array())
-plt.title('Noisy Data')
-plt.colorbar()
-plt.show()
-
-# Setup and run the CGLS algorithm  
-#alpha = 50
-#Grad = Gradient(ig)
-#
-## Form Tikhonov as a Block CGLS structure
-#op_CGLS = BlockOperator( Aop, alpha * Grad, shape=(2,1))
-#block_data = BlockDataContainer(noisy_data, Grad.range_geometry().allocate())
-#
-#x_init = ig.allocate()      
-#cgls = CGLS(x_init=x_init, operator=op_CGLS, data=block_data)
-#cgls.max_iteration = 1000
-#cgls.update_objective_interval = 200
-#cgls.run(1000,verbose=False)
-
-#%%
-# Show results
-plt.figure(figsize=(5,5))
-plt.imshow(cgls.get_output().as_array())
-plt.title('CGLS reconstruction')
-plt.colorbar()
-plt.show()
-
-            
diff --git a/Wrappers/Python/demos/CGLS_examples/CGLS_Tomography.py b/Wrappers/Python/demos/CGLS_examples/CGLS_Tomography.py
new file mode 100644
index 0000000..2ac050f
--- /dev/null
+++ b/Wrappers/Python/demos/CGLS_examples/CGLS_Tomography.py
@@ -0,0 +1,211 @@
+#========================================================================
+# Copyright 2019 Science Technology Facilities Council
+# Copyright 2019 University of Manchester
+#
+# This work is part of the Core Imaging Library developed by Science Technology
+# Facilities Council and University of Manchester
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+#=========================================================================
+
+""" 
+Conjugate Gradient for (Regularized) Least Squares for Tomography
+
+
+Problem:     min_u alpha * || L x ||^{2}_{2} + || A u - g ||_{2}^{2}
+
+             A: Projection operator
+             g: Sinogram
+             L: Identity or Gradient Operator
+
+"""
+
+
+from ccpi.framework import ImageGeometry, ImageData, \
+                            AcquisitionGeometry, BlockDataContainer, AcquisitionData
+
+import numpy as np 
+import numpy                          
+import matplotlib.pyplot as plt
+
+from ccpi.optimisation.algorithms import CGLS
+from ccpi.optimisation.operators import BlockOperator, Gradient, Identity
+       
+import tomophantom
+from tomophantom import TomoP2D
+from ccpi.astra.operators import AstraProjectorSimple 
+import os
+
+
+# Load  Shepp-Logan phantom 
+model = 1 # select a model number from the library
+N = 64 # set dimension of the phantom
+path = os.path.dirname(tomophantom.__file__)
+path_library2D = os.path.join(path, "Phantom2DLibrary.dat")
+phantom_2D = TomoP2D.Model(model, N, path_library2D)
+
+ig = ImageGeometry(voxel_num_x = N, voxel_num_y = N)
+data = ImageData(phantom_2D)
+
+detectors =  N
+angles = np.linspace(0, np.pi, 180, dtype=np.float32)
+
+ag = AcquisitionGeometry('parallel','2D', angles, detectors)
+
+device = input('Available device: GPU==1 / CPU==0 ')
+
+if device =='1':
+    dev = 'gpu'
+else:
+    dev = 'cpu'
+
+Aop = AstraProjectorSimple(ig, ag, dev)    
+sin = Aop.direct(data)
+
+np.random.seed(10)
+noisy_data = AcquisitionData( sin.as_array() + np.random.normal(0,1,ag.shape))
+
+# Show Ground Truth and Noisy Data
+plt.figure(figsize=(10,10))
+plt.subplot(2,1,1)
+plt.imshow(data.as_array())
+plt.title('Ground Truth')
+plt.colorbar()
+plt.subplot(2,1,2)
+plt.imshow(noisy_data.as_array())
+plt.title('Noisy Data')
+plt.colorbar()
+plt.show()
+
+# Setup and run the simple CGLS algorithm  
+x_init = ig.allocate()  
+
+cgls1 = CGLS(x_init = x_init, operator = Aop, data = noisy_data)
+cgls1.max_iteration = 20
+cgls1.update_objective_interval = 5
+cgls1.run(20, verbose = True)
+
+# Setup and run the regularised CGLS algorithm  (Tikhonov with Identity)
+
+x_init = ig.allocate()  
+alpha1 = 50
+op1 = Identity(ig)
+
+block_op1 = BlockOperator( Aop, alpha1 * op1, shape=(2,1))
+block_data1 = BlockDataContainer(noisy_data, op1.range_geometry().allocate())
+   
+cgls2 = CGLS(x_init = x_init, operator = block_op1, data = block_data1)
+cgls2.max_iteration = 200
+cgls2.update_objective_interval = 10
+cgls2.run(200, verbose=True)
+
+# Setup and run the regularised CGLS algorithm  (Tikhonov with Gradient)
+
+x_init = ig.allocate() 
+alpha2 = 25
+op2 = Gradient(ig)
+
+block_op2 = BlockOperator( Aop, alpha2 * op2, shape=(2,1))
+block_data2 = BlockDataContainer(noisy_data, op2.range_geometry().allocate())
+   
+cgls3 = CGLS(x_init=x_init, operator = block_op2, data = block_data2)
+cgls3.max_iteration = 200
+cgls3.update_objective_interval = 5
+cgls3.run(200, verbose = True)
+
+# Show results
+plt.figure(figsize=(8,8))
+
+plt.subplot(2,2,1)
+plt.imshow(data.as_array())
+plt.title('Ground Truth')
+
+plt.subplot(2,2,2)
+plt.imshow(cgls1.get_output().as_array())
+plt.title('GGLS')
+
+plt.subplot(2,2,3)
+plt.imshow(cgls2.get_output().as_array())
+plt.title('Regularised GGLS L = {} * Identity'.format(alpha1))
+
+plt.subplot(2,2,4)
+plt.imshow(cgls3.get_output().as_array())
+plt.title('Regularised GGLS L =  {} * Gradient'.format(alpha2))
+
+plt.show()
+
+
+
+print('Compare CVX vs Regularised CG with L = Gradient')
+
+from ccpi.optimisation.operators import SparseFiniteDiff
+import astra
+import numpy
+
+try:
+    from cvxpy import *
+    cvx_not_installable = True
+except ImportError:
+    cvx_not_installable = False
+
+
+if cvx_not_installable:
+    
+    ##Construct problem    
+    u = Variable(N*N)
+    #q = Variable()
+    
+    DY = SparseFiniteDiff(ig, direction=0, bnd_cond='Neumann')
+    DX = SparseFiniteDiff(ig, direction=1, bnd_cond='Neumann')
+
+    regulariser = alpha2**2 * sum_squares(norm(vstack([DX.matrix() * vec(u), DY.matrix() * vec(u)]), 2, axis = 0))
+    
+    # create matrix representation for Astra operator
+
+    vol_geom = astra.create_vol_geom(N, N)
+    proj_geom = astra.create_proj_geom('parallel', 1.0, detectors, angles)
+
+    proj_id = astra.create_projector('line', proj_geom, vol_geom)
+
+    matrix_id = astra.projector.matrix(proj_id)
+
+    ProjMat = astra.matrix.get(matrix_id)
+    
+    fidelity = sum_squares( ProjMat * u - noisy_data.as_array().ravel()) 
+        
+    solver = MOSEK
+    obj =  Minimize( regulariser +  fidelity)
+    prob = Problem(obj)
+    result = prob.solve(verbose = True, solver = solver)    
+
+
+plt.figure(figsize=(10,20))
+
+plt.subplot(1,3,1)
+plt.imshow(np.reshape(u.value, (N, N)))    
+plt.colorbar()
+
+plt.subplot(1,3,2)
+plt.imshow(cgls3.get_output().as_array())    
+plt.colorbar()
+
+plt.subplot(1,3,3)
+plt.imshow(np.abs(cgls3.get_output().as_array() - np.reshape(u.value, (N, N)) ))    
+plt.colorbar()
+
+plt.show()
+
+print('Primal Objective (CVX) {} '.format(obj.value))
+print('Primal Objective (CGLS) {} '.format(cgls3.objective[-1]))
+
diff --git a/Wrappers/Python/demos/CGLS_examples/LinearSystem.py b/Wrappers/Python/demos/CGLS_examples/LinearSystem.py
new file mode 100644
index 0000000..cc398cb
--- /dev/null
+++ b/Wrappers/Python/demos/CGLS_examples/LinearSystem.py
@@ -0,0 +1,82 @@
+#========================================================================
+# Copyright 2019 Science Technology Facilities Council
+# Copyright 2019 University of Manchester
+#
+# This work is part of the Core Imaging Library developed by Science Technology
+# Facilities Council and University of Manchester
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+#=========================================================================
+
+import numpy
+from ccpi.optimisation.operators import *
+from ccpi.optimisation.algorithms import *
+from ccpi.optimisation.functions import *
+from ccpi.framework import *
+
+# Problem dimension.
+m = 200
+n = 200
+
+numpy.random.seed(10)
+
+# Create matrix A and data b
+Amat = numpy.asarray( numpy.random.randn(m, n), dtype = numpy.float32)
+bmat = numpy.asarray( numpy.random.randn(m), dtype=numpy.float32)
+
+# Geometries for data and solution
+vgb = VectorGeometry(m)
+vgx = VectorGeometry(n)
+
+b = vgb.allocate(0, dtype=numpy.float32)
+b.fill(bmat)
+A = LinearOperatorMatrix(Amat)
+
+# Setup and run CGLS
+x_init = vgx.allocate()
+
+cgls = CGLS(x_init = x_init, operator = A, data = b)
+cgls.max_iteration = 2000
+cgls.update_objective_interval = 200
+cgls.run(2000, verbose = True)
+
+try:
+    from cvxpy import *
+    cvx_not_installable = True
+except ImportError:
+    cvx_not_installable = False
+
+# Compare with CVX
+x = Variable(n)
+obj = Minimize(sum_squares(A.A*x - b.as_array()))
+prob = Problem(obj)
+
+# choose solver
+if 'MOSEK' in installed_solvers():
+    solver = MOSEK
+else:
+    solver = SCS 
+
+result = prob.solve(solver = MOSEK)
+
+
+diff_sol = x.value - cgls.get_output().as_array()
+ 
+print('Error |CVX - CGLS| = {}'.format(numpy.sum(numpy.abs(diff_sol))))
+print('CVX objective = {}'.format(obj.value))
+print('CGLS objective = {}'.format(cgls.objective[-1]))
+
+
+
+
diff --git a/Wrappers/Python/demos/CGLS_examples/Regularised_CGLS_Denoising.py b/Wrappers/Python/demos/CGLS_examples/Regularised_CGLS_Denoising.py
new file mode 100644
index 0000000..c124fa1
--- /dev/null
+++ b/Wrappers/Python/demos/CGLS_examples/Regularised_CGLS_Denoising.py
@@ -0,0 +1,146 @@
+#========================================================================
+# Copyright 2019 Science Technology Facilities Council
+# Copyright 2019 University of Manchester
+#
+# This work is part of the Core Imaging Library developed by Science Technology
+# Facilities Council and University of Manchester
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+#=========================================================================
+
+""" 
+Conjugate Gradient for (Regularized) Least Squares for Tomography
+
+
+Problem:     min_u alpha * || L x ||^{2}_{2} + || A u - g ||_{2}^{2}
+
+             A: Identity operator
+             g: Sinogram
+             L: Identity or Gradient Operator
+
+"""
+
+
+from ccpi.framework import ImageGeometry, ImageData, \
+                            AcquisitionGeometry, BlockDataContainer, AcquisitionData
+
+import numpy as np 
+import numpy                          
+import matplotlib.pyplot as plt
+
+from ccpi.optimisation.algorithms import CGLS
+from ccpi.optimisation.operators import BlockOperator, Gradient, Identity
+from ccpi.framework import TestData
+       
+import os, sys
+
+loader = TestData(data_dir=os.path.join(sys.prefix, 'share','ccpi'))
+data = loader.load(TestData.SHAPES)
+ig = data.geometry
+ag = ig
+
+noisy_data = ImageData(TestData.random_noise(data.as_array(), mode = 'gaussian', seed = 1))
+#noisy_data = ImageData(data.as_array())
+
+# Show Ground Truth and Noisy Data
+plt.figure(figsize=(10,10))
+plt.subplot(2,1,1)
+plt.imshow(data.as_array())
+plt.title('Ground Truth')
+plt.colorbar()
+plt.subplot(2,1,2)
+plt.imshow(noisy_data.as_array())
+plt.title('Noisy Data')
+plt.colorbar()
+plt.show()
+
+# Setup and run the regularised CGLS algorithm  (Tikhonov with Gradient)
+x_init = ig.allocate() 
+alpha = 2
+op = Gradient(ig)
+
+block_op = BlockOperator( Identity(ig), alpha * op, shape=(2,1))
+block_data = BlockDataContainer(noisy_data, op.range_geometry().allocate())
+   
+cgls = CGLS(x_init=x_init, operator = block_op, data = block_data)
+cgls.max_iteration = 200
+cgls.update_objective_interval = 5
+cgls.run(200, verbose = True)
+
+# Show results
+plt.figure(figsize=(20,10))
+plt.subplot(3,1,1)
+plt.imshow(data.as_array())
+plt.title('Ground Truth')
+plt.subplot(3,1,2)
+plt.imshow(noisy_data.as_array())
+plt.title('Noisy')
+plt.subplot(3,1,3)
+plt.imshow(cgls.get_output().as_array())
+plt.title('Regularised GGLS with Gradient')
+plt.show()
+
+#%%
+
+print('Compare CVX vs Regularised CG with L = Gradient')
+
+from ccpi.optimisation.operators import SparseFiniteDiff
+
+
+try:
+    from cvxpy import *
+    cvx_not_installable = True
+except ImportError:
+    cvx_not_installable = False
+
+
+if cvx_not_installable:
+    
+    ##Construct problem    
+    u = Variable(ig.shape)
+    #q = Variable()
+    
+    DY = SparseFiniteDiff(ig, direction=0, bnd_cond='Neumann')
+    DX = SparseFiniteDiff(ig, direction=1, bnd_cond='Neumann')
+
+    regulariser = alpha**2 * sum_squares(norm(vstack([DX.matrix() * vec(u), DY.matrix() * vec(u)]), 2, axis = 0))
+        
+    fidelity = sum_squares( u - noisy_data.as_array()) 
+        
+    # choose solver
+    if 'MOSEK' in installed_solvers():
+        solver = MOSEK
+    else:
+        solver = SCS 
+        
+    obj =  Minimize( regulariser +  fidelity)
+    prob = Problem(obj)
+    result = prob.solve(verbose = True, solver = MOSEK)    
+
+
+plt.figure(figsize=(20,20))
+plt.subplot(3,1,1)
+plt.imshow(np.reshape(u.value, ig.shape))    
+plt.colorbar()
+plt.subplot(3,1,2)
+plt.imshow(cgls.get_output().as_array())    
+plt.colorbar()
+plt.subplot(3,1,3)
+plt.imshow(np.abs(cgls.get_output().as_array() - np.reshape(u.value, ig.shape) ))    
+plt.colorbar()
+plt.show()
+
+print('Primal Objective (CVX) {} '.format(obj.value))
+print('Primal Objective (CGLS) {} '.format(cgls.objective[-1]))
+
diff --git a/Wrappers/Python/demos/CompareAlgorithms/CGLS_FISTA_PDHG_LeastSquares.py b/Wrappers/Python/demos/CompareAlgorithms/CGLS_FISTA_PDHG_LeastSquares.py
index 672d4bc..09e350b 100644
--- a/Wrappers/Python/demos/CompareAlgorithms/CGLS_FISTA_PDHG_LeastSquares.py
+++ b/Wrappers/Python/demos/CompareAlgorithms/CGLS_FISTA_PDHG_LeastSquares.py
@@ -20,8 +20,8 @@
 #=========================================================================
 
 """ 
-Compare solutions of FISTA & PDHG 
-                   & CGLS  & Astra Built-in algorithms for Least Squares
+Compare solutions of FISTA & PDHG & CGLS  
+                    & Astra Built-in algorithms for Least Squares
 
 
 Problem:     min_x || A x - g ||_{2}^{2}
@@ -32,7 +32,7 @@ Problem:     min_x || A x - g ||_{2}^{2}
 """
 
 
-from ccpi.framework import ImageData, TestData, AcquisitionGeometry
+from ccpi.framework import ImageData, ImageGeometry, AcquisitionGeometry
 
 import numpy as np 
 import numpy                          
@@ -43,42 +43,48 @@ from ccpi.optimisation.algorithms import PDHG, CGLS, FISTA
 from ccpi.optimisation.functions import ZeroFunction, L2NormSquared, FunctionOperatorComposition
 from ccpi.astra.ops import AstraProjectorSimple
 import astra   
-import os, sys
 
+import tomophantom
+from tomophantom import TomoP2D
+import os
 
-# Load Data 
-loader = TestData(data_dir=os.path.join(sys.prefix, 'share','ccpi'))  
-                     
-N = 50
-M = 50
-data = loader.load(TestData.SIMPLE_PHANTOM_2D, size=(N,M), scale=(0,1))
-ig = data.geometry
+# Load  Shepp-Logan phantom 
+model = 1 # select a model number from the library
+N = 128 # set dimension of the phantom
+path = os.path.dirname(tomophantom.__file__)
+path_library2D = os.path.join(path, "Phantom2DLibrary.dat")
+phantom_2D = TomoP2D.Model(model, N, path_library2D)
 
-detectors = N
-angles = np.linspace(0, np.pi, N, dtype=np.float32)
+ig = ImageGeometry(voxel_num_x = N, voxel_num_y = N)
+data = ImageData(phantom_2D)
+
+detectors =  N
+angles = np.linspace(0, np.pi, 180, dtype=np.float32)
 
-device = input('Available device: GPU==1 / CPU==0 ')
 ag = AcquisitionGeometry('parallel','2D', angles, detectors)
-if device=='1':
+
+device = input('Available device: GPU==1 / CPU==0 ')
+
+if device =='1':
     dev = 'gpu'
 else:
     dev = 'cpu'
-    
-Aop = AstraProjectorSimple(ig, ag, dev)
-sin = Aop.direct(data)
 
-noisy_data = sin 
+Aop = AstraProjectorSimple(ig, ag, dev)    
+sinogram = Aop.direct(data)
+
+
 
 ###############################################################################
 # Setup and run Astra CGLS algorithm
 vol_geom = astra.create_vol_geom(N, N)
 proj_geom = astra.create_proj_geom('parallel', 1.0, detectors, angles)
-proj_id = astra.create_projector('linear', proj_geom, vol_geom)
+proj_id = astra.create_projector('line', proj_geom, vol_geom)
 
-# Create a sinogram from a phantom
-sinogram_id, sinogram = astra.create_sino(data.as_array(), proj_id)
+# Create a sinogram id
+sinogram_id = astra.data2d.create('-sino', proj_geom, sinogram.as_array())
 
-# Create a data object for the reconstruction
+# Create a data id
 rec_id = astra.data2d.create('-vol', vol_geom)
 
 cgls_astra = astra.astra_dict('CGLS')
@@ -89,45 +95,51 @@ cgls_astra['ProjectorId'] = proj_id
 # Create the algorithm object from the configuration structure
 alg_id = astra.algorithm.create(cgls_astra)
 
-astra.algorithm.run(alg_id, 1000)
+astra.algorithm.run(alg_id, 500)
+recon_cgls_astra = ImageData(astra.data2d.get(rec_id))
 
-recon_cgls_astra = astra.data2d.get(rec_id)
 
-###############################################################################
-# Setup and run the CGLS algorithm  
-x_init = ig.allocate()             
-cgls = CGLS(x_init=x_init, operator=Aop, data=noisy_data)
-cgls.max_iteration = 1000
-cgls.update_objective_interval = 200
-cgls.run(1000, verbose=False)
+#%%
+
+# Setup and run the simple CGLS algorithm  
+x_init = ig.allocate()  
+
+cgls = CGLS(x_init = x_init, operator = Aop, data = sinogram)
+cgls.max_iteration = 500
+cgls.update_objective_interval = 100
+cgls.run(500, verbose = True)
+
+#%%
 
 ###############################################################################
 # Setup and run the PDHG algorithm 
 operator = Aop
-f = L2NormSquared(b = noisy_data)
+f = L2NormSquared(b = sinogram)
 g = ZeroFunction()
 
 ## Compute operator Norm
 normK = operator.norm()
 
 ## Primal & dual stepsizes
-sigma = 1
+sigma = 0.02
 tau = 1/(sigma*normK**2)
 
+
 pdhg = PDHG(f=f,g=g,operator=operator, tau=tau, sigma=sigma, memopt=True)
 pdhg.max_iteration = 1000
-pdhg.update_objective_interval = 200
+pdhg.update_objective_interval = 100
 pdhg.run(1000, verbose=True)
 
+#%%
 ###############################################################################
 # Setup and run the FISTA algorithm 
-fidelity = FunctionOperatorComposition(L2NormSquared(b=noisy_data), Aop)
+fidelity = FunctionOperatorComposition(L2NormSquared(b=sinogram), Aop)
 regularizer = ZeroFunction()
 
 fista = FISTA(x_init=x_init , f=fidelity, g=regularizer)
-fista.max_iteration = 1000
-fista.update_objective_interval = 200
-fista.run(1000, verbose=True)
+fista.max_iteration = 500
+fista.update_objective_interval = 100
+fista.run(500, verbose = True)
 
 #%% Show results
 
@@ -150,40 +162,38 @@ plt.colorbar()
 plt.title('PDHG reconstruction')
 
 plt.subplot(2,2,4)
-plt.imshow(recon_cgls_astra)
+plt.imshow(recon_cgls_astra.as_array())
 plt.colorbar()
 plt.title('CGLS astra')
 
-diff1 = pdhg.get_output() - cgls.get_output()
-diff2 = fista.get_output() - cgls.get_output()
-diff3 = ImageData(recon_cgls_astra) - cgls.get_output()
+diff1 = pdhg.get_output() - recon_cgls_astra
+diff2 = fista.get_output() - recon_cgls_astra
+diff3 = cgls.get_output() - recon_cgls_astra
 
 plt.figure(figsize=(15,15))
 
 plt.subplot(3,1,1)
 plt.imshow(diff1.abs().as_array())
-plt.title('Diff PDHG vs CGLS')
+plt.title('Diff PDHG vs CGLS astra')
 plt.colorbar()
 
 plt.subplot(3,1,2)
 plt.imshow(diff2.abs().as_array())
-plt.title('Diff FISTA vs CGLS')
+plt.title('Diff FISTA vs CGLS astra')
 plt.colorbar()
 
 plt.subplot(3,1,3)
 plt.imshow(diff3.abs().as_array())
-plt.title('Diff CLGS astra vs CGLS')
+plt.title('Diff CLGS vs CGLS astra')
 plt.colorbar()
 
 
-#%%
+cgls_astra_obj = fidelity(ImageData(recon_cgls_astra))
 
 print('Primal Objective (FISTA) {} '.format(fista.objective[-1]))
 print('Primal Objective (CGLS) {} '.format(cgls.objective[-1]))
 print('Primal Objective (PDHG) {} '.format(pdhg.objective[-1][0]))
+print('Primal Objective (CGLS_astra) {} '.format(cgls_astra_obj))
 
 
-true_obj = (Aop.direct(cglsd.get_output())-noisy_data).squared_norm()
-print('True objective {}'.format(true_obj))
-
 
diff --git a/Wrappers/Python/demos/PDHG_examples/GatherAll/phantom.mat b/Wrappers/Python/demos/PDHG_examples/GatherAll/phantom.mat
deleted file mode 100755
index c465bbe..0000000
--- a/Wrappers/Python/demos/PDHG_examples/GatherAll/phantom.mat
+++ /dev/null
diff --git a/Wrappers/Python/test/test_algorithms.py b/Wrappers/Python/test/test_algorithms.py
index 3bb3d57..4bcc95a 100755
--- a/Wrappers/Python/test/test_algorithms.py
+++ b/Wrappers/Python/test/test_algorithms.py
@@ -72,7 +72,7 @@ class TestAlgorithms(unittest.TestCase):
         identity = Identity(ig)
         
         alg = CGLS(x_init=x_init, operator=identity, data=b)
-        alg.max_iteration = 1
+        alg.max_iteration = 200
         alg.run(20, verbose=True)
         self.assertNumpyArrayAlmostEqual(alg.x.as_array(), b.as_array())