Merge remote-tracking branch 'origin/composite_operator_datacontainer' into pdhg_fix

15 files changed, 1031 insertions, 331 deletions

diff --git a/Wrappers/Python/ccpi/framework/framework.py b/Wrappers/Python/ccpi/framework/framework.py
index 2453986..af4139b 100755
--- a/Wrappers/Python/ccpi/framework/framework.py
+++ b/Wrappers/Python/ccpi/framework/framework.py
@@ -997,6 +997,7 @@ class AcquisitionData(DataContainer):
                 
             
 class DataProcessor(object):
+    
     '''Defines a generic DataContainer processor
     
     accepts DataContainer as inputs and 
@@ -1042,6 +1043,7 @@ class DataProcessor(object):
         raise NotImplementedError('Implement basic checks for input DataContainer')
         
     def get_output(self, out=None):
+        
         for k,v in self.__dict__.items():
             if v is None and k != 'output':
                 raise ValueError('Key {0} is None'.format(k))
diff --git a/Wrappers/Python/ccpi/optimisation/algorithms/PDHG.py b/Wrappers/Python/ccpi/optimisation/algorithms/PDHG.py
index 835c979..439149c 100644
--- a/Wrappers/Python/ccpi/optimisation/algorithms/PDHG.py
+++ b/Wrappers/Python/ccpi/optimisation/algorithms/PDHG.py
@@ -8,6 +8,7 @@ Created on Mon Feb  4 16:18:06 2019
 from ccpi.optimisation.algorithms import Algorithm
 from ccpi.framework import ImageData, DataContainer
 import numpy as np
+import numpy
 import time
 from ccpi.optimisation.operators import BlockOperator
 from ccpi.framework import BlockDataContainer
@@ -108,26 +109,6 @@ class PDHG(Algorithm):
         self.loss.append([p1,d1,p1-d1])
 
 
-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):
-            assertNumpyArrayEqual(
-                container1.get_item(col).as_array(), 
-                container2.get_item(col).as_array()
-                )
-        else:
-            assertBlockDataContainerEqual(container1.get_item(col),container2.get_item(col))
-
-def assertNumpyArrayEqual(first, second):
-    res = True
-    try:
-        np.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):
         
@@ -158,6 +139,7 @@ def PDHG_old(f, g, operator, tau = None, sigma = None, opt = None, **kwargs):
     
     y_tmp = y_old.copy()
     y = y_tmp.copy()
+
         
     # relaxation parameter
     theta = 1
@@ -170,79 +152,79 @@ def PDHG_old(f, g, operator, tau = None, sigma = None, opt = None, **kwargs):
     
     
     for i in range(niter):
-        if memopt:
-    #        # Gradient descent, Dual problem solution
-    #        y_tmp = y_old + sigma * operator.direct(xbar)
-            #y_tmp = operator.direct(xbar)
-            operator.direct(xbar, out=y_tmp)
-            y_tmp *= sigma
-            y_tmp +=y_old        
-                            
-            y = f.proximal_conjugate(y_tmp, sigma)
-            #f.proximal_conjugate(y_tmp, sigma, out=y)
 
-            # Gradient ascent, Primal problem solution
-    #        x_tmp = x_old - tau * operator.adjoint(y)
-            
-            #x_tmp = operator.adjoint(y)
-            operator.adjoint(y, out=x_tmp)
-            x_tmp *=-tau
-            x_tmp +=x_old
+        
+        if not memopt:
             
-            #x = g.proximal(x_tmp, tau)
-            g.proximal(x_tmp, tau, out=x)
+            y_old += sigma * operator.direct(xbar)            
+            y = f.proximal_conjugate(y_old, sigma)
             
-            #Update
-    #        xbar = x + theta * (x - x_old)
-            x.subtract(x_old, out=xbar)
+            x_old -= tau*operator.adjoint(y)
+            x = g.proximal(x_old, tau)
+             
+            xbar.fill(x)
+            xbar -= x_old
             xbar *= theta
             xbar += x
-                                    
+            
             x_old.fill(x)
-            y_old.fill(y)
+            y_old.fill(y)            
+                               
         else:
             
-    #        # Gradient descent, Dual problem solution
-            y_tmp1 = y_old + sigma * operator.direct(xbar)
-            # y_tmp = operator.direct(xbar)
-            operator.direct(xbar, out=y_tmp)
+            operator.direct(xbar, out = y_tmp)             
             y_tmp *= sigma
-            y_tmp +=y_old
-            #print ("y_tmp1 equale y_tmp?")
-            #assertBlockDataContainerEqual(y_tmp1, y_tmp)
+            y_tmp += y_old                      
+            f.proximal_conjugate(y_tmp, sigma, out=y)
 
-            y = f.proximal_conjugate(y_tmp, sigma)
-            #f.proximal_conjugate(y_tmp, sigma, out=y)
-            #print ("y1 equale y?")
-            #assertBlockDataContainerEqual(y1, y)
-            # Gradient ascent, Primal problem solution
-            x_tmp1 = x_old - tau * operator.adjoint(y)
-            
-            # x_tmp = operator.adjoint(y)
-            operator.adjoint(y, out=x_tmp)
-            x_tmp *=-tau
-            x_tmp +=x_old
-            
-            assertNumpyArrayEqual(x_tmp.as_array(),x_tmp1.as_array())
+            operator.adjoint(y, out = x_tmp)   
+            x_tmp *= -tau
+            x_tmp += x_old
 
-            x = g.proximal(x_tmp, tau)
-            # g.proximal(x_tmp, tau, out=x)
+            g.proximal(x_tmp, tau, out = x)
             
-            #Update
-            xbar = x + theta * (x - x_old)
-            # xbar = x - x_old
-            # xbar *= theta
-            # xbar += x
-                                    
-            x_old = x
-            y_old = y
-        
-        
-        
-        
-                
-
+            xbar = x - x_old
+            xbar *= theta
+            xbar += x
+                        
+                        
+            x_old.fill(x)
+            y_old.fill(y)
 
+            
+#            pass
+#        
+##        # Gradient descent, Dual problem solution
+##        y_tmp = y_old + sigma * operator.direct(xbar)
+#        y_tmp = operator.direct(xbar)
+#        y_tmp *= sigma
+#        y_tmp +=y_old        
+#                        
+#        y = f.proximal_conjugate(y_tmp, sigma)
+##        f.proximal_conjugate(y_tmp, sigma, out = y)
+#
+#        # Gradient ascent, Primal problem solution
+##        x_tmp = x_old - tau * operator.adjoint(y)
+#        
+#        x_tmp = operator.adjoint(y)
+#        x_tmp *=-tau
+#        x_tmp +=x_old
+#        
+#        x = g.proximal(x_tmp, tau)
+##        g.proximal(x_tmp, tau, out = x)
+#        
+#        #Update
+##        xbar = x + theta * (x - x_old)
+#        xbar = x - x_old
+#        xbar *= theta
+#        xbar += x
+#                                
+#        x_old = x
+#        y_old = y
+#        
+##        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 bbf1d29..8cce290 100644
--- a/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py
@@ -7,12 +7,13 @@ Created on Fri Mar  8 10:01:31 2019
 """
 
 import numpy as np
-#from ccpi.optimisation.funcs import Function
+
 from ccpi.optimisation.functions import Function
 from ccpi.framework import BlockDataContainer
 from numbers import Number
 
 class BlockFunction(Function):
+    
     '''A Block vector of Functions
     
     .. math::
@@ -52,6 +53,7 @@ class BlockFunction(Function):
     
     def proximal_conjugate(self, x, tau, out = None):
         '''proximal_conjugate does not take into account the BlockOperator'''
+
         if out is not None:
             if isinstance(tau, Number):
                 for i in range(self.length):
