memopt test

12 files changed, 445 insertions, 183 deletions

diff --git a/Wrappers/Python/ccpi/framework/BlockDataContainer.py b/Wrappers/Python/ccpi/framework/BlockDataContainer.py
index 888950d..69b6931 100755
--- a/Wrappers/Python/ccpi/framework/BlockDataContainer.py
+++ b/Wrappers/Python/ccpi/framework/BlockDataContainer.py
@@ -186,9 +186,15 @@ class BlockDataContainer(object):
         return self.clone()

     def clone(self):

         return type(self)(*[el.copy() for el in self.containers], shape=self.shape)

-    def fill(self, x):

-        for el,ot in zip(self.containers, x.containers):

-            el.fill(ot)

+

+    def fill(self, other):

+        if isinstance (other, BlockDataContainer):

+            if not self.is_compatible(other):

+                raise ValueError('Incompatible containers')

+            for el,ot in zip(self.containers, other.containers):

+                el.fill(ot)

+        else:

+            return ValueError('Cannot fill with object provided {}'.format(type(other)))

     

     def __add__(self, other):

         return self.add( other )

diff --git a/Wrappers/Python/ccpi/optimisation/algorithms/PDHG.py b/Wrappers/Python/ccpi/optimisation/algorithms/PDHG.py
index ac37e13..2a69857 100644
--- a/Wrappers/Python/ccpi/optimisation/algorithms/PDHG.py
+++ b/Wrappers/Python/ccpi/optimisation/algorithms/PDHG.py
@@ -24,6 +24,7 @@ class PDHG(Algorithm):
         self.g        = kwargs.get('g', None)
         self.tau      = kwargs.get('tau', None)
         self.sigma    = kwargs.get('sigma', None)
+        self.memopt   = kwargs.get('memopt', False)
 
         if self.f is not None and self.operator is not None and \
            self.g is not None:
@@ -45,65 +46,73 @@ class PDHG(Algorithm):
         self.y_old = self.operator.range_geometry().allocate()
         
         self.xbar = self.x_old.copy()
-        #x_tmp = x_old
+        
         self.x = self.x_old.copy()
         self.y = self.y_old.copy()
-        #y_tmp = y_old
+        if self.memopt:
+            self.y_tmp = self.y_old.copy()
+            self.x_tmp = self.x_old.copy()
         #y = y_tmp
             
         # relaxation parameter
         self.theta = 1
 
     def update(self):
-        # Gradient descent, Dual problem solution
-        self.y_old += self.sigma * self.operator.direct(self.xbar)
-        self.y = self.f.proximal_conjugate(self.y_old, self.sigma)
-        
-        # Gradient ascent, Primal problem solution
-        self.x_old -= self.tau * self.operator.adjoint(self.y)
-        self.x = self.g.proximal(self.x_old, self.tau)
-        
-        #Update
-        #xbar = x + theta * (x - x_old)
-        self.xbar.fill(self.x)
-        self.xbar -= self.x_old 
-        self.xbar *= self.theta
-        self.xbar += self.x
-                        
-#        self.x_old.fill(self.x)
-#        self.y_old.fill(self.y)
-        self.y_old = self.y.copy()
-        self.x_old = self.x.copy()
-        #self.y = self.y_old
+        if self.memopt:
+            # Gradient descent, Dual problem solution
+            # self.y_old += self.sigma * self.operator.direct(self.xbar)
+            self.operator.direct(self.xbar, out=self.y_tmp)
+            self.y_tmp *= self.sigma
+            self.y_old += self.y_tmp
+
+            #self.y = self.f.proximal_conjugate(self.y_old, self.sigma)
+            self.f.proximal_conjugate(self.y_old, self.sigma, out=self.y)
+            
+            # Gradient ascent, Primal problem solution
+            self.operator.adjoint(self.y, out=self.x_tmp)
+            self.x_tmp *= self.tau
+            self.x_old -= self.x_tmp
+            
+            self.g.proximal(self.x_old, self.tau, out=self.x)
+            
+            #Update
+            self.x.subtract(self.x_old, out=self.xbar)
+            #self.xbar -= self.x_old 
+            self.xbar *= self.theta
+            self.xbar += self.x
+                            
+            self.x_old.fill(self.x)
+            self.y_old.fill(self.y)
+            #self.y_old = self.y.copy()
+            #self.x_old = self.x.copy()
+        else:
+            # Gradient descent, Dual problem solution
+            self.y_old += self.sigma * self.operator.direct(self.xbar)
+            self.y = self.f.proximal_conjugate(self.y_old, self.sigma)
+            
+            # Gradient ascent, Primal problem solution
+            self.x_old -= self.tau * self.operator.adjoint(self.y)
+            self.x = self.g.proximal(self.x_old, self.tau)
+            
+            #Update
+            #xbar = x + theta * (x - x_old)
+            self.xbar.fill(self.x)
+            self.xbar -= self.x_old 
+            self.xbar *= self.theta
+            self.xbar += self.x
+                            
+            self.x_old.fill(self.x)
+            self.y_old.fill(self.y)
+            #self.y_old = self.y.copy()
+            #self.x_old = self.x.copy()
+            #self.y = self.y_old
 
     def update_objective(self):
