Implementation of FBPD algorithm and TV regularizer func (#49)

* Added FBPD alg for LS+1norm in simple_demo

* Replaced dir_op/adj_op by operator in func

* WIP: implement finite diff op

* FBPD TV working: TV and Norm2 funcs, FiniteDiff2D op #46

* Remove unnecessary FISTA call #46

5 files changed, 188 insertions, 11 deletions

diff --git a/Wrappers/Python/ccpi/reconstruction/algs.py b/Wrappers/Python/ccpi/reconstruction/algs.py
index d957298..be0abb5 100644
--- a/Wrappers/Python/ccpi/reconstruction/algs.py
+++ b/Wrappers/Python/ccpi/reconstruction/algs.py
@@ -95,7 +95,7 @@ def FBPD(x_init, f=None, g=None, h=None, opt=None):
 
     # initialization
     x = x_init
-    y = h.dir_op(x);
+    y = h.op.direct(x);
     
     timing = numpy.zeros(max_iter)
     criter = numpy.zeros(max_iter)
@@ -107,16 +107,16 @@ def FBPD(x_init, f=None, g=None, h=None, opt=None):
     
         # primal forward-backward step
         x_old = x;
-        x = x - tau * ( g.grad(x) + h.adj_op(y) );
+        x = x - tau * ( g.grad(x) + h.op.adjoint(y) );
         x = f.prox(x, tau);
     
         # dual forward-backward step
-        y = y + sigma * h.dir_op(2*x - x_old);
+        y = y + sigma * h.op.direct(2*x - x_old);
         y = y - sigma * h.prox(inv_sigma*y, inv_sigma);   
 
         # time and criterion
         timing[it] = time.time() - t
-        criter[it] = f.fun(x) + g.fun(x) + h.fun(h.dir_op(x));
+        criter[it] = f.fun(x) + g.fun(x) + h.fun(h.op.direct(x));
            
         # stopping rule
         #if np.linalg.norm(x - x_old) < tol * np.linalg.norm(x_old) and it > 10:
diff --git a/Wrappers/Python/ccpi/reconstruction/funcs.py b/Wrappers/Python/ccpi/reconstruction/funcs.py
index c6f29d2..7a93b6e 100644
--- a/Wrappers/Python/ccpi/reconstruction/funcs.py
+++ b/Wrappers/Python/ccpi/reconstruction/funcs.py
@@ -17,16 +17,50 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+from ccpi.reconstruction.ops import Identity, FiniteDiff2D
+import numpy
+
 
 class BaseFunction(object):
-    def __init__(self):   pass
+    def __init__(self):
+        self.op = Identity()
     def fun(self,x):      return 0
     def grad(self,x):     return 0
     def prox(self,x,tau): return x
-    def dir_op(self,x):   return x
-    def adj_op(self,x):   return x
 
-# Define a class for 2-norm
+class Norm2(BaseFunction):
+    
+    def __init__(self, 
+                 gamma=1.0, 
+                 direction=None):
+        super(Norm2, self).__init__()
+        self.gamma     = gamma;
+        self.direction = direction; 
+    
+    def fun(self, x):
+        
+        xx = numpy.sqrt(numpy.sum(numpy.square(x.as_array()), self.direction, keepdims=True))
+        p  = numpy.sum(self.gamma*xx)        
+        
+        return p
+    
+    def prox(self, x, tau):
+
+        xx = numpy.sqrt(numpy.sum( numpy.square(x.as_array()), self.direction, keepdims=True ))
+        xx = numpy.maximum(0, 1 - tau*self.gamma / xx)
+        p  = x.as_array() * xx
+        
+        return type(x)(p,geometry=x.geometry)
+
+class TV2D(Norm2):
+    
+    def __init__(self, gamma):
+        super(TV2D,self).__init__(gamma, 2)
+        self.op = FiniteDiff2D()
+        self.L = self.op.get_max_sing_val()
+        
+
+# Define a class for squared 2-norm
 class Norm2sq(BaseFunction):
     '''
     f(x) = c*||A*x-b||_2^2
diff --git a/Wrappers/Python/ccpi/reconstruction/ops.py b/Wrappers/Python/ccpi/reconstruction/ops.py
index 82f18ac..c21ff06 100644
--- a/Wrappers/Python/ccpi/reconstruction/ops.py
+++ b/Wrappers/Python/ccpi/reconstruction/ops.py
@@ -93,6 +93,39 @@ class Identity(Operator):
     def get_max_sing_val(self):
         return self.s1
 
+class FiniteDiff2D(Operator):
+    def __init__(self):
+        self.s1 = 8.0
+        super(FiniteDiff2D, self).__init__()
+        
+    def direct(self,x):
+        '''Forward differences with Neumann BC.'''
+        d1 = numpy.zeros_like(x.as_array())
+        d1[:,:-1] = x.as_array()[:,1:] - x.as_array()[:,:-1]
+        d2 = numpy.zeros_like(x.as_array())
+        d2[:-1,:] = x.as_array()[1:,:] - x.as_array()[:-1,:]
+        d = numpy.stack((d1,d2),2)
+        
+        return type(x)(d,geometry=x.geometry)
+    
+    def adjoint(self,x):
+        '''Backward differences, Newumann BC.'''
+        Nrows, Ncols, Nchannels = x.as_array().shape
+        zer = numpy.zeros((Nrows,1))
+        xxx = x.as_array()[:,:-1,0]
+        h = numpy.concatenate((zer,xxx), 1) - numpy.concatenate((xxx,zer), 1)
+        
+        zer = numpy.zeros((1,Ncols))
+        xxx = x.as_array()[:-1,:,1]
+        v = numpy.concatenate((zer,xxx), 0) - numpy.concatenate((xxx,zer), 0)
+        return type(x)(h + v,geometry=x.geometry)
+    
+    def size(self):
+        return NotImplemented
+    
+    def get_max_sing_val(self):
+        return self.s1
+
 
 def PowerMethodNonsquare(op,numiters):
     # Initialise random
diff --git a/Wrappers/Python/test/simple_demo_tv.py b/Wrappers/Python/test/simple_demo_tv.py
new file mode 100644
index 0000000..0ae550a
--- /dev/null
+++ b/Wrappers/Python/test/simple_demo_tv.py
@@ -0,0 +1,100 @@
+
+
+import sys
+
+sys.path.append("..")
+
+from ccpi.framework import *
+from ccpi.reconstruction.algs import *
+from ccpi.reconstruction.funcs import *
+from ccpi.reconstruction.ops import *
+from ccpi.reconstruction.astra_ops import *
+from ccpi.reconstruction.geoms import *
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+test_case = 2   # 1=parallel2D, 2=cone2D
+
+# Set up phantom
+N = 128
+
+x = np.zeros((N,N))
+x[round(N/4):round(3*N/4),round(N/4):round(3*N/4)] = 1.0
+x[round(N/8):round(7*N/8),round(3*N/8):round(5*N/8)] = 2.0
+
+plt.imshow(x)
+plt.show()
+
+vg = VolumeGeometry(N,N,None, 1,1,None)
+
+Phantom = VolumeData(x,geometry=vg)
+
+# Set up measurement geometry
+angles_num = 20; # angles number
+
+if test_case==1:
+    angles = np.linspace(0,np.pi,angles_num,endpoint=False)
+elif test_case==2:
+    angles = np.linspace(0,2*np.pi,angles_num,endpoint=False)
+else:
+    NotImplemented
+
+det_w = 1.0
+det_num = N
+SourceOrig = 200
+OrigDetec = 0
+
+# Parallelbeam geometry test
+if test_case==1:
+    pg = SinogramGeometry('parallel',
+                          '2D',
+                          angles,
+                          det_num,det_w)
+elif test_case==2:
+    pg = SinogramGeometry('cone',
+                          '2D',
+                          angles,
+                          det_num,
+                          det_w,
+                          dist_source_center=SourceOrig, 
+                          dist_center_detector=OrigDetec)
+
+# ASTRA operator using volume and sinogram geometries
+Aop = AstraProjectorSimple(vg, pg, 'gpu')
+
+# Unused old astra projector without geometry
+# Aop_old = AstraProjector(det_w, det_num, SourceOrig, 
+#                      OrigDetec, angles, 
+#                      N,'fanbeam','gpu') 
+
+# Try forward and backprojection
+b = Aop.direct(Phantom)
+out2 = Aop.adjoint(b)
+
+plt.imshow(b.array)
+plt.show()
+
+plt.imshow(out2.array)
+plt.show()
+
+# Create least squares object instance with projector and data.
+f = Norm2sq(Aop,b,c=0.5)
+
+# Initial guess
+x_init = VolumeData(np.zeros(x.shape),geometry=vg)
+
+# Now least squares plus 1-norm regularization
+lam = 1
+g0 = TV2D(lam)
+
+
+# Run FBPD=Forward Backward Primal Dual method on least squares plus 1-norm
+opt = {'tol': 1e-4, 'iter': 10000}
+x_fbpd1, it_fbpd1, timing_fbpd1, criter_fbpd1 = FBPD(x_init,None,f,g0,opt=opt)
+
+plt.imshow(x_fbpd1.array)
+plt.show()
+
+plt.semilogy(criter_fbpd1)
+plt.show()
\ No newline at end of file
diff --git a/Wrappers/Python/wip/simple_demo.py b/Wrappers/Python/wip/simple_demo.py
index 7c5f43e..21796b8 100644
--- a/Wrappers/Python/wip/simple_demo.py
+++ b/Wrappers/Python/wip/simple_demo.py
@@ -56,7 +56,7 @@ elif test_case==2:
                           dist_center_detector=OrigDetec)
 
 # ASTRA operator using volume and sinogram geometries
-Aop = AstraProjectorSimple(vg, pg, 'cpu')
+Aop = AstraProjectorSimple(vg, pg, 'gpu')
 
 # Unused old astra projector without geometry
 # Aop_old = AstraProjector(det_w, det_num, SourceOrig, 
@@ -95,5 +95,15 @@ x_fista1, it1, timing1, criter1 = FISTA(x_init, f, g0)
 plt.imshow(x_fista1.array)
 plt.show()
 
-# Delete projector
-#Aop.delete()
+plt.semilogy(criter1)
+plt.show()
+
+# Run FBPD=Forward Backward Primal Dual method on least squares plus 1-norm
+opt = {'tol': 1e-4, 'iter': 10000}
+x_fbpd1, it_fbpd1, timing_fbpd1, criter_fbpd1 = FBPD(x_init,None,f,g0,opt=opt)
+
+plt.imshow(x_fbpd1.array)
+plt.show()
+
+plt.semilogy(criter_fbpd1)
+plt.show()
\ No newline at end of file