@@ -71,6 +73,7 @@ class BlockFunction(Function):
                     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'''
@@ -86,4 +89,96 @@ class BlockFunction(Function):
     
     def gradient(self,x, out=None):
         '''FIXME: gradient returns pass'''
-        pass
\ No newline at end of file
+        pass
+    
+    
+if __name__ == '__main__':
+    
+    M, N, K = 2,3,5
+    
+    from ccpi.optimisation.functions import L2NormSquared, MixedL21Norm
+    from ccpi.framework import ImageGeometry, BlockGeometry
+    from ccpi.optimisation.operators import Gradient, Identity, BlockOperator
+    import numpy
+    
+    
+    ig = ImageGeometry(M, N)
+    BG = BlockGeometry(ig, ig)
+    
+    u = ig.allocate('random_int')
+    B = BlockOperator( Gradient(ig), Identity(ig) )
+    
+    U = B.direct(u)
+    b = ig.allocate('random_int')
+    
+    f1 =  10 * MixedL21Norm()
+    f2 =  0.5 * L2NormSquared(b=b)    
+    
+    f = BlockFunction(f1, f2)
+    tau = 0.3
+    
+    print( " without out " )
+    res_no_out = f.proximal_conjugate( U, tau)
+    res_out = B.range_geometry().allocate()
+    f.proximal_conjugate( U, tau, out = res_out)
+    
+    numpy.testing.assert_array_almost_equal(res_no_out[0][0].as_array(), \
+                                            res_out[0][0].as_array(), decimal=4) 
+    
+    numpy.testing.assert_array_almost_equal(res_no_out[0][1].as_array(), \
+                                            res_out[0][1].as_array(), decimal=4) 
+
+    numpy.testing.assert_array_almost_equal(res_no_out[1].as_array(), \
+                                            res_out[1].as_array(), decimal=4)     
+    
+    
+    
+    
+    
+    
+    
+    ##########################################################################
+    
+    
+    
+    
+    
+    
+    
+#    zzz = B.range_geometry().allocate('random_int')
+#    www = B.range_geometry().allocate()
+#    www.fill(zzz)
+    
+#    res[0].fill(z)
+    
+    
+    
+    
+#    f.proximal_conjugate(z, sigma, out = res)
+    
+#    print(z1[0][0].as_array())    
+#    print(res[0][0].as_array())
+    
+    
+    
+    
+#    U = BG.allocate('random_int')
+#    RES = BG.allocate()
+#    f = BlockFunction(f1, f2)
+#    
+#    z = f.proximal_conjugate(U, 0.2)
+#    f.proximal_conjugate(U, 0.2, out = RES)
+#    
+#    print(z[0].as_array())
+#    print(RES[0].as_array())
+#    
+#    print(z[1].as_array())
+#    print(RES[1].as_array())    
+    
+    
+    
+    
+    
+
+
+    
\ No newline at end of file
diff --git a/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py b/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py
index 3c06641..903dafa 100644
--- a/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py
@@ -20,38 +20,31 @@
 import numpy
 from ccpi.optimisation.functions import Function
 from ccpi.optimisation.functions.ScaledFunction import ScaledFunction
-from ccpi.framework import DataContainer, ImageData, ImageGeometry  
 
-############################   L2NORM FUNCTION   #############################
 class L2NormSquared(Function):
     
-    def __init__(self, **kwargs):
-        
-        ''' L2NormSquared class
-            f : ImageGeometry --> R
-            
-            Cases: f(x) = ||x||^{2}_{2}
-                   f(x) = || x - b ||^{2}_{2}     
-        
-        '''
+    ''' 
+    
+    Class: L2NormSquared 
         
-        #TODO need x, b to live in the same geometry if b is not None
-                        
+    Cases:  a) f(x) = ||x||^{2}
+    
+            b) f(x) = ||x - b||^{2}, b
+                             
+    '''    
+    
+    def __init__(self, **kwargs):
+                                
         super(L2NormSquared, self).__init__()
         self.b = kwargs.get('b',None)  
 
     def __call__(self, x):
-        ''' Evaluates L2NormSq at point x'''
         
+        ''' Evaluate L2NormSquared at x: f(x) '''
+            
         y = x
         if self.b is not None: 
-#            x.subtract(self.b, out = x)
             y = x - self.b
-#        else:
-#            y
-#        if out is None:
-#            return x.squared_norm()
-#        else:
         try:
             return y.squared_norm()
         except AttributeError as ae:
@@ -60,41 +53,43 @@ class L2NormSquared(Function):
         
             
         
-    def gradient(self, x, out=None):
-        ''' Evaluates gradient of L2NormSq at point x'''
+    def gradient(self, x, out=None):        
+        
+        ''' Evaluate gradient of L2NormSquared at x: f'(x) '''
+                
         if out is not None:
+            
             out.fill(x)
             if self.b is not None:
                 out -= self.b
             out *= 2
+            
         else:
+            
             y = x
             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):
-        ''' Evaluate convex conjugate of L2NormSq'''
+    def convex_conjugate(self, x):
+        
+        ''' Evaluate convex conjugate of L2NormSquared at x: f^{*}(x)'''
             
         tmp = 0
+        
         if self.b is not None:
-#            tmp = (self.b * x).sum()
             tmp = (x * self.b).sum()
             
-        if out is None:
-            # FIXME: this is a number
-            return (1./4.) * x.squared_norm() + tmp
-        else:
-            # FIXME: this is a DataContainer
-            out.fill((1./4.) * x.squared_norm() + tmp)
-                    
+        return (1./4.) * x.squared_norm() + tmp
+
 
     def proximal(self, x, tau, out = None):
 