-        self.loss.append([self.f(self.operator.direct(self.x)) + self.g(self.x),
-            -(self.f.convex_conjugate(self.y) + self.g.convex_conjugate(- 1 * self.operator.adjoint(self.y)))
-        ])
+        p1 = self.f(self.operator.direct(self.x)) + self.g(self.x)
+        d1 = -(self.f.convex_conjugate(self.y) + self.g(-1*self.operator.adjoint(self.y)))
 
-    
-    
-def assertBlockDataContainerEqual( container1, container2):
-    print ("assert Block Data Container Equal")
-    assert issubclass(container1.__class__, container2.__class__)
-    for col in range(container1.shape[0]):
-        if issubclass(container1.get_item(col).__class__, DataContainer):
-            print ("Checking col ", col)
-            assertNumpyArrayEqual(
-                container1.get_item(col).as_array(), 
-                container2.get_item(col).as_array()
-                )
-        else:
-            assertBlockDataContainerEqual(container1.get_item(col),container2.get_item(col))    
+        self.loss.append([p1,d1,p1-d1])
 
-def assertNumpyArrayEqual(first, second):
-    res = True
-    try:
-        numpy.testing.assert_array_equal(first, second)
-    except AssertionError as err:
-        res = False
-        print(err)
-    assert res
 
 
 def PDHG_old(f, g, operator, tau = None, sigma = None, opt = None, **kwargs):
@@ -122,7 +131,10 @@ def PDHG_old(f, g, operator, tau = None, sigma = None, opt = None, **kwargs):
     show_iter = opt['show_iter'] if 'show_iter' in opt.keys() else False 
     stop_crit = opt['stop_crit'] if 'stop_crit' in opt.keys() else False 
 
-
+    if memopt:
+        print ("memopt")
+    else:
+        print("no memopt")
     x_old = operator.domain_geometry().allocate()
     y_old = operator.range_geometry().allocate()       
             
@@ -131,7 +143,8 @@ def PDHG_old(f, g, operator, tau = None, sigma = None, opt = None, **kwargs):
     x = x_old.copy()
     
     y_tmp = y_old.copy()
-    y = y_old.copy()
+    y = y_tmp.copy()
+
         
     # relaxation parameter
     theta = 1
@@ -144,6 +157,7 @@ def PDHG_old(f, g, operator, tau = None, sigma = None, opt = None, **kwargs):
     
     
     for i in range(niter):
+
         
         if not memopt:
             
@@ -214,13 +228,6 @@ def PDHG_old(f, g, operator, tau = None, sigma = None, opt = None, **kwargs):
 ##        operator.direct(xbar, out = y_tmp)
 ##        y_tmp *= sigma
 ##        y_tmp +=y_old
-        
-        
-        
-        
-                
-
-
 #        if isinstance(f, FunctionOperatorComposition):
 #        p1 = f(x) + g(x)
 #        else:
diff --git a/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py b/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py
index 14105b5..9917d99 100644
--- a/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py
@@ -53,29 +53,27 @@ class BlockFunction(Function):
     
     def proximal_conjugate(self, x, tau, out = None):
         '''proximal_conjugate does not take into account the BlockOperator'''
-        
-        if out is None:
-        
-            out = [None]*self.length
+
+        if out is not None:
             if isinstance(tau, Number):
                 for i in range(self.length):
-                    out[i] = self.functions[i].proximal_conjugate(x.get_item(i), tau)
+                    self.functions[i].proximal_conjugate(x.get_item(i), tau, out=out.get_item(i))
             else:
                 for i in range(self.length):
-                    out[i] = self.functions[i].proximal_conjugate(x.get_item(i), tau.get_item(i))
-    
-            return BlockDataContainer(*out) 
-        
-        else:
+                    self.functions[i].proximal_conjugate(x.get_item(i), tau.get_item(i),out=out.get_item(i))
             
+        else:
+                
+            out = [None]*self.length
             if isinstance(tau, Number):
                 for i in range(self.length):
-                    self.functions[i].proximal_conjugate(x.get_item(i), tau, out = out.get_item(i))
+                    out[i] = self.functions[i].proximal_conjugate(x.get_item(i), tau)
             else:
                 for i in range(self.length):
