Merge pull request #250 from vais-ral/KL

Implementation of KullbackLeibler with memopt and tests

3 files changed, 117 insertions, 49 deletions

diff --git a/Wrappers/Python/ccpi/optimisation/functions/KullbackLeibler.py b/Wrappers/Python/ccpi/optimisation/functions/KullbackLeibler.py
index 18af154..e7e41f7 100644
--- a/Wrappers/Python/ccpi/optimisation/functions/KullbackLeibler.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/KullbackLeibler.py
@@ -19,46 +19,119 @@
 
 import numpy
 from ccpi.optimisation.functions import Function
-from ccpi.optimisation.functions.ScaledFunction import ScaledFunction
-from ccpi.framework import DataContainer, ImageData, ImageGeometry  
+from ccpi.optimisation.functions.ScaledFunction import ScaledFunction 
+from ccpi.framework import ImageData, ImageGeometry
+import functools
 
 class KullbackLeibler(Function):
     
-    def __init__(self,data,**kwargs):
+    ''' Assume that data > 0
+                
+    '''
+    
+    def __init__(self,data, **kwargs):
         
         super(KullbackLeibler, self).__init__()
         
         self.b = data        
         self.bnoise = kwargs.get('bnoise', 0)
 
-        self.sum_value = self.b + self.bnoise        
-        if  (self.sum_value.as_array()<0).any():
-            self.sum_value = numpy.inf
                                                 
     def __call__(self, x):
         
+        # TODO check
+        
+        self.sum_value = x + self.bnoise        
+        if  (self.sum_value.as_array()<0).any():
+            self.sum_value = numpy.inf
+        
         if self.sum_value==numpy.inf:
             return numpy.inf
         else:
-            return numpy.sum( x.as_array() - self.b.as_array() * numpy.log(self.sum_value.as_array()))
+            tmp = self.sum_value
+            #tmp.fill( numpy.log(tmp.as_array()) )            
+            self.log(tmp)
+            return (x - self.b * tmp ).sum()
+            
+#            return numpy.sum( x.as_array() - self.b.as_array() * numpy.log(self.sum_value.as_array()))
 
+    def log(self, datacontainer):
+        '''calculates the in-place log of the datacontainer'''
+        if not functools.reduce(lambda x,y: x and y>0,
+                                datacontainer.as_array().ravel(), True):
+            raise ValueError('KullbackLeibler. Cannot calculate log of negative number')
+        datacontainer.fill( numpy.log(datacontainer.as_array()) )
         
-    def gradient(self, x):
+    def gradient(self, x, out=None):
         
         #TODO Division check
-        return 1 - self.b/(x + self.bnoise)
-    
-    def convex_conjugate(self, x, out=None):
-        pass
+        if out is None:
+            return 1 - self.b/(x + self.bnoise)
+        else:
+            x.add(self.bnoise, out=out)
+            self.b.divide(out, out=out)
+            out.subtract(1, out=out)
+            out *= -1
+            
+    def convex_conjugate(self, x):
+        
+        tmp = self.b/( 1 - x )
+        self.log(tmp)
+        return (self.b * ( tmp - 1 ) - self.bnoise * (x - 1)).sum()
+#        return self.b * ( ImageData(numpy.log(self.b/(1-x)) - 1 )) - self.bnoise * (x - 1)
     
     def proximal(self, x, tau, out=None):
         
-        z = x + tau * self.bnoise
-        return (z + 1) - ((z-1)**2 + 4 * tau * self.b).sqrt()
-
+        if out is None:        
+            return 0.5 *( (x - self.bnoise - tau) + ( (x + self.bnoise - tau)**2 + 4*tau*self.b   ) .sqrt() )
+        else:
+            tmp =  0.5 *( (x - self.bnoise - tau) + 
+                        ( (x + self.bnoise - tau)**2 + 4*tau*self.b   ) .sqrt()
+                        )
+            x.add(self.bnoise, out=out)
+            out -= tau
+            out *= out
+            tmp = self.b * (4 * tau)
+            out.add(tmp, out=out)
+            out.sqrt(out=out)
+            
+            x.subtract(self.bnoise, out=tmp)
+            tmp -= tau
+            
+            out += tmp
+            
+            out *= 0.5
+            
+            
     
     def proximal_conjugate(self, x, tau, out=None):
-        pass
+
+                
+        if out is None:
+            z = x + tau * self.bnoise
+            return (z + 1) - ((z-1)**2 + 4 * tau * self.b).sqrt()
+        else:
+            z_m = x + tau * self.bnoise - 1
+            self.b.multiply(4*tau, out=out)
+            z_m.multiply(z_m, out=z_m)
+            out += z_m
+            out.sqrt(out=out)
+            # z = z_m + 2
+            z_m.sqrt(out=z_m)
+            z_m += 2
+            out *= -1
+            out += z_m
+        
+    
+    def __rmul__(self, scalar):
+        
+        ''' Multiplication of L2NormSquared with a scalar
+        
+        Returns: ScaledFunction
+                        
+        '''
+        
+        return ScaledFunction(self, scalar)     
         
         
     