-        ''' The proximal operator ( prox_\{tau * f\}(x) ) evaluates i.e., 
-                argmin_x { 0.5||x - u||^{2} + tau f(x) }
+        ''' Evaluate Proximal Operator of tau * f(\cdot) at x:
+        
+            prox_{tau*f(\cdot)}(x) = \argmin_{z} \frac{1}{2}|| z - x ||^{2}_{2} + tau * f(z)
+        
         '''        
         
         if out is None:
@@ -108,27 +103,37 @@ class L2NormSquared(Function):
                 out -= self.b
             out /= (1+2*tau)
             if self.b is not None:
-                out += self.b
-                #out.fill((x - self.b)/(1+2*tau) + self.b)
-            #else:
-            #    out.fill(x/(1+2*tau))                
+                out += self.b             
 
     
     def proximal_conjugate(self, x, tau, out=None):
         
+        ''' Evaluate Proximal Operator of tau * f^{*}(\cdot) at x (i.e., the convex conjugate of f) :
+        
+            prox_{tau*f(\cdot)}(x) = \argmin_{z} \frac{1}{2}|| z - x ||^{2}_{2} + tau * f^{*}(z)
+        
+        '''         
+        
         if out is None:
             if self.b is not None:
-                # change the order cannot add ImageData + NestedBlock
                 return (x - tau*self.b)/(1 + tau/2) 
             else:
-                return x/(1 + tau/2 )
+                return x/(1 + tau/2)
         else:
             if self.b is not None:
-                out.fill((x - tau*self.b)/(1 + tau/2))
+                out.fill( (x - tau*self.b)/(1 + tau/2) )
             else:
-                out.fill(x/(1 + tau/2 ))
+                out.fill( x/(1 + tau/2) )
                                         
     def __rmul__(self, scalar):
+        
+        ''' Allows multiplication of L2NormSquared with a scalar
+        
+        Returns: ScaledFunction
+        
+                
+        '''
+        
         return ScaledFunction(self, scalar)        
 
 
@@ -227,7 +232,29 @@ if __name__ == '__main__':
                                             (u/(1 + tau/(2*scalar) )).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)    
+                                            ((u - tau * b)/(1 + tau/(2*scalar) )).as_array(), decimal=4)   
+    
+    
+    
+    print( " ####### check without out ######### " )
+          
+          
+    u_out_no_out = ig.allocate('random_int')         
+    res_no_out = f_scaled_data.proximal_conjugate(u_out_no_out, 0.5)          
+    print(res_no_out.as_array())
+    
+    print( " ####### check with out ######### " ) 
+          
+    res_out = ig.allocate()        
+    f_scaled_data.proximal_conjugate(u_out_no_out, 0.5, out = res_out)
+    
+    print(res_out.as_array())   
+
+    numpy.testing.assert_array_almost_equal(res_no_out.as_array(), \
+                                            res_out.as_array(), decimal=4)        
+          
+          
+        
     
     
     
diff --git a/Wrappers/Python/ccpi/optimisation/functions/MixedL21Norm.py b/Wrappers/Python/ccpi/optimisation/functions/MixedL21Norm.py
index c5be084..24c47f4 100755
--- a/Wrappers/Python/ccpi/optimisation/functions/MixedL21Norm.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/MixedL21Norm.py
@@ -79,7 +79,7 @@ class MixedL21Norm(Function):
         pass
     
     def proximal_conjugate(self, x, tau, out=None): 
-        
+
         if self.SymTensor:
             
             param = [1]*x.shape[0]
@@ -89,7 +89,9 @@ class MixedL21Norm(Function):
             res = BlockDataContainer(*frac) 
             
             return res
+        
         else:
+
             if out is None:                                        
                 tmp = [ el*el for el in x.containers]
                 res = sum(tmp).sqrt().maximum(1.0) 
@@ -106,20 +108,19 @@ class MixedL21Norm(Function):
     def __rmul__(self, scalar):
         return ScaledFunction(self, scalar) 
 
-#class MixedL21Norm_tensor(Function):
-#    
-#    def __init__(self):
-#        print("feerf")
-#    
+
 #
 if __name__ == '__main__':
     
     M, N, K = 2,3,5
-    ig = ImageGeometry(voxel_num_x=M, voxel_num_y = N)
-    u1 = ig.allocate('random_int')
-    u2 = ig.allocate('random_int')
+    from ccpi.framework import BlockGeometry
+    import numpy
+    
+    ig = ImageGeometry(M, N)
+    
+    BG = BlockGeometry(ig, ig)
     
-    U = BlockDataContainer(u1, u2, shape=(2,1))
+    U = BG.allocate('random_int')
     
     # Define no scale and scaled
     f_no_scaled = MixedL21Norm() 
@@ -129,9 +130,31 @@ if __name__ == '__main__':
     
     a1 = f_no_scaled(U)
     a2 = f_scaled(U)    
+    print(a1, 2*a2)
+        
+    
+    print( " ####### check without out ######### " )
+          
+          
+    u_out_no_out = BG.allocate('random_int')         
+    res_no_out = f_scaled.proximal_conjugate(u_out_no_out, 0.5)          
+    print(res_no_out[0].as_array())
+    
+    print( " ####### check with out ######### " ) 
+#          
+    res_out = BG.allocate()        
+    f_scaled.proximal_conjugate(u_out_no_out, 0.5, out = res_out)
+#    
+    print(res_out[0].as_array())   
+#
+    numpy.testing.assert_array_almost_equal(res_no_out[0].as_array(), \
+                                            res_out[0].as_array(), decimal=4)
+
+    numpy.testing.assert_array_almost_equal(res_no_out[1].as_array(), \
+                                            res_out[1].as_array(), decimal=4)     
+#    
+    
     
-    z1 = f_no_scaled.proximal_conjugate(U, 1)
-    z2 = f_scaled.proximal_conjugate(U, 1)
     
 
     
diff --git a/Wrappers/Python/ccpi/optimisation/functions/ScaledFunction.py b/Wrappers/Python/ccpi/optimisation/functions/ScaledFunction.py
index 9e2ba0c..464b944 100755
--- a/Wrappers/Python/ccpi/optimisation/functions/ScaledFunction.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/ScaledFunction.py
@@ -59,11 +59,9 @@ class ScaledFunction(object):
         '''This returns the proximal operator for the function at x, tau

         '''

         if out is None:

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

+            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))

-            #self.function.proximal_conjugate(x/self.scalar, tau/self.scalar, out=out)

-            #out *= self.scalar

 

     def grad(self, x):

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

@@ -91,3 +89,59 @@ class ScaledFunction(object):
             return self.function.proximal(x, tau*self.scalar)     

         else:

             out.fill( self.function.proximal(x, tau*self.scalar) )

+            

+if __name__ == '__main__':        

+

+    from ccpi.optimisation.functions import L2NormSquared, MixedL21Norm

+    from ccpi.framework import ImageGeometry, BlockGeometry

+    

+    M, N, K = 2,3,5

+    ig = ImageGeometry(voxel_num_x=M, voxel_num_y = N, voxel_num_z = K)

+    

+    u = ig.allocate('random_int')

+    b = ig.allocate('random_int') 

+    

+    BG = BlockGeometry(ig, ig)

+    U = BG.allocate('random_int')

+        

+    f2 = 0.5 * L2NormSquared(b=b)

+    f1 = 30 * MixedL21Norm()

+    tau = 0.355

+    

+    res_no_out1 =  f1.proximal_conjugate(U, tau)

+    res_no_out2 =  f2.proximal_conjugate(u, tau)    

+    

+    

+#    print( " ######## with out ######## ")

+    res_out1 = BG.allocate()

+    res_out2 = ig.allocate()

+    

+    f1.proximal_conjugate(U, tau, out = res_out1)

+    f2.proximal_conjugate(u, tau, out = res_out2)

+

+

+    numpy.testing.assert_array_almost_equal(res_no_out1[0].as_array(), \

+                                            res_out1[0].as_array(), decimal=4) 

+    

+    numpy.testing.assert_array_almost_equal(res_no_out2.as_array(), \

+                                            res_out2.as_array(), decimal=4)     

+      

+          

+          

+    

+    

+    

+    

+    

+    

+    

+    

+    

+    

+    

+    

+

+    

+

+

+    

diff --git a/Wrappers/Python/ccpi/optimisation/functions/ZeroFun.py b/Wrappers/Python/ccpi/optimisation/functions/ZeroFun.py
index 88d9b64..6d21acb 100644
--- a/Wrappers/Python/ccpi/optimisation/functions/ZeroFun.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/ZeroFun.py
@@ -45,14 +45,17 @@ class ZeroFun(Function):
             else:
                 return x.maximum(0).sum() + x.maximum(0).sum()
     
-    def proximal(self,x,tau, out=None):
+    def proximal(self, x, tau, out=None):
         if out is None:
             return x.copy()
         else:
             out.fill(x)
         
-    def proximal_conjugate(self, x, tau):
-        return 0
+    def proximal_conjugate(self, x, tau, out = None):
+        if out is None:
+            return 0
+        else:
+            return 0
 
     def domain_geometry(self):
         pass
diff --git a/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator.py b/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator.py
index 954f022..835f96d 100644
--- a/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator.py
+++ b/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator.py
@@ -52,40 +52,34 @@ class FiniteDiff(LinearOperator):
         
         if out is None:        
             out = np.zeros_like(x_asarr)
-            fd_arr = out
         else:
-            fd_arr = out.as_array()
-            # set the array to 0
-            fd_arr[:] = 0
-        
-#        if out is None:        
-#            out = self.gm_domain.allocate().as_array()
-#        
-#        fd_arr = out.as_array()
-#        fd_arr = self.gm_domain.allocate().as_array()
-          
+            out = out.as_array()
+            out[:]=0
+
+                  
+
         ######################## Direct for 2D  ###############################
         if x_sz == 2:
             
             if self.direction == 1:
                 
-                np.subtract( x_asarr[:,1:], x_asarr[:,0:-1], out = fd_arr[:,0:-1] )
+                np.subtract( x_asarr[:,1:], x_asarr[:,0:-1], out = out[:,0:-1] )
                 
                 if self.bnd_cond == 'Neumann':
                     pass
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,0], x_asarr[:,-1], out = fd_arr[:,-1] )
+                    np.subtract( x_asarr[:,0], x_asarr[:,-1], out = out[:,-1] )
                 else: 
                     raise ValueError('No valid boundary conditions')
                 
             if self.direction == 0:
                 
-                np.subtract( x_asarr[1:], x_asarr[0:-1], out = fd_arr[0:-1,:] )
+                np.subtract( x_asarr[1:], x_asarr[0:-1], out = out[0:-1,:] )
 
                 if self.bnd_cond == 'Neumann':
                     pass                                        
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[0,:], x_asarr[-1,:], out = fd_arr[-1,:] ) 
+                    np.subtract( x_asarr[0,:], x_asarr[-1,:], out = out[-1,:] ) 
                 else:    
                     raise ValueError('No valid boundary conditions') 
                     
@@ -94,35 +88,35 @@ class FiniteDiff(LinearOperator):
                     
             if self.direction == 0:  
                 
-                np.subtract( x_asarr[1:,:,:], x_asarr[0:-1,:,:], out = fd_arr[0:-1,:,:] )
+                np.subtract( x_asarr[1:,:,:], x_asarr[0:-1,:,:], out = out[0:-1,:,:] )
                 
                 if self.bnd_cond == 'Neumann':
                     pass
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[0,:,:], x_asarr[-1,:,:], out = fd_arr[-1,:,:] ) 
+                    np.subtract( x_asarr[0,:,:], x_asarr[-1,:,:], out = out[-1,:,:] ) 
                 else:    
                     raise ValueError('No valid boundary conditions')                      
                                                              
             if self.direction == 1:
                 
-                np.subtract( x_asarr[:,1:,:], x_asarr[:,0:-1,:], out = fd_arr[:,0:-1,:] ) 
+                np.subtract( x_asarr[:,1:,:], x_asarr[:,0:-1,:], out = out[:,0:-1,:] ) 
                 
                 if self.bnd_cond == 'Neumann':
                     pass
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,0,:], x_asarr[:,-1,:], out = fd_arr[:,-1,:] )
+                    np.subtract( x_asarr[:,0,:], x_asarr[:,-1,:], out = out[:,-1,:] )
                 else:    
                     raise ValueError('No valid boundary conditions')                      
                                 
              
             if self.direction == 2:
                 
-                np.subtract( x_asarr[:,:,1:], x_asarr[:,:,0:-1], out = fd_arr[:,:,0:-1] ) 
+                np.subtract( x_asarr[:,:,1:], x_asarr[:,:,0:-1], out = out[:,:,0:-1] ) 
                 
                 if self.bnd_cond == 'Neumann':
                     pass
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,:,0], x_asarr[:,:,-1], out = fd_arr[:,:,-1] )
+                    np.subtract( x_asarr[:,:,0], x_asarr[:,:,-1], out = out[:,:,-1] )
                 else:    
                     raise ValueError('No valid boundary conditions')  
                     
@@ -130,42 +124,42 @@ class FiniteDiff(LinearOperator):
         elif x_sz == 4:
                     
             if self.direction == 0:                            
-                np.subtract( x_asarr[1:,:,:,:], x_asarr[0:-1,:,:,:], out = fd_arr[0:-1,:,:,:] )
+                np.subtract( x_asarr[1:,:,:,:], x_asarr[0:-1,:,:,:], out = out[0:-1,:,:,:] )
                 
                 if self.bnd_cond == 'Neumann':
                     pass
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[0,:,:,:], x_asarr[-1,:,:,:], out = fd_arr[-1,:,:,:] )
+                    np.subtract( x_asarr[0,:,:,:], x_asarr[-1,:,:,:], out = out[-1,:,:,:] )
                 else:    
                     raise ValueError('No valid boundary conditions') 
                                                 
             if self.direction == 1:
-                np.subtract( x_asarr[:,1:,:,:], x_asarr[:,0:-1,:,:], out = fd_arr[:,0:-1,:,:] ) 
+                np.subtract( x_asarr[:,1:,:,:], x_asarr[:,0:-1,:,:], out = out[:,0:-1,:,:] ) 
                 
                 if self.bnd_cond == 'Neumann':
                     pass
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,0,:,:], x_asarr[:,-1,:,:], out = fd_arr[:,-1,:,:] )
+                    np.subtract( x_asarr[:,0,:,:], x_asarr[:,-1,:,:], out = out[:,-1,:,:] )
                 else:    
                     raise ValueError('No valid boundary conditions')                 
                 
             if self.direction == 2:
-                np.subtract( x_asarr[:,:,1:,:], x_asarr[:,:,0:-1,:], out = fd_arr[:,:,0:-1,:] ) 
+                np.subtract( x_asarr[:,:,1:,:], x_asarr[:,:,0:-1,:], out = out[:,:,0:-1,:] ) 
                 
                 if self.bnd_cond == 'Neumann':
                     pass
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,:,0,:], x_asarr[:,:,-1,:], out = fd_arr[:,:,-1,:] )
+                    np.subtract( x_asarr[:,:,0,:], x_asarr[:,:,-1,:], out = out[:,:,-1,:] )
                 else:    
                     raise ValueError('No valid boundary conditions')                   
                 
             if self.direction == 3:
-                np.subtract( x_asarr[:,:,:,1:], x_asarr[:,:,:,0:-1], out = fd_arr[:,:,:,0:-1] )                 
+                np.subtract( x_asarr[:,:,:,1:], x_asarr[:,:,:,0:-1], out = out[:,:,:,0:-1] )                 
 
                 if self.bnd_cond == 'Neumann':
                     pass
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,:,:,0], x_asarr[:,:,:,-1], out = fd_arr[:,:,:,-1] )
+                    np.subtract( x_asarr[:,:,:,0], x_asarr[:,:,:,-1], out = out[:,:,:,-1] )
                 else:    
                     raise ValueError('No valid boundary conditions')                   
                                 
@@ -179,51 +173,42 @@ class FiniteDiff(LinearOperator):
     def adjoint(self, x, out=None):
         
         x_asarr = x.as_array()
-        #x_asarr = x
         x_sz = len(x.shape)
         
         if out is None:        
             out = np.zeros_like(x_asarr)
-            fd_arr = out
         else:
-            #out *= 0
-            fd_arr = out.as_array()
-            fd_arr[:] = 0
-        
-#        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 = out.as_array()        
+            out[:]=0
+
         
         ######################## Adjoint for 2D  ###############################
         if x_sz == 2:        
         
             if self.direction == 1:
                 
-                np.subtract( x_asarr[:,1:], x_asarr[:,0:-1], out = fd_arr[:,1:] )
+                np.subtract( x_asarr[:,1:], x_asarr[:,0:-1], out = out[:,1:] )
                 
                 if self.bnd_cond == 'Neumann':
-                    np.subtract( x_asarr[:,0], 0, out = fd_arr[:,0] )
-                    np.subtract( -x_asarr[:,-2], 0, out = fd_arr[:,-1] )
+                    np.subtract( x_asarr[:,0], 0, out = out[:,0] )
+                    np.subtract( -x_asarr[:,-2], 0, out = out[:,-1] )
                     
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,0], x_asarr[:,-1], out = fd_arr[:,0] )                                        
+                    np.subtract( x_asarr[:,0], x_asarr[:,-1], out = out[:,0] )                                        
                     
                 else:   
                     raise ValueError('No valid boundary conditions') 
                                     
             if self.direction == 0:
                 
-                np.subtract( x_asarr[1:,:], x_asarr[0:-1,:], out = fd_arr[1:,:] )
+                np.subtract( x_asarr[1:,:], x_asarr[0:-1,:], out = out[1:,:] )
                 
                 if self.bnd_cond == 'Neumann':
-                    np.subtract( x_asarr[0,:], 0, out = fd_arr[0,:] )
-                    np.subtract( -x_asarr[-2,:], 0, out = fd_arr[-1,:] ) 
+                    np.subtract( x_asarr[0,:], 0, out = out[0,:] )
+                    np.subtract( -x_asarr[-2,:], 0, out = out[-1,:] ) 
                     
                 elif self.bnd_cond == 'Periodic':  
-                    np.subtract( x_asarr[0,:], x_asarr[-1,:], out = fd_arr[0,:] ) 
+                    np.subtract( x_asarr[0,:], x_asarr[-1,:], out = out[0,:] ) 
                     
                 else:   
                     raise ValueError('No valid boundary conditions')     
@@ -233,35 +218,35 @@ class FiniteDiff(LinearOperator):
                 
             if self.direction == 0:          
                   
-                np.subtract( x_asarr[1:,:,:], x_asarr[0:-1,:,:], out = fd_arr[1:,:,:] )
+                np.subtract( x_asarr[1:,:,:], x_asarr[0:-1,:,:], out = out[1:,:,:] )
                 
                 if self.bnd_cond == 'Neumann':
-                    np.subtract( x_asarr[0,:,:], 0, out = fd_arr[0,:,:] )
-                    np.subtract( -x_asarr[-2,:,:], 0, out = fd_arr[-1,:,:] )
+                    np.subtract( x_asarr[0,:,:], 0, out = out[0,:,:] )
+                    np.subtract( -x_asarr[-2,:,:], 0, out = out[-1,:,:] )
                 elif self.bnd_cond == 'Periodic':                     
-                    np.subtract( x_asarr[0,:,:], x_asarr[-1,:,:], out = fd_arr[0,:,:] )
+                    np.subtract( x_asarr[0,:,:], x_asarr[-1,:,:], out = out[0,:,:] )
                 else:   
                     raise ValueError('No valid boundary conditions')                     
                                     
             if self.direction == 1:
-                np.subtract( x_asarr[:,1:,:], x_asarr[:,0:-1,:], out = fd_arr[:,1:,:] )
+                np.subtract( x_asarr[:,1:,:], x_asarr[:,0:-1,:], out = out[:,1:,:] )
                 
                 if self.bnd_cond == 'Neumann':
-                    np.subtract( x_asarr[:,0,:], 0, out = fd_arr[:,0,:] )
-                    np.subtract( -x_asarr[:,-2,:], 0, out = fd_arr[:,-1,:] )
+                    np.subtract( x_asarr[:,0,:], 0, out = out[:,0,:] )
+                    np.subtract( -x_asarr[:,-2,:], 0, out = out[:,-1,:] )
                 elif self.bnd_cond == 'Periodic':                     
-                    np.subtract( x_asarr[:,0,:], x_asarr[:,-1,:], out = fd_arr[:,0,:] )
+                    np.subtract( x_asarr[:,0,:], x_asarr[:,-1,:], out = out[:,0,:] )
                 else:   
                     raise ValueError('No valid boundary conditions')                                 
                 
             if self.direction == 2:
-                np.subtract( x_asarr[:,:,1:], x_asarr[:,:,0:-1], out = fd_arr[:,:,1:] )
+                np.subtract( x_asarr[:,:,1:], x_asarr[:,:,0:-1], out = out[:,:,1:] )
                 
                 if self.bnd_cond == 'Neumann':
-                    np.subtract( x_asarr[:,:,0], 0, out = fd_arr[:,:,0] ) 
-                    np.subtract( -x_asarr[:,:,-2], 0, out = fd_arr[:,:,-1] ) 
+                    np.subtract( x_asarr[:,:,0], 0, out = out[:,:,0] ) 
+                    np.subtract( -x_asarr[:,:,-2], 0, out = out[:,:,-1] ) 
                 elif self.bnd_cond == 'Periodic':                     
-                    np.subtract( x_asarr[:,:,0], x_asarr[:,:,-1], out = fd_arr[:,:,0] )
+                    np.subtract( x_asarr[:,:,0], x_asarr[:,:,-1], out = out[:,:,0] )
                 else:   
                     raise ValueError('No valid boundary conditions')                                 
         
@@ -269,51 +254,51 @@ class FiniteDiff(LinearOperator):
         elif x_sz == 4:                
                 
             if self.direction == 0:                            
-                np.subtract( x_asarr[1:,:,:,:], x_asarr[0:-1,:,:,:], out = fd_arr[1:,:,:,:] )
+                np.subtract( x_asarr[1:,:,:,:], x_asarr[0:-1,:,:,:], out = out[1:,:,:,:] )
                 
                 if self.bnd_cond == 'Neumann':
-                    np.subtract( x_asarr[0,:,:,:], 0, out = fd_arr[0,:,:,:] )
-                    np.subtract( -x_asarr[-2,:,:,:], 0, out = fd_arr[-1,:,:,:] )
+                    np.subtract( x_asarr[0,:,:,:], 0, out = out[0,:,:,:] )
+                    np.subtract( -x_asarr[-2,:,:,:], 0, out = out[-1,:,:,:] )
                     
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[0,:,:,:], x_asarr[-1,:,:,:], out = fd_arr[0,:,:,:] )
+                    np.subtract( x_asarr[0,:,:,:], x_asarr[-1,:,:,:], out = out[0,:,:,:] )
                 else:    
                     raise ValueError('No valid boundary conditions') 
                                 
             if self.direction == 1:
-                np.subtract( x_asarr[:,1:,:,:], x_asarr[:,0:-1,:,:], out = fd_arr[:,1:,:,:] )
+                np.subtract( x_asarr[:,1:,:,:], x_asarr[:,0:-1,:,:], out = out[:,1:,:,:] )
                 
                 if self.bnd_cond == 'Neumann':
-                   np.subtract( x_asarr[:,0,:,:], 0, out = fd_arr[:,0,:,:] )
-                   np.subtract( -x_asarr[:,-2,:,:], 0, out = fd_arr[:,-1,:,:] )
+                   np.subtract( x_asarr[:,0,:,:], 0, out = out[:,0,:,:] )
+                   np.subtract( -x_asarr[:,-2,:,:], 0, out = out[:,-1,:,:] )
                     
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,0,:,:], x_asarr[:,-1,:,:], out = fd_arr[:,0,:,:] )
+                    np.subtract( x_asarr[:,0,:,:], x_asarr[:,-1,:,:], out = out[:,0,:,:] )
                 else:    
                     raise ValueError('No valid boundary conditions') 
                     
                 
             if self.direction == 2:
-                np.subtract( x_asarr[:,:,1:,:], x_asarr[:,:,0:-1,:], out = fd_arr[:,:,1:,:] )
+                np.subtract( x_asarr[:,:,1:,:], x_asarr[:,:,0:-1,:], out = out[:,:,1:,:] )
                 
                 if self.bnd_cond == 'Neumann':
-                    np.subtract( x_asarr[:,:,0,:], 0, out = fd_arr[:,:,0,:] ) 
-                    np.subtract( -x_asarr[:,:,-2,:], 0, out = fd_arr[:,:,-1,:] ) 
+                    np.subtract( x_asarr[:,:,0,:], 0, out = out[:,:,0,:] ) 
+                    np.subtract( -x_asarr[:,:,-2,:], 0, out = out[:,:,-1,:] ) 
                     
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,:,0,:], x_asarr[:,:,-1,:], out = fd_arr[:,:,0,:] )
+                    np.subtract( x_asarr[:,:,0,:], x_asarr[:,:,-1,:], out = out[:,:,0,:] )
                 else:    
                     raise ValueError('No valid boundary conditions')                 
                 
             if self.direction == 3:
-                np.subtract( x_asarr[:,:,:,1:], x_asarr[:,:,:,0:-1], out = fd_arr[:,:,:,1:] )
+                np.subtract( x_asarr[:,:,:,1:], x_asarr[:,:,:,0:-1], out = out[:,:,:,1:] )
                 
                 if self.bnd_cond == 'Neumann':
-                    np.subtract( x_asarr[:,:,:,0], 0, out = fd_arr[:,:,:,0] ) 
-                    np.subtract( -x_asarr[:,:,:,-2], 0, out = fd_arr[:,:,:,-1] )   
+                    np.subtract( x_asarr[:,:,:,0], 0, out = out[:,:,:,0] ) 
+                    np.subtract( -x_asarr[:,:,:,-2], 0, out = out[:,:,:,-1] )   
                     
                 elif self.bnd_cond == 'Periodic':
-                    np.subtract( x_asarr[:,:,:,0], x_asarr[:,:,:,-1], out = fd_arr[:,:,:,0] )
+                    np.subtract( x_asarr[:,:,:,0], x_asarr[:,:,:,-1], out = out[:,:,:,0] )
                 else:    
                     raise ValueError('No valid boundary conditions')                  
                               
@@ -332,8 +317,7 @@ class FiniteDiff(LinearOperator):
         return self.gm_domain
        
     def norm(self):
-        x0 = self.gm_domain.allocate()
-        x0.fill( np.random.random_sample(x0.shape) )
+        x0 = self.gm_domain.allocate('random_int')
         self.s1, sall, svec = PowerMethodNonsquare(self, 25, x0)
         return self.s1
     
@@ -341,23 +325,46 @@ class FiniteDiff(LinearOperator):
 if __name__ == '__main__':
     
     from ccpi.framework import ImageGeometry
+    import numpy
     
     N, M = 2, 3
 
     ig = ImageGeometry(N, M)
 
 
-    FD = FiniteDiff(ig, direction = 0, bnd_cond = 'Neumann')
+    FD = FiniteDiff(ig, direction = 1, bnd_cond = 'Neumann')
     u = FD.domain_geometry().allocate('random_int')
-    
-    
+        
     res = FD.domain_geometry().allocate()
+    res1 = FD.range_geometry().allocate()
     FD.direct(u, out=res)
-    print(res.as_array())
-#    z = FD.direct(u)
-    
+
+    z = FD.direct(u)    
 #    print(z.as_array(), res.as_array())
 
+    for i in range(10):
+#        
+        z1 = FD.direct(u) 
+        FD.direct(u, out=res)
+        
+        u = ig.allocate('random_int')
+        res = u
+        z1  = u
+        numpy.testing.assert_array_almost_equal(z1.as_array(), \
+                                                res.as_array(), decimal=4)
+        
+#        print(z1.as_array(), res.as_array())
+        z2 = FD.adjoint(z1) 
+        FD.adjoint(z1, out=res1)  
+        numpy.testing.assert_array_almost_equal(z2.as_array(), \
+                                                res1.as_array(), decimal=4)        
+        
+        
+        
+
+        
+        
+        
     
 #    w = G.range_geometry().allocate('random_int')
     
diff --git a/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator_old.py b/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator_old.py
new file mode 100644
index 0000000..387fb4b
--- /dev/null
+++ b/Wrappers/Python/ccpi/optimisation/operators/FiniteDifferenceOperator_old.py
@@ -0,0 +1,374 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Mar  1 22:51:17 2019
+
+@author: evangelos
+"""
+
+from ccpi.optimisation.operators import LinearOperator
+from ccpi.optimisation.ops import PowerMethodNonsquare
+from ccpi.framework import ImageData, BlockDataContainer
+import numpy as np
+
+class FiniteDiff(LinearOperator):
+    
+    # Works for Neum/Symmetric &  periodic boundary conditions
+    # TODO add central differences???
+    # TODO not very well optimised, too many conditions
+    # TODO add discretisation step, should get that from imageGeometry
+    
+    # Grad_order = ['channels', 'direction_z', 'direction_y', 'direction_x']
+    # Grad_order = ['channels', 'direction_y', 'direction_x']
+    # Grad_order = ['direction_z', 'direction_y', 'direction_x']
+    # Grad_order = ['channels', 'direction_z', 'direction_y', 'direction_x']
+    
+    def __init__(self, gm_domain, gm_range=None, direction=0, bnd_cond = 'Neumann'):
+        ''''''
+        super(FiniteDiff, self).__init__() 
+        '''FIXME: domain and range should be geometries'''
+        self.gm_domain = gm_domain
+        self.gm_range = gm_range
+        
+        self.direction = direction
+        self.bnd_cond = bnd_cond
+        
+        # Domain Geometry = Range Geometry if not stated
+        if self.gm_range is None:
+            self.gm_range = self.gm_domain
+        # check direction and "length" of geometry
+        if self.direction + 1 > len(self.gm_domain.shape):
+            raise ValueError('Gradient directions more than geometry domain')      
+        
+        #self.voxel_size = kwargs.get('voxel_size',1)
+        # this wrongly assumes a homogeneous voxel size
+        self.voxel_size = self.gm_domain.voxel_size_x
+
+
+    def direct(self, x, out=None):
+        
+        x_asarr = x.as_array()
+        x_sz = len(x.shape)
+        
+        if out is None:        
+            out = np.zeros_like(x_asarr)
+            fd_arr = out
+        else:
+            fd_arr = out.as_array()   
+#            fd_arr[:]=0
+        
+#        if out is None:        
+#            out = self.gm_domain.allocate().as_array()
+#        
+#        fd_arr = out.as_array()
+#        fd_arr = self.gm_domain.allocate().as_array()
+          
+        ######################## Direct for 2D  ###############################
+        if x_sz == 2:
+            
+            if self.direction == 1:
+                
+                np.subtract( x_asarr[:,1:], x_asarr[:,0:-1], out = fd_arr[:,0:-1] )
+                
+                if self.bnd_cond == 'Neumann':
+                    pass                                        
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,0], x_asarr[:,-1], out = fd_arr[:,-1] )
+                else: 
+                    raise ValueError('No valid boundary conditions')
+                
+            if self.direction == 0:
+                
+                np.subtract( x_asarr[1:], x_asarr[0:-1], out = fd_arr[0:-1,:] )
+
+                if self.bnd_cond == 'Neumann':
+                    pass                                        
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[0,:], x_asarr[-1,:], out = fd_arr[-1,:] ) 
+                else:    
+                    raise ValueError('No valid boundary conditions') 
+                    
+        ######################## Direct for 3D  ###############################                        
+        elif x_sz == 3:
+                    
+            if self.direction == 0:  
+                
+                np.subtract( x_asarr[1:,:,:], x_asarr[0:-1,:,:], out = fd_arr[0:-1,:,:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    pass
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[0,:,:], x_asarr[-1,:,:], out = fd_arr[-1,:,:] ) 
+                else:    
+                    raise ValueError('No valid boundary conditions')                      
+                                                             
+            if self.direction == 1:
+                
+                np.subtract( x_asarr[:,1:,:], x_asarr[:,0:-1,:], out = fd_arr[:,0:-1,:] ) 
+                
+                if self.bnd_cond == 'Neumann':
+                    pass
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,0,:], x_asarr[:,-1,:], out = fd_arr[:,-1,:] )
+                else:    
+                    raise ValueError('No valid boundary conditions')                      
+                                
+             
+            if self.direction == 2:
+                
+                np.subtract( x_asarr[:,:,1:], x_asarr[:,:,0:-1], out = fd_arr[:,:,0:-1] ) 
+                
+                if self.bnd_cond == 'Neumann':
+                    pass
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,:,0], x_asarr[:,:,-1], out = fd_arr[:,:,-1] )
+                else:    
+                    raise ValueError('No valid boundary conditions')  
+                    
+        ######################## Direct for 4D  ###############################
+        elif x_sz == 4:
+                    
+            if self.direction == 0:                            
+                np.subtract( x_asarr[1:,:,:,:], x_asarr[0:-1,:,:,:], out = fd_arr[0:-1,:,:,:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    pass
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[0,:,:,:], x_asarr[-1,:,:,:], out = fd_arr[-1,:,:,:] )
+                else:    
+                    raise ValueError('No valid boundary conditions') 
+                                                
+            if self.direction == 1:
+                np.subtract( x_asarr[:,1:,:,:], x_asarr[:,0:-1,:,:], out = fd_arr[:,0:-1,:,:] ) 
+                
+                if self.bnd_cond == 'Neumann':
+                    pass
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,0,:,:], x_asarr[:,-1,:,:], out = fd_arr[:,-1,:,:] )
+                else:    
+                    raise ValueError('No valid boundary conditions')                 
+                
+            if self.direction == 2:
+                np.subtract( x_asarr[:,:,1:,:], x_asarr[:,:,0:-1,:], out = fd_arr[:,:,0:-1,:] ) 
+                
+                if self.bnd_cond == 'Neumann':
+                    pass
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,:,0,:], x_asarr[:,:,-1,:], out = fd_arr[:,:,-1,:] )
+                else:    
+                    raise ValueError('No valid boundary conditions')                   
+                
+            if self.direction == 3:
+                np.subtract( x_asarr[:,:,:,1:], x_asarr[:,:,:,0:-1], out = fd_arr[:,:,:,0:-1] )                 
+
+                if self.bnd_cond == 'Neumann':
+                    pass
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,:,:,0], x_asarr[:,:,:,-1], out = fd_arr[:,:,:,-1] )
+                else:    
+                    raise ValueError('No valid boundary conditions')                   
+                                
+        else:
+            raise NotImplementedError                
+         
+#        res = out #/self.voxel_size 
+        return type(x)(out)
+
+                    
+    def adjoint(self, x, out=None):
+        
+        x_asarr = x.as_array()
+        #x_asarr = x
+        x_sz = len(x.shape)
+        
+        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()
+        
+        ######################## Adjoint for 2D  ###############################
+        if x_sz == 2:        
+        
+            if self.direction == 1:
+                
+                np.subtract( x_asarr[:,1:], x_asarr[:,0:-1], out = fd_arr[:,1:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    np.subtract( x_asarr[:,0], 0, out = fd_arr[:,0] )
+                    np.subtract( -x_asarr[:,-2], 0, out = fd_arr[:,-1] )
+                    
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,0], x_asarr[:,-1], out = fd_arr[:,0] )                                        
+                    
+                else:   
+                    raise ValueError('No valid boundary conditions') 
+                                    
+            if self.direction == 0:
+                
+                np.subtract( x_asarr[1:,:], x_asarr[0:-1,:], out = fd_arr[1:,:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    np.subtract( x_asarr[0,:], 0, out = fd_arr[0,:] )
+                    np.subtract( -x_asarr[-2,:], 0, out = fd_arr[-1,:] ) 
+                    
+                elif self.bnd_cond == 'Periodic':  
+                    np.subtract( x_asarr[0,:], x_asarr[-1,:], out = fd_arr[0,:] ) 
+                    
+                else:   
+                    raise ValueError('No valid boundary conditions')     
+        
+        ######################## Adjoint for 3D  ###############################        
+        elif x_sz == 3:                
+                
+            if self.direction == 0:          
+                  
+                np.subtract( x_asarr[1:,:,:], x_asarr[0:-1,:,:], out = fd_arr[1:,:,:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    np.subtract( x_asarr[0,:,:], 0, out = fd_arr[0,:,:] )
+                    np.subtract( -x_asarr[-2,:,:], 0, out = fd_arr[-1,:,:] )
+                elif self.bnd_cond == 'Periodic':                     
+                    np.subtract( x_asarr[0,:,:], x_asarr[-1,:,:], out = fd_arr[0,:,:] )
+                else:   
+                    raise ValueError('No valid boundary conditions')                     
+                                    
+            if self.direction == 1:
+                np.subtract( x_asarr[:,1:,:], x_asarr[:,0:-1,:], out = fd_arr[:,1:,:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    np.subtract( x_asarr[:,0,:], 0, out = fd_arr[:,0,:] )
+                    np.subtract( -x_asarr[:,-2,:], 0, out = fd_arr[:,-1,:] )
+                elif self.bnd_cond == 'Periodic':                     
+                    np.subtract( x_asarr[:,0,:], x_asarr[:,-1,:], out = fd_arr[:,0,:] )
+                else:   
+                    raise ValueError('No valid boundary conditions')                                 
+                
+            if self.direction == 2:
+                np.subtract( x_asarr[:,:,1:], x_asarr[:,:,0:-1], out = fd_arr[:,:,1:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    np.subtract( x_asarr[:,:,0], 0, out = fd_arr[:,:,0] ) 
+                    np.subtract( -x_asarr[:,:,-2], 0, out = fd_arr[:,:,-1] ) 
+                elif self.bnd_cond == 'Periodic':                     
+                    np.subtract( x_asarr[:,:,0], x_asarr[:,:,-1], out = fd_arr[:,:,0] )
+                else:   
+                    raise ValueError('No valid boundary conditions')                                 
+        
+        ######################## Adjoint for 4D  ###############################        
+        elif x_sz == 4:                
+                
+            if self.direction == 0:                            
+                np.subtract( x_asarr[1:,:,:,:], x_asarr[0:-1,:,:,:], out = fd_arr[1:,:,:,:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    np.subtract( x_asarr[0,:,:,:], 0, out = fd_arr[0,:,:,:] )
+                    np.subtract( -x_asarr[-2,:,:,:], 0, out = fd_arr[-1,:,:,:] )
+                    
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[0,:,:,:], x_asarr[-1,:,:,:], out = fd_arr[0,:,:,:] )
+                else:    
+                    raise ValueError('No valid boundary conditions') 
+                                
+            if self.direction == 1:
+                np.subtract( x_asarr[:,1:,:,:], x_asarr[:,0:-1,:,:], out = fd_arr[:,1:,:,:] )
+                
+                if self.bnd_cond == 'Neumann':
+                   np.subtract( x_asarr[:,0,:,:], 0, out = fd_arr[:,0,:,:] )
+                   np.subtract( -x_asarr[:,-2,:,:], 0, out = fd_arr[:,-1,:,:] )
+                    
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,0,:,:], x_asarr[:,-1,:,:], out = fd_arr[:,0,:,:] )
+                else:    
+                    raise ValueError('No valid boundary conditions') 
+                    
+                
+            if self.direction == 2:
+                np.subtract( x_asarr[:,:,1:,:], x_asarr[:,:,0:-1,:], out = fd_arr[:,:,1:,:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    np.subtract( x_asarr[:,:,0,:], 0, out = fd_arr[:,:,0,:] ) 
+                    np.subtract( -x_asarr[:,:,-2,:], 0, out = fd_arr[:,:,-1,:] ) 
+                    
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,:,0,:], x_asarr[:,:,-1,:], out = fd_arr[:,:,0,:] )
+                else:    
+                    raise ValueError('No valid boundary conditions')                 
+                
+            if self.direction == 3:
+                np.subtract( x_asarr[:,:,:,1:], x_asarr[:,:,:,0:-1], out = fd_arr[:,:,:,1:] )
+                
+                if self.bnd_cond == 'Neumann':
+                    np.subtract( x_asarr[:,:,:,0], 0, out = fd_arr[:,:,:,0] ) 
+                    np.subtract( -x_asarr[:,:,:,-2], 0, out = fd_arr[:,:,:,-1] )   
+                    
+                elif self.bnd_cond == 'Periodic':
+                    np.subtract( x_asarr[:,:,:,0], x_asarr[:,:,:,-1], out = fd_arr[:,:,:,0] )
+                else:    
+                    raise ValueError('No valid boundary conditions')                  
+                              
+        else:
+            raise NotImplementedError
+            
+        out *= -1 #/self.voxel_size
+        return type(x)(out)
+            
+    def range_geometry(self):
+        '''Returns the range geometry'''
+        return self.gm_range
+    
+    def domain_geometry(self):
+        '''Returns the domain geometry'''
+        return self.gm_domain
+       
+    def norm(self):
+        x0 = self.gm_domain.allocate()
+        x0.fill( np.random.random_sample(x0.shape) )
+        self.s1, sall, svec = PowerMethodNonsquare(self, 25, x0)
+        return self.s1
+    
+    
+if __name__ == '__main__':
+    
+    from ccpi.framework import ImageGeometry
+    import numpy
+    
+    N, M = 2, 3
+
+    ig = ImageGeometry(N, M)
+
+
+    FD = FiniteDiff(ig, direction = 0, bnd_cond = 'Neumann')
+    u = FD.domain_geometry().allocate('random_int')
+    
+    
+    res = FD.domain_geometry().allocate()
+    FD.direct(u, out=res)
+
+    z = FD.direct(u)    
+    print(z.as_array(), res.as_array())
+
+    for i in range(10):
+        
+        z1 = FD.direct(u) 
+        FD.direct(u, out=res)
+        numpy.testing.assert_array_almost_equal(z1.as_array(), \
+                                                res.as_array(), decimal=4)
+        
+
+        
+        
+        
+    
+#    w = G.range_geometry().allocate('random_int')
+    
+
+    
+    
\ No newline at end of file
diff --git a/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py b/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py
index d655653..4935435 100644
--- a/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py
+++ b/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py
@@ -49,22 +49,21 @@ class Gradient(LinearOperator):
         
                 
         if out is not None:
+            
             for i in range(self.gm_range.shape[0]):
-                self.FD.direction=self.ind[i]
+                self.FD.direction = self.ind[i]
                 self.FD.direct(x, out = out[i])
-#                FiniteDiff(self.gm_domain, direction = self.ind[i], bnd_cond = self.bnd_cond).direct(x, out=out[i])
-            return out
         else:
             tmp = self.gm_range.allocate()        
             for i in range(tmp.shape[0]):
                 self.FD.direction=self.ind[i]
                 tmp.get_item(i).fill(self.FD.direct(x))
-#                tmp.get_item(i).fill(FiniteDiff(self.gm_domain, direction = self.ind[i], bnd_cond = self.bnd_cond).direct(x))
             return tmp    
         
     def adjoint(self, x, out=None):
         
         if out is not None:
+
             tmp = self.gm_domain.allocate()            
             for i in range(x.shape[0]):
                 self.FD.direction=self.ind[i] 
@@ -77,6 +76,7 @@ class Gradient(LinearOperator):
             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))
             return tmp    
             
@@ -96,7 +96,9 @@ class Gradient(LinearOperator):
     def __rmul__(self, scalar):
         return ScaledOperator(self, scalar) 
     
-    
+    ###########################################################################
+    ###############  For preconditioning ######################################
+    ###########################################################################
     def matrix(self):
         
         tmp = self.gm_range.allocate()
@@ -220,17 +222,22 @@ if __name__ == '__main__':
 #
     u = G.domain_geometry().allocate('random_int')
     w = G.range_geometry().allocate('random_int')
-#
-#
+    
+    
+    print( "################   without out #############")
+          
+    print( (G.direct(u)*w).sum(),  (u*G.adjoint(w)).sum() ) 
+        
+
+    print( "################   with out #############")
+    
     res = G.range_geometry().allocate()
-#    
-    G.direct(u, out=res)
-    z = G.direct(u)
-#    
-    print(res[0].as_array())
-    print(z[0].as_array())
-#
-##    LHS = (G.direct(u)*w).sum()
-##    RHS = (u * G.adjoint(w)).sum()
+    res1 = G.domain_geometry().allocate()
+    G.direct(u, out = res)          
+    G.adjoint(w, out = res1)
+          
+    print( (res*w).sum(),  (u*res1).sum() )     
     
-#    
+    
+    
+
diff --git a/Wrappers/Python/ccpi/optimisation/operators/ScaledOperator.py b/Wrappers/Python/ccpi/optimisation/operators/ScaledOperator.py
index 3203dde..ba0049e 100644
--- a/Wrappers/Python/ccpi/optimisation/operators/ScaledOperator.py
+++ b/Wrappers/Python/ccpi/optimisation/operators/ScaledOperator.py
@@ -3,12 +3,10 @@ import numpy
 
 class ScaledOperator(object):
     '''ScaledOperator