-                    self.functions[i].proximal_conjugate(x.get_item(i), tau.get_item(i), out = out)
-               
-        
+                    out[i] = self.functions[i].proximal_conjugate(x.get_item(i), tau.get_item(i))
+            
+            return BlockDataContainer(*out) 
+
     
     def proximal(self, x, tau, out = None):
         '''proximal does not take into account the BlockOperator'''
diff --git a/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py b/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py
index d5e527a..9e0e424 100644
--- a/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py
@@ -69,10 +69,10 @@ class L2NormSquared(Function):
             out *= 2
         else:
             y = x
-        if self.b is not None:
-#            x.subtract(self.b, out=x)
-            y = x - self.b
-        return 2*y
+            if self.b is not None:
+    #            x.subtract(self.b, out=x)
+                y = x - self.b
+            return 2*y
         
                                                        
     def convex_conjugate(self, x, out=None):
diff --git a/Wrappers/Python/ccpi/optimisation/functions/MixedL21Norm.py b/Wrappers/Python/ccpi/optimisation/functions/MixedL21Norm.py
index 639f7bf..0ce7d8a 100755
--- a/Wrappers/Python/ccpi/optimisation/functions/MixedL21Norm.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/MixedL21Norm.py
@@ -82,57 +82,51 @@ class MixedL21Norm(Function):
         
         if self.SymTensor:
             
-            if out is None:
+            param = [1]*x.shape[0]
+            param[-1] = 2
+            tmp = [param[i]*(x[i] ** 2) for i in range(x.shape[0])]
+            frac = [x[i]/(sum(tmp).sqrt()).maximum(1.0) for i in range(x.shape[0])]
+            res = BlockDataContainer(*frac) 
             
-                param = [1]*x.shape[0]
-                param[-1] = 2
-                tmp = [param[i]*(x[i] ** 2) for i in range(x.shape[0])]
-                frac = [x[i]/(sum(tmp).sqrt()).maximum(1.0) for i in range(x.shape[0])]
-                res = BlockDataContainer(*frac) 
-                return res
-            else:
-                
-                
-                pass
-                        
-#                tmp2 = np.sqrt(x.as_array()[0]**2 +  x.as_array()[1]**2 +  2*x.as_array()[2]**2)/self.alpha
-#                res = x.divide(ImageData(tmp2).maximum(1.0))                                
+            return res
+        
         else:
-            
-            if out is None:
-                
-                tmp = [ el*el for el in x]
-                res = (sum(tmp).sqrt()).maximum(1.0) 
-                
-#                res = sum(x**2).sqrt().maximum(1.0) 
-                
-                #frac = [x[i]/res for i in range(x.shape[0])]
-                #res = BlockDataContainer(*frac)
-                
-                return x/res 
-            
-            else:
-                
-                tmp = x**2 #[ el*el for el in x]
-                res = (sum(tmp).sqrt()).maximum(1.0) 
-                
-#                res = sum(x**2).sqrt().maximum(1.0)
-#                frac = [x[i]/res for i in range(x.shape[0])]
-#                res = BlockDataContainer(*frac)
-#                res = sum(x**2).sqrt().maximum(1.0) 
-                out.fill(x/res)
-#                return res 
+#             pass
 
+    
+# #                tmp2 = np.sqrt(x.as_array()[0]**2 +  x.as_array()[1]**2 +  2*x.as_array()[2]**2)/self.alpha
+# #                res = x.divide(ImageData(tmp2).maximum(1.0))                                
+         if out is None:
+            
+            tmp = [ el*el for el in x]
+            res = (sum(tmp).sqrt()).maximum(1.0) 
+            frac = [x[i]/res for i in range(x.shape[0])]
+            res = BlockDataContainer(*frac)
+                                                   
+            return res
+        
+         else:
+             
+            tmp = [ el*el for el in x]
+            res = (sum(tmp).sqrt()).maximum(1.0)             
+            out.fill(x/res)             
+             
+             
+             
+        #     tmp = [ el*el for el in x]
+        #     res = (sum(tmp).sqrt()).maximum(1.0) 
+        #     #frac = [x[i]/res for i in range(x.shape[0])]
+        #     for i in range(x.shape[0]):
+        #         a = out.get_item(i)
+        #         b = x.get_item(i)
+        #         b /= res
+        #         a.fill( b )
 
     
     def __rmul__(self, scalar):
         return ScaledFunction(self, scalar) 
 
-#class MixedL21Norm_tensor(Function):
-#    
-#    def __init__(self):
-#        print("feerf")
-#    
+
 #
 if __name__ == '__main__':
     