@@ -79,4 +152,4 @@ if __name__ == '__main__':
 
 
     
-        
\ 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 1946d67..2bb4ca7 100644
--- a/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/L2NormSquared.py
@@ -93,8 +93,8 @@ class L2NormSquared(Function):
             if self.b is None:
                 return x/(1+2*tau)
             else:
-                tmp = x
-                tmp -= self.b
+                tmp = x.subtract(self.b)
+                #tmp -= self.b
                 tmp /= (1+2*tau)
                 tmp += self.b
                 return tmp
diff --git a/Wrappers/Python/test/test_functions.py b/Wrappers/Python/test/test_functions.py
index b428c12..af419c7 100644
--- a/Wrappers/Python/test/test_functions.py
+++ b/Wrappers/Python/test/test_functions.py
@@ -9,7 +9,7 @@ Created on Sat Mar  2 19:24:37 2019
 
 import numpy as np
 #from ccpi.optimisation.funcs import Function
-from ccpi.optimisation.functions import Function
+from ccpi.optimisation.functions import Function, KullbackLeibler
 from ccpi.framework import DataContainer, ImageData, ImageGeometry 
 from ccpi.optimisation.operators import  Identity
 from ccpi.optimisation.operators import BlockOperator
@@ -99,6 +99,7 @@ class TestFunction(unittest.TestCase):
     
     def test_L2NormSquared(self):
         # TESTS for L2 and scalar * L2
+        print ("Test L2NormSquared")
     
         M, N, K = 2,3,5
         ig = ImageGeometry(voxel_num_x=M, voxel_num_y = N, voxel_num_z = K)
@@ -324,7 +325,7 @@ class TestFunction(unittest.TestCase):
         
         a1 = f_no_scaled(U)
         a2 = f_scaled(U)
-        self.assertNumpyArrayAlmostEqual(a1.as_array(),a2.as_array())
+        self.assertNumpyArrayAlmostEqual(a1,a2)
         
         
         tmp = [ el**2 for el in U.containers ]
@@ -341,32 +342,26 @@ class TestFunction(unittest.TestCase):
         f_no_scaled.proximal_conjugate(U, 1, out=z3)
         self.assertBlockDataContainerEqual(z3,z1)
 
-#    
-#    f1 = L2NormSq(alpha=1, b=noisy_data)
-#    print(f1(noisy_data))
-#    
-#    f2 =  L2NormSq(alpha=5, b=noisy_data).composition_with(op2)
-#    print(f2(noisy_data))
-#    
-#    print(f1.gradient(noisy_data).as_array())
-#    print(f2.gradient(noisy_data).as_array())
-##    
-#    print(f1.proximal(noisy_data,1).as_array())
-#    print(f2.proximal(noisy_data,1).as_array())
-#    
-#    
-#    f3 = mixed_L12Norm(alpha = 1).composition_with(op1)
-#    print(f3(noisy_data))
-#            
-#    print(ImageData(op1.direct(noisy_data).power(2).sum(axis=0)).sqrt().sum())
-#    
-#    print( 5*(op2.direct(d) - noisy_data).power(2).sum(), f2(d))
-#    
-#    from functions import mixed_L12Norm as mixed_L12Norm_old
-#    
-#    print(mixed_L12Norm_old(op1,None,alpha)(noisy_data))
-    
-    
-    #        
-
+    def test_KullbackLeibler(self):
+        print ("test_KullbackLeibler")
+        N, M = 2,3
+        ig  = ImageGeometry(N, M)
+        data = ig.allocate(ImageGeometry.RANDOM_INT)
+        x = ig.allocate(ImageGeometry.RANDOM_INT)
+        bnoise = ig.allocate(ImageGeometry.RANDOM_INT)
+        
+        out = ig.allocate()
+        
+        f = KullbackLeibler(data, bnoise=bnoise)
+        
+        grad = f.gradient(x)
+        f.gradient(x, out=out)
+        numpy.testing.assert_array_equal(grad.as_array(), out.as_array())
+        
+        prox = f.proximal(x,1.2)
+        f.proximal(x, 1.2, out=out)
+        numpy.testing.assert_array_equal(prox.as_array(), out.as_array())
         
+        proxc = f.proximal_conjugate(x,1.2)
+        f.proximal_conjugate(x, 1.2, out=out)
+        numpy.testing.assert_array_equal(proxc.as_array(), out.as_array())
\ No newline at end of file