-
     A class to represent the scalar multiplication of an Operator with a scalar.
     It holds an operator and a scalar. Basically it returns the multiplication
     of the result of direct and adjoint of the operator with the scalar.
     For the rest it behaves like the operator it holds.
-
     Args:
        operator (Operator): a Operator or LinearOperator
        scalar (Number): a scalar multiplier
@@ -50,3 +48,4 @@ class ScaledOperator(object):
         return self.operator.domain_geometry()
     def is_linear(self):
         return self.operator.is_linear()
+
diff --git a/Wrappers/Python/wip/pdhg_TV_denoising.py b/Wrappers/Python/wip/pdhg_TV_denoising.py
index f569fa7..f276b46 100755
--- a/Wrappers/Python/wip/pdhg_TV_denoising.py
+++ b/Wrappers/Python/wip/pdhg_TV_denoising.py
@@ -23,11 +23,14 @@ from timeit import default_timer as timer
 def dt(steps):
     return steps[-1] - steps[-2]
 
+#%%
+
 
 # ############################################################################
 # Create phantom for TV denoising
 
-N = 512
+N = 200
+
 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
@@ -62,9 +65,9 @@ if method == '0':
     #### Create functions
       
     f1 = alpha * MixedL21Norm()
-    f2 = 0.5 * L2NormSquared(b = noisy_data)
-    
-    f = BlockFunction(f1, f2 )                                        
+    f2 = 0.5 * L2NormSquared(b = noisy_data)    
+    f = BlockFunction(f1, f2)  
+                                      
     g = ZeroFun()
     
 else:
@@ -82,75 +85,106 @@ else:
 # Compute operator Norm
 normK = operator.norm()
 print ("normK", normK)
+
 # Primal & dual stepsizes
 sigma = 1
 tau = 1/(sigma*normK**2)
 
-# opt = {'niter':2000, 'memopt': True}
+opt = {'niter':1000}
+opt1 = {'niter':1000, 'memopt': True}
 
-# res, time, primal, dual, pdgap = PDHG_old(f, g, operator, tau = tau, sigma = sigma, opt = opt) 
- 
+t1 = timer()
+res, time, primal, dual, pdgap = PDHG_old(f, g, operator, tau = tau, sigma = sigma, opt = opt) 
+print(timer()-t1)
 
+print("with memopt \n")
 
-# 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 = 100
-
-
-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,3,1)
-plt.imshow(noisy_data.as_array())
-plt.colorbar()
-plt.subplot(1,3,2)
-plt.imshow(sol)
+t2 = timer()
+res1, time1, primal1, dual1, pdgap1 = PDHG_old(f, g, operator, tau = tau, sigma = sigma, opt = opt1) 
+print(timer()-t2)
+
+plt.figure(figsize=(5,5))
+plt.imshow(res.as_array())
 plt.colorbar()
-plt.subplot(1,3,3)
-plt.imshow(pdhgo.get_output().as_array())
+plt.show()
+ 
+plt.figure(figsize=(5,5))
+plt.imshow(res1.as_array())
 plt.colorbar()