diff --git a/Wrappers/Python/ccpi/optimisation/functions/ScaledFunction.py b/Wrappers/Python/ccpi/optimisation/functions/ScaledFunction.py
index c3d5ab9..464b944 100755
--- a/Wrappers/Python/ccpi/optimisation/functions/ScaledFunction.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/ScaledFunction.py
@@ -61,8 +61,7 @@ class ScaledFunction(object):
         if out is None:

             return self.scalar * self.function.proximal_conjugate(x/self.scalar, tau/self.scalar)

         else:

-            out.fill(self.scalar * self.function.proximal_conjugate(x/self.scalar, tau/self.scalar))

-            

+            out.fill(self.scalar*self.function.proximal_conjugate(x/self.scalar, tau/self.scalar))

 

     def grad(self, x):

         '''Alias of gradient(x,None)'''

diff --git a/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator.py b/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator.py
index 8e73ff2..835f96d 100644
--- a/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator.py
+++ b/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator.py
@@ -57,6 +57,7 @@ class FiniteDiff(LinearOperator):
             out[:]=0
 
                   
+
         ######################## Direct for 2D  ###############################
         if x_sz == 2:
             
@@ -65,7 +66,7 @@ class FiniteDiff(LinearOperator):
                 np.subtract( x_asarr[:,1:], x_asarr[:,0:-1], out = out[:,0:-1] )
                 
                 if self.bnd_cond == 'Neumann':
-                    pass                                        
+                    pass
                 elif self.bnd_cond == 'Periodic':
                     np.subtract( x_asarr[:,0], x_asarr[:,-1], out = out[:,-1] )
                 else: 
@@ -178,19 +179,8 @@ class FiniteDiff(LinearOperator):
             out = np.zeros_like(x_asarr)
         else:
             out = out.as_array()        
-        
-#        if out is None:        
-#            out = np.zeros_like(x_asarr)
-#            fd_arr = out
-#        else:
-#            fd_arr = out.as_array()          
-        
-#        if out is None:        
-#            out = self.gm_domain.allocate().as_array()
-#            fd_arr = out
-#        else:
-#            fd_arr = out.as_array()
-##        fd_arr = self.gm_domain.allocate().as_array()
+            out[:]=0
+
         
         ######################## Adjoint for 2D  ###############################
         if x_sz == 2:        
diff --git a/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py b/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py
index e535847..4935435 100644
--- a/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py
+++ b/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py
@@ -63,18 +63,21 @@ class Gradient(LinearOperator):
     def adjoint(self, x, out=None):
         
         if out is not None:
-          
-            tmp = self.gm_domain.allocate()   
-            
+
+            tmp = self.gm_domain.allocate()            
             for i in range(x.shape[0]):
                 self.FD.direction=self.ind[i] 
                 self.FD.adjoint(x.get_item(i), out = tmp)
-                out+=tmp               
+                if i == 0:
+                    out.fill(tmp)
+                else:
+                    out += tmp
         else:            
             tmp = self.gm_domain.allocate()
             for i in range(x.shape[0]):
                 self.FD.direction=self.ind[i]
-                tmp+=self.FD.adjoint(x.get_item(i))
+
+                tmp += self.FD.adjoint(x.get_item(i))
             return tmp    
             
     
diff --git a/Wrappers/Python/test/test_BlockDataContainer.py b/Wrappers/Python/test/test_BlockDataContainer.py
index 2ee0e94..2fca23c 100755
--- a/Wrappers/Python/test/test_BlockDataContainer.py
+++ b/Wrappers/Python/test/test_BlockDataContainer.py
@@ -14,7 +14,7 @@ from ccpi.optimisation.funcs import Norm2sq, Norm1
 from ccpi.framework import ImageGeometry, AcquisitionGeometry

 from ccpi.framework import ImageData, AcquisitionData

 #from ccpi.optimisation.algorithms import GradientDescent

-from ccpi.framework import BlockDataContainer

+from ccpi.framework import BlockDataContainer, DataContainer

 #from ccpi.optimisation.Algorithms import CGLS

 import functools

 

@@ -402,3 +402,52 @@ class TestBlockDataContainer(unittest.TestCase):
 

         c5 = d.get_item(0).power(2).sum()

 

+    def test_BlockDataContainer_fill(self):

+        print ("test block data container")

+        ig0 = ImageGeometry(2,3,4)

+        ig1 = ImageGeometry(2,3,5)

+        

+        data0 = ImageData(geometry=ig0)

+        data1 = ImageData(geometry=ig1) + 1

+        

+        data2 = ImageData(geometry=ig0) + 2

+        data3 = ImageData(geometry=ig1) + 3

+        

