fix memopt proximal

2 files changed, 21 insertions, 9 deletions

diff --git a/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py b/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py
index bf627a5..0836324 100644
--- a/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/BlockFunction.py
@@ -93,16 +93,28 @@ class BlockFunction(Function):
         
         
         '''
+        
+        if out is None:
                 
-        out = [None]*self.length
-        if isinstance(tau, Number):
-            for i in range(self.length):
-                out[i] = self.functions[i].proximal(x.get_item(i), tau)
+            out = [None]*self.length
+            if isinstance(tau, Number):
+                for i in range(self.length):
+                    out[i] = self.functions[i].proximal(x.get_item(i), tau)
+            else:
+                for i in range(self.length):
+                    out[i] = self.functions[i].proximal(x.get_item(i), tau.get_item(i))                        
+            
+            return BlockDataContainer(*out)  
+        
         else:
-            for i in range(self.length):
-                out[i] = self.functions[i].proximal(x.get_item(i), tau.get_item(i))                        
-
-        return BlockDataContainer(*out)     
+            if isinstance(tau, Number):
+                for i in range(self.length):
+                    self.functions[i].proximal(x.get_item(i), tau, out[i])
+            else:
+                for i in range(self.length):
+                    self.functions[i].proximal(x.get_item(i), tau.get_item(i), out[i])
+            
+            
     
     def gradient(self,x, out=None):
         
diff --git a/Wrappers/Python/wip/pdhg_TV_denoising.py b/Wrappers/Python/wip/pdhg_TV_denoising.py
index 872d832..a8e25d2 100755
--- a/Wrappers/Python/wip/pdhg_TV_denoising.py
+++ b/Wrappers/Python/wip/pdhg_TV_denoising.py
@@ -6,7 +6,7 @@ Created on Fri Feb 22 14:53:03 2019
 @author: evangelos
 """
 
-from ccpi.framework import ImageData, ImageGeometry, BlockDataContainer
+from ccpi.framework import ImageData, ImageGeometry
 
 import numpy as np 
 import numpy