+plt.show()
 
+
+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 = 100
+#
+#
+#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,3,1)
+#plt.imshow(noisy_data.as_array())
+#plt.colorbar()
+#plt.subplot(1,3,2)
+#plt.imshow(sol)
+#plt.colorbar()
+#plt.subplot(1,3,3)
+#plt.imshow(pdhgo.get_output().as_array())
+#plt.colorbar()
 #
-###
-#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()
 #
+#
+##%%
+##
diff --git a/Wrappers/Python/wip/pdhg_TV_denoising_salt_pepper.py b/Wrappers/Python/wip/pdhg_TV_denoising_salt_pepper.py
index eb7eef4..cec9770 100644
--- a/Wrappers/Python/wip/pdhg_TV_denoising_salt_pepper.py
+++ b/Wrappers/Python/wip/pdhg_TV_denoising_salt_pepper.py
@@ -34,7 +34,7 @@ ig = ImageGeometry(voxel_num_x = N, voxel_num_y = N)
 ag = ig
 
 # Create noisy data. Add Gaussian noise
-n1 = random_noise(data, mode = 's&p', salt_vs_pepper = 0.9)
+n1 = random_noise(data, mode = 's&p', salt_vs_pepper = 0.9, amount=0.2)
 noisy_data = ImageData(n1)
 
 plt.imshow(noisy_data.as_array())