+        cp0 = BlockDataContainer(data0,data1)

+        #cp1 = BlockDataContainer(data2,data3)

+

+        cp2 = BlockDataContainer(data0+1, data1+1)

+

+        data0.fill(data2)

+        self.assertNumpyArrayEqual(data0.as_array(), data2.as_array())

+        data0 = ImageData(geometry=ig0)

+

+        for el,ot in zip(cp0, cp2):

+            print (el.shape, ot.shape)

+        cp0.fill(cp2)

+        self.assertBlockDataContainerEqual(cp0, cp2)

+

+

+    def assertBlockDataContainerEqual(self, container1, container2):

+        print ("assert Block Data Container Equal")

+        self.assertTrue(issubclass(container1.__class__, container2.__class__))

+        for col in range(container1.shape[0]):

+            if issubclass(container1.get_item(col).__class__, DataContainer):

+                print ("Checking col ", col)

+                self.assertNumpyArrayEqual(

+                    container1.get_item(col).as_array(), 

+                    container2.get_item(col).as_array()

+                    )

+            else:

+                self.assertBlockDataContainerEqual(container1.get_item(col),container2.get_item(col))

+

+    def assertNumpyArrayEqual(self, first, second):

+        res = True

+        try:

+            numpy.testing.assert_array_equal(first, second)

+        except AssertionError as err:

+            res = False

+            print(err)

+        self.assertTrue(res)    

+

+

diff --git a/Wrappers/Python/test/test_Operator.py b/Wrappers/Python/test/test_Operator.py
index 6656d34..293fb43 100644
--- a/Wrappers/Python/test/test_Operator.py
+++ b/Wrappers/Python/test/test_Operator.py
@@ -2,7 +2,8 @@ import unittest
 #from ccpi.optimisation.operators import Operator
 from ccpi.optimisation.ops import TomoIdentity
 from ccpi.framework import ImageGeometry, ImageData, BlockDataContainer, DataContainer
-from ccpi.optimisation.operators import BlockOperator, BlockScaledOperator
+from ccpi.optimisation.operators import BlockOperator, BlockScaledOperator,\
+    FiniteDiff
 import numpy
 from timeit import default_timer as timer
 from ccpi.framework import ImageGeometry
@@ -11,7 +12,43 @@ from ccpi.optimisation.operators import Gradient, Identity, SparseFiniteDiff
 def dt(steps):
     return steps[-1] - steps[-2]
 
-class TestOperator(unittest.TestCase):
+class CCPiTestClass(unittest.TestCase):
+    def assertBlockDataContainerEqual(self, container1, container2):
+        print ("assert Block Data Container Equal")
+        self.assertTrue(issubclass(container1.__class__, container2.__class__))
+        for col in range(container1.shape[0]):
+            if issubclass(container1.get_item(col).__class__, DataContainer):
+                print ("Checking col ", col)
+                self.assertNumpyArrayEqual(
+                    container1.get_item(col).as_array(), 
+                    container2.get_item(col).as_array()
+                    )
+            else:
+                self.assertBlockDataContainerEqual(container1.get_item(col),container2.get_item(col))
+    
+    def assertNumpyArrayEqual(self, first, second):
+        res = True
+        try:
+            numpy.testing.assert_array_equal(first, second)
+        except AssertionError as err:
+            res = False
+            print(err)
+        self.assertTrue(res)
+
+    def assertNumpyArrayAlmostEqual(self, first, second, decimal=6):
+        res = True
+        try:
+            numpy.testing.assert_array_almost_equal(first, second, decimal)
+        except AssertionError as err:
+            res = False
+            print(err)
+            print("expected " , second)
+            print("actual " , first)
+
+        self.assertTrue(res)
+
+
+class TestOperator(CCPiTestClass):
     def test_ScaledOperator(self):
         ig = ImageGeometry(10,20,30)
         img = ig.allocate()
@@ -29,6 +66,40 @@ class TestOperator(unittest.TestCase):
         y = Id.direct(img)
         numpy.testing.assert_array_equal(y.as_array(), img.as_array())
 
+    def test_FiniteDifference(self):
+        ##
+        N, M = 2, 3
+
+        ig = ImageGeometry(N, M)
+        Id = Identity(ig)
+
+        FD = FiniteDiff(ig, direction = 0, bnd_cond = 'Neumann')
+        u = FD.domain_geometry().allocate('random_int')
+        
+        
+        res = FD.domain_geometry().allocate(ImageGeometry.RANDOM_INT)
+        FD.adjoint(u, out=res)
+        w = FD.adjoint(u)
+
+        self.assertNumpyArrayEqual(res.as_array(), w.as_array())
+        
+        res = Id.domain_geometry().allocate(ImageGeometry.RANDOM_INT)
+        Id.adjoint(u, out=res)
+        w = Id.adjoint(u)
+
+        self.assertNumpyArrayEqual(res.as_array(), w.as_array())
+        self.assertNumpyArrayEqual(u.as_array(), w.as_array())
+
+        G = Gradient(ig)
+
+        u = G.range_geometry().allocate(ImageGeometry.RANDOM_INT)
+        res = G.domain_geometry().allocate(ImageGeometry.RANDOM_INT)
+        G.adjoint(u, out=res)
+        w = G.adjoint(u)
+        self.assertNumpyArrayEqual(res.as_array(), w.as_array())
+        
+
+
 
 
 class TestBlockOperator(unittest.TestCase):
@@ -90,22 +161,23 @@ class TestBlockOperator(unittest.TestCase):
         print (z1.shape)
         print(z1[0][0].as_array())
         print(res[0][0].as_array())   
+        self.assertBlockDataContainerEqual(z1, res)
+        # for col in range(z1.shape[0]):
+        #     a = z1.get_item(col)
+        #     b = res.get_item(col)
+        #     if isinstance(a, BlockDataContainer):
+        #         for col2 in range(a.shape[0]):
+        #             self.assertNumpyArrayEqual(
+        #                 a.get_item(col2).as_array(), 
+        #                 b.get_item(col2).as_array()
+        #                 )        
+        #     else:
+        #         self.assertNumpyArrayEqual(
+        #             a.as_array(), 
+        #             b.as_array()
+        #             )
+        z1 = B.range_geometry().allocate(ImageGeometry.RANDOM_INT)
 
-        for col in range(z1.shape[0]):
-            a = z1.get_item(col)
-            b = res.get_item(col)
-            if isinstance(a, BlockDataContainer):
-                for col2 in range(a.shape[0]):
-                    self.assertNumpyArrayEqual(
-                        a.get_item(col2).as_array(), 
-                        b.get_item(col2).as_array()
-                        )        
-            else:
-                self.assertNumpyArrayEqual(
-                    a.as_array(), 
-                    b.as_array()
-                    )
-        z1 = B.direct(u)
         res1 = B.adjoint(z1)
         res2 = B.domain_geometry().allocate()
         B.adjoint(z1, out=res2)
@@ -264,7 +336,7 @@ class TestBlockOperator(unittest.TestCase):
         u = ig.allocate('random_int')
         steps = [timer()]
         i = 0
-        n = 25.
+        n = 2.
         t1 = t2 = 0
         res = B.range_geometry().allocate()
             
diff --git a/Wrappers/Python/test/test_functions.py b/Wrappers/Python/test/test_functions.py
index 19cb65f..1891afd 100644
--- a/Wrappers/Python/test/test_functions.py
+++ b/Wrappers/Python/test/test_functions.py
@@ -65,6 +65,9 @@ class TestFunction(unittest.TestCase):
         a3 = 0.5 * d.squared_norm() + d.dot(noisy_data)
         self.assertEqual(a3, g.convex_conjugate(d))
         #print( a3, g.convex_conjugate(d))
+
+        #test proximal conjugate
+
     
     def test_L2NormSquared(self):
         # TESTS for L2 and scalar * L2
@@ -94,7 +97,7 @@ class TestFunction(unittest.TestCase):
         c2 = 1/4. * u.squared_norm()
         numpy.testing.assert_equal(c1, c2)
         
-        #check convex conjuagate with data
+        #check convex conjugate with data
         d1 = f1.convex_conjugate(u)
         d2 = (1./4.) * u.squared_norm() + (u*b).sum()
         numpy.testing.assert_equal(d1, d2)  
@@ -121,10 +124,9 @@ class TestFunction(unittest.TestCase):
         l1 = f1.proximal_conjugate(u, tau)
         l2 = (u - tau * b)/(1 + tau/2 )
         numpy.testing.assert_array_almost_equal(l1.as_array(), l2.as_array(), decimal=4)     
-        
-            
+
         # check scaled function properties
-        
+
         # scalar 
         scalar = 100
         f_scaled_no_data = scalar * L2NormSquared()
@@ -161,7 +163,105 @@ class TestFunction(unittest.TestCase):
         numpy.testing.assert_array_almost_equal(f_scaled_data.proximal_conjugate(u, tau).as_array(), \
                                                 ((u - tau * b)/(1 + tau/(2*scalar) )).as_array(), decimal=4)    
         