@@ -44,10 +44,10 @@ plt.show()
 #%%
 
 # Regularisation Parameter
-alpha = 10
+alpha = 2
 
 #method = input("Enter structure of PDHG (0=Composite or 1=NotComposite): ")
-method = '1'
+method = '0'
 if method == '0':
 
     # Create operators
@@ -78,15 +78,27 @@ else:
     ###########################################################################
 #%%
     
-# Compute operator Norm
-normK = operator.norm()
-print ("normK", normK)
-# Primal & dual stepsizes
-#sigma = 1
-#tau = 1/(sigma*normK**2)
-
-sigma = 1/normK
-tau = 1/normK
+diag_precon =  True
+
+if diag_precon:
+    
+    def tau_sigma_precond(operator):
+        
+        tau = 1/operator.sum_abs_row()
+        sigma = 1/ operator.sum_abs_col()
+               
+        return tau, sigma
+
+    tau, sigma = tau_sigma_precond(operator)
+             
+else:
+    # Compute operator Norm
+    normK = operator.norm()
+    print ("normK", normK)
+    # Primal & dual stepsizes
+    sigma = 1/normK
+    tau = 1/normK
+#    tau = 1/(sigma*normK**2)
 
 opt = {'niter':2000}
 
diff --git a/Wrappers/Python/wip/pdhg_TV_tomography2D.py b/Wrappers/Python/wip/pdhg_TV_tomography2D.py
index f06f166..159f2ea 100644
--- a/Wrappers/Python/wip/pdhg_TV_tomography2D.py
+++ b/Wrappers/Python/wip/pdhg_TV_tomography2D.py
@@ -21,6 +21,7 @@ from ccpi.optimisation.functions import ZeroFun, L2NormSquared, \
 
 from ccpi.astra.ops import AstraProjectorSimple
 from skimage.util import random_noise