+    def test_L2NormSquaredOut(self):
+        # TESTS for L2 and scalar * L2
+    
+        M, N, K = 2,3,5
+        ig = ImageGeometry(voxel_num_x=M, voxel_num_y = N, voxel_num_z = K)
+        u = ig.allocate(ImageGeometry.RANDOM_INT)
+        b = ig.allocate(ImageGeometry.RANDOM_INT) 
+        
+        # check grad/call no data
+        f = L2NormSquared()
+        a1 = f.gradient(u)
+        a2 = a1 * 0.
+        f.gradient(u, out=a2)
+        numpy.testing.assert_array_almost_equal(a1.as_array(), a2.as_array(), decimal=4)
+        #numpy.testing.assert_equal(f(u), u.squared_norm())
+
+        # check grad/call with data
+        f1 = L2NormSquared(b=b)
+        b1 = f1.gradient(u)
+        b2 = b1 * 0.
+        f1.gradient(u, out=b2)
+            
+        numpy.testing.assert_array_almost_equal(b1.as_array(), b2.as_array(), decimal=4)
+        #numpy.testing.assert_equal(f1(u), (u-b).squared_norm())
+        
+        # check proximal no data
+        tau = 5
+        e1 = f.proximal(u, tau)
+        e2 = e1 * 0.
+        f.proximal(u, tau, out=e2)
+        numpy.testing.assert_array_almost_equal(e1.as_array(), e2.as_array(), decimal=4)
+        
+        # check proximal with data
+        tau = 5
+        h1 = f1.proximal(u, tau)
+        h2 = h1 * 0.
+        f1.proximal(u, tau, out=h2)
+        numpy.testing.assert_array_almost_equal(h1.as_array(), h2.as_array(), decimal=4)    
+        
+        # check proximal conjugate no data
+        tau = 0.2
+        k1 = f.proximal_conjugate(u, tau)
+        k2 = k1 * 0.
+        f.proximal_conjugate(u, tau, out=k2)
+
+        numpy.testing.assert_array_almost_equal(k1.as_array(), k2.as_array(), decimal=4) 
+        
+        # check proximal conjugate with data
+        l1 = f1.proximal_conjugate(u, tau)
+        l2 = l1 * 0.
+        f1.proximal_conjugate(u, tau, out=l2)
+        numpy.testing.assert_array_almost_equal(l1.as_array(), l2.as_array(), decimal=4)     
+
+        # check scaled function properties
+
+        # scalar 
+        scalar = 100
+        f_scaled_no_data = scalar * L2NormSquared()
+        f_scaled_data = scalar * L2NormSquared(b=b)
+        
+        # grad
+        w = f_scaled_no_data.gradient(u)
+        ww = w * 0
+        f_scaled_no_data.gradient(u, out=ww)
+
+        numpy.testing.assert_array_almost_equal(w.as_array(), 
+            ww.as_array(), decimal=4)
+
+        # numpy.testing.assert_array_almost_equal(f_scaled_data.gradient(u).as_array(), scalar*f1.gradient(u).as_array(), decimal=4)
+        
+        # # conj
+        # numpy.testing.assert_almost_equal(f_scaled_no_data.convex_conjugate(u), \
+        #                         f.convex_conjugate(u/scalar) * scalar, decimal=4)
+        
+        # numpy.testing.assert_almost_equal(f_scaled_data.convex_conjugate(u), \
+        #                         scalar * f1.convex_conjugate(u/scalar), decimal=4)
+        
+        # # proximal
+        w = f_scaled_no_data.proximal(u, tau)
+        ww = w * 0
+        f_scaled_no_data.proximal(u, tau, out=ww)
+        numpy.testing.assert_array_almost_equal(w.as_array(), \
+                                                ww.as_array())
+        
+        
+        # numpy.testing.assert_array_almost_equal(f_scaled_data.proximal(u, tau).as_array(), \
+        #                                         f1.proximal(u, tau*scalar).as_array())
+                                
+        
+        # proximal conjugate
+        w = f_scaled_no_data.proximal_conjugate(u, tau)
+        ww = w * 0
+        f_scaled_no_data.proximal_conjugate(u, tau, out=ww)
+        numpy.testing.assert_array_almost_equal(w.as_array(), \
+                                                ww.as_array(), decimal=4)
         
+        # numpy.testing.assert_array_almost_equal(f_scaled_data.proximal_conjugate(u, tau).as_array(), \
+        #                                         ((u - tau * b)/(1 + tau/(2*scalar) )).as_array(), decimal=4)    
+           
     def test_Norm2sq_as_FunctionOperatorComposition(self):
         M, N, K = 2,3,5
         ig = ImageGeometry(voxel_num_x=M, voxel_num_y = N, voxel_num_z = K)
diff --git a/Wrappers/Python/wip/pdhg_TV_denoising.py b/Wrappers/Python/wip/pdhg_TV_denoising.py
index d0cb198..598acb0 100755
--- a/Wrappers/Python/wip/pdhg_TV_denoising.py
+++ b/Wrappers/Python/wip/pdhg_TV_denoising.py
@@ -19,12 +19,16 @@ from ccpi.optimisation.functions import ZeroFun, L2NormSquared, \
 
 from skimage.util import random_noise
 
+from timeit import default_timer as timer
+def dt(steps):
+    return steps[-1] - steps[-2]
 
 
 # ############################################################################
 # Create phantom for TV denoising
 
-N = 500
+N = 100
+
 data = np.zeros((N,N))
 data[round(N/4):round(3*N/4),round(N/4):round(3*N/4)] = 0.5
 data[round(N/8):round(7*N/8),round(3*N/8):round(5*N/8)] = 1
@@ -36,8 +40,8 @@ ag = ig
 n1 = random_noise(data, mode = 'gaussian', mean=0, var = 0.05, seed=10)
 noisy_data = ImageData(n1)
 
-plt.imshow(noisy_data.as_array())
-plt.show()
+#plt.imshow(noisy_data.as_array())
+#plt.show()
 
 #%%
 
@@ -84,6 +88,7 @@ print ("normK", normK)
 sigma = 1
 tau = 1/(sigma*normK**2)
 
+<<<<<<< HEAD
 opt = {'niter':100}
 opt1 = {'niter':100, 'memopt': True}
 
@@ -108,34 +113,73 @@ plt.figure(figsize=(5,5))
 plt.imshow(np.abs(res1.as_array()-res.as_array()))
 plt.colorbar()
 plt.show()
-
+#=======
+## opt = {'niter':2000, 'memopt': True}
+#
+## res, time, primal, dual, pdgap = PDHG_old(f, g, operator, tau = tau, sigma = sigma, opt = opt) 
+# 
+#>>>>>>> origin/pdhg_fix
+#
+#
+## opt = {'niter':2000, 'memopt': False}
+## res1, time1, primal1, dual1, pdgap1 = PDHG_old(f, g, operator, tau = tau, sigma = sigma, opt = opt) 
+#
+## plt.figure(figsize=(5,5))
+## plt.subplot(1,3,1)
+## plt.imshow(res.as_array())
+## plt.title('memopt')
+## plt.colorbar()
+## plt.subplot(1,3,2)
+## plt.imshow(res1.as_array())
+## plt.title('no memopt')
+## plt.colorbar()
+## plt.subplot(1,3,3)
+## plt.imshow((res1 - res).abs().as_array())
+## plt.title('diff')
+## plt.colorbar()
+## plt.show()
 #pdhg = PDHG(f=f,g=g,operator=operator, tau=tau, sigma=sigma)
 #pdhg.max_iteration = 2000
-#pdhg.update_objective_interval = 10
+#pdhg.update_objective_interval = 100
+#
 #
-#pdhg.run(2000)
+#pdhgo = PDHG(f=f,g=g,operator=operator, tau=tau, sigma=sigma, memopt=True)
+#pdhgo.max_iteration = 2000
+#pdhgo.update_objective_interval = 100
 #
-#    
+#steps = [timer()]
+#pdhgo.run(200)
+#steps.append(timer())
+#t1 = dt(steps)
 #
+#pdhg.run(200)
+#steps.append(timer())
+#t2 = dt(steps)
+#
+#print ("Time difference {} {} {}".format(t1,t2,t2-t1))
 #sol = pdhg.get_output().as_array()
 ##sol = result.as_array()
 ##
 #fig = plt.figure()
-#plt.subplot(1,2,1)
+#plt.subplot(1,3,1)
 #plt.imshow(noisy_data.as_array())
-##plt.colorbar()
-#plt.subplot(1,2,2)
+#plt.colorbar()
+#plt.subplot(1,3,2)
 #plt.imshow(sol)
-##plt.colorbar()
+#plt.colorbar()
+#plt.subplot(1,3,3)
+#plt.imshow(pdhgo.get_output().as_array())
+#plt.colorbar()
+#
 #plt.show()
+###
 ##
+####
+##plt.plot(np.linspace(0,N,N), data[int(N/2),:], label = 'GTruth')
+##plt.plot(np.linspace(0,N,N), sol[int(N/2),:], label = 'Recon')
+##plt.legend()
+##plt.show()
 #
-###
-#plt.plot(np.linspace(0,N,N), data[int(N/2),:], label = 'GTruth')
-#plt.plot(np.linspace(0,N,N), sol[int(N/2),:], label = 'Recon')
-#plt.legend()
-#plt.show()
-
-
-#%%
 #
+##%%
+##