+from timeit import default_timer as timer
 
 
 #%%###############################################################################
@@ -118,15 +119,37 @@ else:
 #
 #pdhg.run(5000)
     
-opt = {'niter':2000}
-#
+opt = {'niter':300}
+opt1 = {'niter':300, 'memopt': True}
+
+
+t1 = timer()
 res, time, primal, dual, pdgap = PDHG_old(f, g, operator, tau = tau, sigma = sigma, opt = opt) 
 
+print(timer()-t1)
 plt.figure(figsize=(5,5))
 plt.imshow(res.as_array())
 plt.colorbar()
-plt.show() 
-
+plt.show()
+ 
+#%%
+print("with memopt \n")
+#
+t2 = timer()
+res1, time1, primal1, dual1, pdgap1 = PDHG_old(f, g, operator, tau = tau, sigma = sigma, opt = opt1) 
+#print(timer()-t2)
+#
+#
+plt.figure(figsize=(5,5))
+plt.imshow(res1.as_array())
+plt.colorbar()
+plt.show()
+#
+#%%
+plt.figure(figsize=(5,5))
+plt.imshow(np.abs(res1.as_array()-res.as_array()))
+plt.colorbar()
+plt.show()
 #%%
 #sol = pdhg.get_output().as_array()
 #fig = plt.figure()
diff --git a/Wrappers/Python/wip/test_profile.py b/Wrappers/Python/wip/test_profile.py
index 7be19f9..f14c0c3 100644
--- a/Wrappers/Python/wip/test_profile.py
+++ b/Wrappers/Python/wip/test_profile.py
@@ -9,18 +9,76 @@ Created on Mon Apr  8 13:57:46 2019
 # profile direct, adjoint, gradient
 
 from ccpi.framework import ImageGeometry
-from ccpi.optimisation.operators import Gradient
+from ccpi.optimisation.operators import Gradient, BlockOperator, Identity
+from ccpi.optimisation.functions import MixedL21Norm, L2NormSquared, BlockFunction
+import numpy
 
-N, M = 500, 500
+N, M, K = 2, 3, 2
 
 ig = ImageGeometry(N, M)
+b = ig.allocate('random_int')
 
 G = Gradient(ig)
+Id = Identity(ig)
 
-u = G.domain_geometry().allocate('random_int')
-w = G.range_geometry().allocate('random_int')
+#operator = BlockOperator(G, Id)
+operator = G
 
-for i in range(500):
+f1 = MixedL21Norm()
+f2 = L2NormSquared(b = b)
+
+f = BlockFunction( f1, f2)
+
+
+x_old = operator.domain_geometry().allocate()
+y_old = operator.range_geometry().allocate('random_int') 
+  
+
+xbar = operator.domain_geometry().allocate('random_int')
+
+x_tmp = x_old.copy()
+x = x_old.copy()
+    
+y_tmp = operator.range_geometry().allocate()
+y = y_old.copy()
+
+y1 = y.copy()
+
+sigma = 20
+
+for i in range(100):
+    
+    operator.direct(xbar, out = y_tmp)  
+    y_tmp *= sigma
+    y_tmp += y_old
+
+    
+    y_tmp1 =  sigma * operator.direct(xbar) + y_old
+    
+    print(i)    
+    print(" y_old :", y_old[0].as_array(), "\n")
+    print(" y_tmp[0] :", y_tmp[0].as_array(),"\n")
+    print(" y_tmp1[0] :", y_tmp1[0].as_array())
+    
+    
+    numpy.testing.assert_array_equal(y_tmp[0].as_array(), \
+                                            y_tmp1[0].as_array())
+    
+    numpy.testing.assert_array_equal(y_tmp[1].as_array(), \
+                                            y_tmp1[1].as_array()) 
+    
+    
+    y1 = f.proximal_conjugate(y_tmp1, sigma)
+    f.proximal_conjugate(y_tmp, sigma, y)
+    
+
+    
+    
+ 
+    
+
+    
+    
+    
     
-    res = G.adjoint(w)
     
\ No newline at end of file