Merge branch 'add_data' into demos

18 files changed, 346 insertions, 159 deletions

diff --git a/Wrappers/Python/ccpi/data/__init__.py b/Wrappers/Python/ccpi/data/__init__.py
deleted file mode 100644
index 2884108..0000000
--- a/Wrappers/Python/ccpi/data/__init__.py
+++ /dev/null
@@ -1,66 +0,0 @@
-# -*- coding: utf-8 -*-
-#   This work is part of the Core Imaging Library developed by
-#   Visual Analytics and Imaging System Group of the Science Technology
-#   Facilities Council, STFC
-
-#   Copyright 2018 Edoardo Pasca
-
-#   Licensed under the Apache License, Version 2.0 (the "License");
-#   you may not use this file except in compliance with the License.
-#   You may obtain a copy of the License at
-
-#       http://www.apache.org/licenses/LICENSE-2.0
-
-#   Unless required by applicable law or agreed to in writing, software
-#   distributed under the License is distributed on an "AS IS" BASIS,
-#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#   See the License for the specific language governing permissions and
-#   limitations under the License.
-
-
-from ccpi.framework import ImageData
-import numpy
-from PIL import Image
-import os
-import os.path 
-
-data_dir = os.path.abspath(os.path.dirname(__file__))
-          
-def camera(**kwargs):
-
-    tmp = Image.open(os.path.join(data_dir, 'camera.png'))
-    
-    size = kwargs.get('size',(512, 512))
-    
-    data = numpy.array(tmp.resize(size))
-        
-    data = data/data.max()
-    
-    return ImageData(data) 
-
-
-def boat(**kwargs):
-
-    tmp = Image.open(os.path.join(data_dir, 'boat.tiff'))
-    
-    size = kwargs.get('size',(512, 512))
-    
-    data = numpy.array(tmp.resize(size))
-        
-    data = data/data.max()
-    
-    return ImageData(data)  
-
-
-def peppers(**kwargs):
-
-    tmp = Image.open(os.path.join(data_dir, 'peppers.tiff'))
-    
-    size = kwargs.get('size',(512, 512))
-    
-    data = numpy.array(tmp.resize(size))
-        
-    data = data/data.max()
-    
-    return ImageData(data)    
-    
diff --git a/Wrappers/Python/ccpi/framework/TestData.py b/Wrappers/Python/ccpi/framework/TestData.py
new file mode 100755
index 0000000..752bc13
--- /dev/null
+++ b/Wrappers/Python/ccpi/framework/TestData.py
@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-

+from ccpi.framework import ImageData, ImageGeometry

+import numpy

+from PIL import Image

+import os

+import os.path 

+

+data_dir = os.path.abspath(os.path.join(

+        os.path.dirname(__file__),

+        '../data/')

+)

+

+class TestData(object):

+    BOAT = 'boat.tiff'

+    CAMERA = 'camera.png'

+    PEPPERS = 'peppers.tiff'

+    RESOLUTION_CHART = 'resolution_chart.tiff'

+    SIMPLE_PHANTOM_2D = 'simple_jakobs_phantom'

+    

+    def __init__(self, **kwargs):

+        self.data_dir = kwargs.get('data_dir', data_dir)

+        

+    def load(self, which, size=(512,512), scale=(0,1), **kwargs):

+        if which not in [TestData.BOAT, TestData.CAMERA, 

+                         TestData.PEPPERS, TestData.RESOLUTION_CHART,

+                         TestData.SIMPLE_PHANTOM_2D]:

+            raise ValueError('Unknown TestData {}.'.format(which))

+        if which == TestData.SIMPLE_PHANTOM_2D:

+            N = size[0]

+            M = size[1]

+            sdata = numpy.zeros((N,M))

+            sdata[round(N/4):round(3*N/4),round(N/4):round(3*N/4)] = 0.5

+            sdata[round(M/8):round(7*M/8),round(3*M/8):round(5*M/8)] = 1

+            ig = ImageGeometry(voxel_num_x = N, voxel_num_y = M, dimension_labels=[ImageGeometry.HORIZONTAL_X, ImageGeometry.HORIZONTAL_Y])

+            data = ig.allocate()

+            data.fill(sdata)

+        else:

+            tmp = Image.open(os.path.join(self.data_dir, which))

+            print (tmp)

+            bands = tmp.getbands()

+            if len(bands) > 1:

+                ig = ImageGeometry(voxel_num_x=size[0], voxel_num_y=size[1], channels=len(bands), 

+                dimension_labels=[ImageGeometry.HORIZONTAL_X, ImageGeometry.HORIZONTAL_Y, ImageGeometry.CHANNEL])

+                data = ig.allocate()

+            else:

+                ig = ImageGeometry(voxel_num_x = size[0], voxel_num_y = size[1], dimension_labels=[ImageGeometry.HORIZONTAL_X, ImageGeometry.HORIZONTAL_Y])

+                data = ig.allocate()

+            data.fill(numpy.array(tmp.resize((size[1],size[0]))))

+            if scale is not None:

+                dmax = data.as_array().max()

+                dmin = data.as_array().min()

+                # scale 0,1

+                data = (data -dmin) / (dmax - dmin)

+                if scale != (0,1):

+                    #data = (data-dmin)/(dmax-dmin) * (scale[1]-scale[0]) +scale[0])

+                    data *= (scale[1]-scale[0])

+                    data += scale[0]

+        print ("data.geometry", data.geometry)

+        return data

+

+    def camera(**kwargs):

+    

+        tmp = Image.open(os.path.join(data_dir, 'camera.png'))

+        

+        size = kwargs.get('size',(512, 512))

+        

+        data = numpy.array(tmp.resize(size))

+            

+        data = data/data.max()

+        

+        return ImageData(data) 

+    

+    

+    def boat(**kwargs):

+    

+        tmp = Image.open(os.path.join(data_dir, 'boat.tiff'))

+        

+        size = kwargs.get('size',(512, 512))

+        

+        data = numpy.array(tmp.resize(size))

+            

+        data = data/data.max()

+        

+        return ImageData(data)  

+    

+    

+    def peppers(**kwargs):

+    

+        tmp = Image.open(os.path.join(data_dir, 'peppers.tiff'))

+        

+        size = kwargs.get('size',(512, 512))

+        

+        data = numpy.array(tmp.resize(size))

+            

+        data = data/data.max()

+        

+        return ImageData(data)    

+    

diff --git a/Wrappers/Python/ccpi/framework/__init__.py b/Wrappers/Python/ccpi/framework/__init__.py
index 229edb5..8926897 100755
--- a/Wrappers/Python/ccpi/framework/__init__.py
+++ b/Wrappers/Python/ccpi/framework/__init__.py
@@ -24,3 +24,5 @@ from .framework import DataProcessor
 from .framework import AX, PixelByPixelDataProcessor, CastDataContainer

 from .BlockDataContainer import BlockDataContainer

 from .BlockGeometry import BlockGeometry

+

+from .TestData import TestData

diff --git a/Wrappers/Python/ccpi/framework/framework.py b/Wrappers/Python/ccpi/framework/framework.py
index dbe7d0a..3840f2c 100755
--- a/Wrappers/Python/ccpi/framework/framework.py
+++ b/Wrappers/Python/ccpi/framework/framework.py
@@ -63,7 +63,8 @@ class ImageGeometry(object):
                  center_x=0, 
                  center_y=0, 
                  center_z=0, 
-                 channels=1):
+                 channels=1, 
+                 **kwargs):
         
         self.voxel_num_x = voxel_num_x
         self.voxel_num_y = voxel_num_y
@@ -80,25 +81,44 @@ class ImageGeometry(object):
         if self.channels > 1:            
             if self.voxel_num_z>1:
                 self.length = 4
-                self.shape = (self.channels, self.voxel_num_z, self.voxel_num_y, self.voxel_num_x)
+                shape = (self.channels, self.voxel_num_z, self.voxel_num_y, self.voxel_num_x)
                 dim_labels = [ImageGeometry.CHANNEL, ImageGeometry.VERTICAL,
                 ImageGeometry.HORIZONTAL_Y, ImageGeometry.HORIZONTAL_X]
             else:
                 self.length = 3
-                self.shape = (self.channels, self.voxel_num_y, self.voxel_num_x)
+                shape = (self.channels, self.voxel_num_y, self.voxel_num_x)
                 dim_labels = [ImageGeometry.CHANNEL, ImageGeometry.HORIZONTAL_Y, ImageGeometry.HORIZONTAL_X]
         else:
             if self.voxel_num_z>1:
                 self.length = 3
-                self.shape = (self.voxel_num_z, self.voxel_num_y, self.voxel_num_x)
+                shape = (self.voxel_num_z, self.voxel_num_y, self.voxel_num_x)
                 dim_labels = [ImageGeometry.VERTICAL, ImageGeometry.HORIZONTAL_Y,
                  ImageGeometry.HORIZONTAL_X]
             else:
                 self.length = 2  
-                self.shape = (self.voxel_num_y, self.voxel_num_x)
+                shape = (self.voxel_num_y, self.voxel_num_x)
                 dim_labels = [ImageGeometry.HORIZONTAL_Y, ImageGeometry.HORIZONTAL_X]
         
-        self.dimension_labels = dim_labels
+        labels = kwargs.get('dimension_labels', None)
+        if labels is None:
+            self.shape = shape
+            self.dimension_labels = dim_labels
+        else:
+            order = self.get_order_by_label(labels, dim_labels)
+            if order != [0,1,2]:
+                # resort
+                self.shape = tuple([shape[i] for i in order])
+                self.dimension_labels = labels
+                
+    def get_order_by_label(self, dimension_labels, default_dimension_labels):
+        order = []
+        for i, el in enumerate(dimension_labels):
+            for j, ek in enumerate(default_dimension_labels):
+                if el == ek:
+                    order.append(j)
+                    break
+        return order
+
         
     def get_min_x(self):
         return self.center_x - 0.5*self.voxel_num_x*self.voxel_size_x
@@ -146,7 +166,10 @@ class ImageGeometry(object):
         return repres
     def allocate(self, value=0, dimension_labels=None, **kwargs):
         '''allocates an ImageData according to the size expressed in the instance'''
-        out = ImageData(geometry=self)
+        if dimension_labels is None:
+            out = ImageData(geometry=self, dimension_labels=self.dimension_labels)
+        else:
+            out = ImageData(geometry=self, dimension_labels=dimension_labels)
         if isinstance(value, Number):
             if value != 0:
                 out += value
@@ -164,9 +187,7 @@ class ImageGeometry(object):
                 out.fill(numpy.random.randint(max_value,size=self.shape))
             else:
                 raise ValueError('Value {} unknown'.format(value))
-        if dimension_labels is not None:
-            if dimension_labels != self.dimension_labels:
-                return out.subset(dimensions=dimension_labels)
+
         return out
     # The following methods return 2 members of the class, therefore I 
     # don't think we need to implement them. 
@@ -244,6 +265,8 @@ class AcquisitionGeometry(object):
         self.channels = channels
         self.angle_unit=kwargs.get(AcquisitionGeometry.ANGLE_UNIT, 
                                AcquisitionGeometry.DEGREE)
+
+        # default labels
         if channels > 1:
             if pixel_num_v > 1:
                 shape = (channels, num_of_angles , pixel_num_v, pixel_num_h)
@@ -262,9 +285,31 @@ class AcquisitionGeometry(object):
             else:
                 shape = (num_of_angles, pixel_num_h)
                 dim_labels = [AcquisitionGeometry.ANGLE, AcquisitionGeometry.HORIZONTAL]
-        self.shape = shape
+        
+        labels = kwargs.get('dimension_labels', None)
+        if labels is None:
+            self.shape = shape
+            self.dimension_labels = dim_labels
+        else:
+            if len(labels) != len(dim_labels):
+                raise ValueError('Wrong number of labels. Expected {} got {}'.format(len(dim_labels), len(labels)))
+            order = self.get_order_by_label(labels, dim_labels)
+            if order != [0,1,2]:
+                # resort
+                self.shape = tuple([shape[i] for i in order])
+                self.dimension_labels = labels
+        
+    def get_order_by_label(self, dimension_labels, default_dimension_labels):
+        order = []
+        for i, el in enumerate(dimension_labels):
+            for j, ek in enumerate(default_dimension_labels):
+                if el == ek:
+                    order.append(j)
+                    break
+        return order
+
+
 
-        self.dimension_labels = dim_labels
         
     def clone(self):
         '''returns a copy of the AcquisitionGeometry'''
@@ -292,7 +337,10 @@ class AcquisitionGeometry(object):
         return repres
     def allocate(self, value=0, dimension_labels=None):
         '''allocates an AcquisitionData according to the size expressed in the instance'''
-        out = AcquisitionData(geometry=self)
+        if dimension_labels is None:
+            out = AcquisitionData(geometry=self, dimension_labels=self.dimension_labels)
+        else:
+            out = AcquisitionData(geometry=self, dimension_labels=dimension_labels)
         if isinstance(value, Number):
             if value != 0:
                 out += value
@@ -310,9 +358,7 @@ class AcquisitionGeometry(object):
                 out.fill(numpy.random.randint(max_value,size=self.shape))
             else:
                 raise ValueError('Value {} unknown'.format(value))
-        if dimension_labels is not None:
-            if dimension_labels != self.dimension_labels:
-                return out.subset(dimensions=dimension_labels)
+        
         return out
     
 class DataContainer(object):
@@ -658,7 +704,7 @@ class DataContainer(object):
                 #       geometry=self.geometry)
                 return out
             else:
-                raise ValueError(message(type(self),"Wrong size for data memory: ", out.shape,self.shape))
+                raise ValueError(message(type(self),"Wrong size for data memory: out {} x2 {} expected {}".format( out.shape,x2.shape ,self.shape)))
         elif issubclass(type(out), DataContainer) and isinstance(x2, (int,float,complex)):
             if self.check_dimensions(out):
                 kwargs['out']=out.as_array()
@@ -806,7 +852,7 @@ class ImageData(DataContainer):
         self.geometry = kwargs.get('geometry', None)
         if array is None:
             if self.geometry is not None:
-                shape, dimension_labels = self.get_shape_labels(self.geometry)
+                shape, dimension_labels = self.get_shape_labels(self.geometry, dimension_labels)
                     
                 array = numpy.zeros( shape , dtype=numpy.float32) 
                 super(ImageData, self).__init__(array, deep_copy,
diff --git a/Wrappers/Python/ccpi/optimisation/functions/KullbackLeibler.py b/Wrappers/Python/ccpi/optimisation/functions/KullbackLeibler.py
index d808e63..4dd77cf 100644
--- a/Wrappers/Python/ccpi/optimisation/functions/KullbackLeibler.py
+++ b/Wrappers/Python/ccpi/optimisation/functions/KullbackLeibler.py
@@ -107,16 +107,20 @@ class KullbackLeibler(Function):
             return 0.5*((z + 1) - ((z-1)**2 + 4 * tau * self.b).sqrt())
         else:
             
-            tmp = x + tau * self.bnoise
-            self.b.multiply(4*tau, out=out)            
-            out.add((tmp-1)**2, out=out)
+            #tmp = x + tau * self.bnoise
+            tmp = tau * self.bnoise
+            tmp += x
+            tmp -= 1
+            
+            self.b.multiply(4*tau, out=out)    
+            
+            out.add((tmp)**2, out=out)
             out.sqrt(out=out)
             out *= -1
-            out.add(tmp+1, out=out)
+            tmp += 2
+            out += tmp
             out *= 0.5
-                    
-        
-    
+
     def __rmul__(self, scalar):
         
         ''' Multiplication of L2NormSquared with a scalar
diff --git a/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py b/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py
index 6ffaf70..d98961b 100644
--- a/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py
+++ b/Wrappers/Python/ccpi/optimisation/operators/GradientOperator.py
@@ -14,23 +14,47 @@ from ccpi.optimisation.operators import FiniteDiff, SparseFiniteDiff
 #%%
 
 class Gradient(LinearOperator):
-    
+    CORRELATION_SPACE = "Space"
+    CORRELATION_SPACECHANNEL = "SpaceChannels"
+    # 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, bnd_cond = 'Neumann', **kwargs):
         
         super(Gradient, self).__init__() 
                 
         self.gm_domain = gm_domain # Domain of Grad Operator
         
-        self.correlation = kwargs.get('correlation','Space')
+        self.correlation = kwargs.get('correlation',Gradient.CORRELATION_SPACE)
         
-        if self.correlation=='Space':
+        if self.correlation==Gradient.CORRELATION_SPACE:
             if self.gm_domain.channels>1:
-                self.gm_range = BlockGeometry(*[self.gm_domain for _ in range(self.gm_domain.length-1)] ) 
-                self.ind = numpy.arange(1,self.gm_domain.length)
-            else:    
+                self.gm_range = BlockGeometry(*[self.gm_domain for _ in range(self.gm_domain.length-1)] )
+                if self.gm_domain.length == 4:
+                    # 3D + Channel
+                    # expected Grad_order = ['channels', 'direction_z', 'direction_y', 'direction_x']
+                    expected_order = [ImageGeometry.CHANNEL, ImageGeometry.VERTICAL, ImageGeometry.HORIZONTAL_Y, ImageGeometry.HORIZONTAL_X]
+                else:
+                    # 2D + Channel
+                    # expected Grad_order = ['channels', 'direction_y', 'direction_x']
+                    expected_order = [ImageGeometry.CHANNEL, ImageGeometry.HORIZONTAL_Y, ImageGeometry.HORIZONTAL_X]
+                order = self.gm_domain.get_order_by_label(self.gm_domain.dimension_labels, expected_order)
+                self.ind = order[1:]
+                #self.ind = numpy.arange(1,self.gm_domain.length)
+            else:
+                # no channel info
                 self.gm_range = BlockGeometry(*[self.gm_domain for _ in range(self.gm_domain.length) ] )
-                self.ind = numpy.arange(self.gm_domain.length)
-        elif self.correlation=='SpaceChannels':
+                if self.gm_domain.length == 3:
+                    # 3D
+                    # expected Grad_order = ['direction_z', 'direction_y', 'direction_x']
+                    expected_order = [ImageGeometry.VERTICAL, ImageGeometry.HORIZONTAL_Y, ImageGeometry.HORIZONTAL_X]
+                else:
+                    # 2D
+                    expected_order = [ImageGeometry.HORIZONTAL_Y, ImageGeometry.HORIZONTAL_X]    
+                self.ind = self.gm_domain.get_order_by_label(self.gm_domain.dimension_labels, expected_order)
+                # self.ind = numpy.arange(self.gm_domain.length)
+        elif self.correlation==Gradient.CORRELATION_SPACECHANNEL:
             if self.gm_domain.channels>1:
                 self.gm_range = BlockGeometry(*[self.gm_domain for _ in range(self.gm_domain.length)])
                 self.ind = range(self.gm_domain.length)
@@ -86,10 +110,11 @@ class Gradient(LinearOperator):
     def range_geometry(self):
         return self.gm_range
                                    
-    def norm(self):
+    def norm(self, **kwargs):
 
         x0 = self.gm_domain.allocate('random')
-        self.s1, sall, svec = LinearOperator.PowerMethod(self, 10, x0)
+        iterations = kwargs.get('iterations', 10)
+        self.s1, sall, svec = LinearOperator.PowerMethod(self, iterations, x0)
         return self.s1
     
     def __rmul__(self, scalar):
@@ -136,14 +161,21 @@ if __name__ == '__main__':
     from ccpi.optimisation.operators import Identity, BlockOperator
     
     
-    M, N = 2, 3
+    M, N = 20, 30
     ig = ImageGeometry(M, N)
     arr = ig.allocate('random_int' )
     
     # check direct of Gradient and sparse matrix
     G = Gradient(ig)
+    norm1 = G.norm(iterations=300)
+    print ("should be sqrt(8) {} {}".format(numpy.sqrt(8), norm1))
     G_sp = G.matrix()
+    ig4 = ImageGeometry(M,N, channels=3)
+    G4 = Gradient(ig4, correlation=Gradient.CORRELATION_SPACECHANNEL)
+    norm4 = G4.norm(iterations=300)
+    print ("should be sqrt(12) {} {}".format(numpy.sqrt(12), norm4))
     
+
     res1 = G.direct(arr)
     res1y = numpy.reshape(G_sp[0].toarray().dot(arr.as_array().flatten('F')), ig.shape, 'F')
     
diff --git a/Wrappers/Python/conda-recipe/meta.yaml b/Wrappers/Python/conda-recipe/meta.yaml
index 6564014..9d03220 100644
--- a/Wrappers/Python/conda-recipe/meta.yaml
+++ b/Wrappers/Python/conda-recipe/meta.yaml
@@ -35,6 +35,7 @@ requirements:
     - scipy
     - matplotlib
     - h5py
+    - pillow
   
 about:
   home: http://www.ccpi.ac.uk
diff --git a/Wrappers/Python/ccpi/data/boat.tiff b/Wrappers/Python/data/boat.tiff
index fc1205a..fc1205a 100644
--- a/Wrappers/Python/ccpi/data/boat.tiff
+++ b/Wrappers/Python/data/boat.tiff
diff --git a/Wrappers/Python/ccpi/data/camera.png b/Wrappers/Python/data/camera.png
index 49be869..49be869 100644
--- a/Wrappers/Python/ccpi/data/camera.png
+++ b/Wrappers/Python/data/camera.png
diff --git a/Wrappers/Python/ccpi/data/peppers.tiff b/Wrappers/Python/data/peppers.tiff
index 8c956f8..8c956f8 100644
--- a/Wrappers/Python/ccpi/data/peppers.tiff
+++ b/Wrappers/Python/data/peppers.tiff
diff --git a/Wrappers/Python/data/resolution_chart.tiff b/Wrappers/Python/data/resolution_chart.tiff
new file mode 100755
index 0000000..d09cef3
--- /dev/null
+++ b/Wrappers/Python/data/resolution_chart.tiff
diff --git a/Wrappers/Python/ccpi/data/test_show_data.py b/Wrappers/Python/data/test_show_data.py
index 7325c27..7325c27 100644
--- a/Wrappers/Python/ccpi/data/test_show_data.py
+++ b/Wrappers/Python/data/test_show_data.py
diff --git a/Wrappers/Python/demos/PDHG_examples/PDHG_TV_Denoising_Gaussian.py b/Wrappers/Python/demos/PDHG_examples/PDHG_TV_Denoising_Gaussian.py
index afdb6a2..cba5bcb 100644
--- a/Wrappers/Python/demos/PDHG_examples/PDHG_TV_Denoising_Gaussian.py
+++ b/Wrappers/Python/demos/PDHG_examples/PDHG_TV_Denoising_Gaussian.py
@@ -1,22 +1,20 @@
 #========================================================================
-# Copyright 2019 Science Technology Facilities Council
-# Copyright 2019 University of Manchester
-#
-# This work is part of the Core Imaging Library developed by Science Technology
-# Facilities Council and University of Manchester
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#         http://www.apache.org/licenses/LICENSE-2.0.txt
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
+#  CCP in Tomographic Imaging (CCPi) Core Imaging Library (CIL).
+
+#   Copyright 2017 UKRI-STFC
+#   Copyright 2017 University of Manchester
+
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+
+#   http://www.apache.org/licenses/LICENSE-2.0
+
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
 #=========================================================================
 """ 
 
@@ -50,42 +48,60 @@ from ccpi.optimisation.algorithms import PDHG
 from ccpi.optimisation.operators import BlockOperator, Identity, Gradient
 from ccpi.optimisation.functions import ZeroFunction, L2NormSquared, \
                       MixedL21Norm, BlockFunction
-      
+
+from ccpi.framework import TestData
+import os, sys
+loader = TestData(data_dir=os.path.join(sys.prefix, 'share','ccpi'))
+
 # Load Data                      
-N = 100
+N = 256
+M = 300
 
-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
-data = ImageData(data)
-ig = ImageGeometry(voxel_num_x = N, voxel_num_y = N)
+
+# user can change the size of the input data
+# you can choose between 
+# TestData.PEPPERS 2D + Channel
+# TestData.BOAT 2D 
+# TestData.CAMERA 2D
+# TestData.RESOLUTION_CHART 2D 
+# TestData.SIMPLE_PHANTOM_2D 2D
+data = loader.load(TestData.PEPPERS, size=(N,M), scale=(0,1))
+
+ig = data.geometry
 ag = ig
                       
 # Create Noisy data. Add Gaussian noise
 np.random.seed(10)
-noisy_data = ImageData( data.as_array() + np.random.normal(0, 0.1, size=ig.shape) )
+noisy_data = ImageData( data.as_array() + np.random.normal(0, 0.1, size=data.shape) )
+
+print ("min {} max {}".format(data.as_array().min(), data.as_array().max()))
 
 # Show Ground Truth and Noisy Data
-plt.figure(figsize=(15,15))
-plt.subplot(2,1,1)
+plt.figure()
+plt.subplot(1,3,1)
 plt.imshow(data.as_array())
 plt.title('Ground Truth')
 plt.colorbar()
-plt.subplot(2,1,2)
+plt.subplot(1,3,2)
 plt.imshow(noisy_data.as_array())
 plt.title('Noisy Data')
 plt.colorbar()
+plt.subplot(1,3,3)
+plt.imshow((data - noisy_data).as_array())
+plt.title('diff')
+plt.colorbar()
+
 plt.show()
 
 # Regularisation Parameter
-alpha = 0.2
+alpha = .1
 
 method = '0'
 
 if method == '0':
 
     # Create operators
-    op1 = Gradient(ig)
+    op1 = Gradient(ig, correlation=Gradient.CORRELATION_SPACE)
     op2 = Identity(ig, ag)
 
     # Create BlockOperator
@@ -114,28 +130,33 @@ tau = 1/(sigma*normK**2)
 
 # Setup and Run the PDHG algorithm
 pdhg = PDHG(f=f,g=g,operator=operator, tau=tau, sigma=sigma)
-pdhg.max_iteration = 3000
-pdhg.update_objective_interval = 200
-pdhg.run(3000, verbose=False)
+pdhg.max_iteration = 10000
+pdhg.update_objective_interval = 100
+pdhg.run(1000, verbose=True)
 
 # Show Results
-plt.figure(figsize=(15,15))
-plt.subplot(3,1,1)
+plt.figure()
+plt.subplot(1,3,1)
 plt.imshow(data.as_array())
 plt.title('Ground Truth')
 plt.colorbar()
-plt.subplot(3,1,2)
+plt.clim(0,1)
+plt.subplot(1,3,2)
 plt.imshow(noisy_data.as_array())
 plt.title('Noisy Data')
 plt.colorbar()
-plt.subplot(3,1,3)
+plt.clim(0,1)
+plt.subplot(1,3,3)
 plt.imshow(pdhg.get_output().as_array())
 plt.title('TV Reconstruction')
+plt.clim(0,1)
 plt.colorbar()
 plt.show()
 
-plt.plot(np.linspace(0,N,N), data.as_array()[int(N/2),:], label = 'GTruth')
-plt.plot(np.linspace(0,N,N), pdhg.get_output().as_array()[int(N/2),:], label = 'TV reconstruction')
+plt.plot(np.linspace(0,N,M), noisy_data.as_array()[int(N/2),:], label = 'Noisy data')
+plt.plot(np.linspace(0,N,M), data.as_array()[int(N/2),:], label = 'GTruth')
+plt.plot(np.linspace(0,N,M), pdhg.get_output().as_array()[int(N/2),:], label = 'TV reconstruction')
+
 plt.legend()
 plt.title('Middle Line Profiles')
 plt.show()
diff --git a/Wrappers/Python/setup.py b/Wrappers/Python/setup.py
index bceea46..8bd33a6 100644
--- a/Wrappers/Python/setup.py
+++ b/Wrappers/Python/setup.py
@@ -31,7 +31,7 @@ if  cil_version == '':
 setup(
     name="ccpi-framework",
     version=cil_version,
-    packages=['ccpi' , 'ccpi.io', 'ccpi.data',
+    packages=['ccpi' , 'ccpi.io',
               'ccpi.framework', 'ccpi.optimisation', 
               'ccpi.optimisation.operators',
               'ccpi.optimisation.algorithms',
@@ -39,6 +39,9 @@ setup(
               'ccpi.processors',
               'ccpi.contrib','ccpi.contrib.optimisation',
               'ccpi.contrib.optimisation.algorithms'],
+    data_files = [('share/ccpi', ['data/boat.tiff', 'data/peppers.tiff',
+                                 'data/camera.png', 
+                                 'data/resolution_chart.tiff'])],
 
     # Project uses reStructuredText, so ensure that the docutils get
     # installed or upgraded on the target machine
@@ -53,8 +56,9 @@ setup(
     # zip_safe = False,
 
     # metadata for upload to PyPI
-    author="Edoardo Pasca",
-    author_email="edoardo.pasca@stfc.ac.uk",
+    author="CCPi developers",
+    maintainer="Edoardo Pasca",
+    maintainer_email="edoardo.pasca@stfc.ac.uk",
     description='CCPi Core Imaging Library - Python Framework Module',
     license="Apache v2.0",
     keywords="Python Framework",
diff --git a/Wrappers/Python/test/test_DataContainer.py b/Wrappers/Python/test/test_DataContainer.py
index e92d4c6..4c53df8 100755
--- a/Wrappers/Python/test/test_DataContainer.py
+++ b/Wrappers/Python/test/test_DataContainer.py
@@ -487,6 +487,13 @@ class TestDataContainer(unittest.TestCase):
         self.assertNumpyArrayEqual(vol1.as_array(), numpy.ones(vol.shape) * 4)
 
         self.assertEqual(vol.number_of_dimensions, 3)
+        
+        ig2 = ImageGeometry (voxel_num_x=2,voxel_num_y=3,voxel_num_z=4, 
+                     dimension_labels=[ImageGeometry.HORIZONTAL_X, ImageGeometry.HORIZONTAL_Y,
+                 ImageGeometry.VERTICAL])
+        data = ig2.allocate()
+        self.assertNumpyArrayEqual(numpy.asarray(data.shape), numpy.asarray(ig2.shape))
+        self.assertNumpyArrayEqual(numpy.asarray(data.shape), data.as_array().shape)
 
     def test_AcquisitionData(self):
         sgeometry = AcquisitionGeometry(dimension=2, angles=numpy.linspace(0, 180, num=10),
@@ -494,6 +501,29 @@ class TestDataContainer(unittest.TestCase):
                                         pixel_num_h=5, channels=2)
         sino = AcquisitionData(geometry=sgeometry)
         self.assertEqual(sino.shape, (2, 10, 3, 5))
+        
+        ag = AcquisitionGeometry (pixel_num_h=2,pixel_num_v=3,channels=4, dimension=2, angles=numpy.linspace(0, 180, num=10),
+                                        geom_type='parallel', )
+        print (ag.shape)
+        print (ag.dimension_labels)
+        
+        data = ag.allocate()
+        self.assertNumpyArrayEqual(numpy.asarray(data.shape), numpy.asarray(ag.shape))
+        self.assertNumpyArrayEqual(numpy.asarray(data.shape), data.as_array().shape)
+        
+        print (data.shape, ag.shape, data.as_array().shape)
+        
+        ag2 = AcquisitionGeometry (pixel_num_h=2,pixel_num_v=3,channels=4, dimension=2, angles=numpy.linspace(0, 180, num=10),
+                                                geom_type='parallel', 
+                                                dimension_labels=[AcquisitionGeometry.VERTICAL ,
+                         AcquisitionGeometry.ANGLE, AcquisitionGeometry.HORIZONTAL, AcquisitionGeometry.CHANNEL])
+        
+        data = ag2.allocate()
+        print (data.shape, ag2.shape, data.as_array().shape)
+        self.assertNumpyArrayEqual(numpy.asarray(data.shape), numpy.asarray(ag2.shape))
+        self.assertNumpyArrayEqual(numpy.asarray(data.shape), data.as_array().shape)
+
+
     def test_ImageGeometry_allocate(self):
         vgeometry = ImageGeometry(voxel_num_x=4, voxel_num_y=3, channels=2)
         image = vgeometry.allocate()
diff --git a/Wrappers/Python/test/test_Gradient.py b/Wrappers/Python/test/test_Gradient.py
index c6b2d2e..6d9d3be 100755
--- a/Wrappers/Python/test/test_Gradient.py
+++ b/Wrappers/Python/test/test_Gradient.py
@@ -84,3 +84,18 @@ class TestGradient(unittest.TestCase):
         res = G2D.direct(u4)  
         print(res[0].as_array())
         print(res[1].as_array())
+
+        M, N = 20, 30
+        ig = ImageGeometry(M, N)
+        arr = ig.allocate('random_int' )
+        
+        # check direct of Gradient and sparse matrix
+        G = Gradient(ig)
+        norm1 = G.norm(iterations=300)
+        print ("should be sqrt(8) {} {}".format(numpy.sqrt(8), norm1))
+        numpy.testing.assert_almost_equal(norm1, numpy.sqrt(8), decimal=1)
+        ig4 = ImageGeometry(M,N, channels=3)
+        G4 = Gradient(ig4, correlation=Gradient.CORRELATION_SPACECHANNEL)
+        norm4 = G4.norm(iterations=300)
+        print ("should be sqrt(12) {} {}".format(numpy.sqrt(12), norm4))
+        numpy.testing.assert_almost_equal(norm4, numpy.sqrt(12), decimal=1)
diff --git a/Wrappers/Python/test/test_NexusReader.py b/Wrappers/Python/test/test_NexusReader.py
index 55543ba..a498d71 100755
--- a/Wrappers/Python/test/test_NexusReader.py
+++ b/Wrappers/Python/test/test_NexusReader.py
@@ -21,67 +21,67 @@ class TestNexusReader(unittest.TestCase):
     def tearDown(self):
         os.remove(self.filename)
 
-
-    def testGetDimensions(self):
+    def testAll(self):
+    # def testGetDimensions(self):
         nr = NexusReader(self.filename)
         self.assertEqual(nr.get_sinogram_dimensions(), (135, 91, 160), "Sinogram dimensions are not correct")
         
-    def testGetProjectionDimensions(self):
+    # def testGetProjectionDimensions(self):
         nr = NexusReader(self.filename)
         self.assertEqual(nr.get_projection_dimensions(), (91,135,160), "Projection dimensions are not correct")        
         
-    def testLoadProjectionWithoutDimensions(self):
+    # def testLoadProjectionWithoutDimensions(self):
         nr = NexusReader(self.filename)
         projections = nr.load_projection()        
         self.assertEqual(projections.shape, (91,135,160), "Loaded projection data dimensions are not correct")        
 
-    def testLoadProjectionWithDimensions(self):
+    # def testLoadProjectionWithDimensions(self):
         nr = NexusReader(self.filename)
         projections = nr.load_projection((slice(0,1), slice(0,135), slice(0,160)))        
         self.assertEqual(projections.shape, (1,135,160), "Loaded projection data dimensions are not correct")        
             
-    def testLoadProjectionCompareSingle(self):
+    # def testLoadProjectionCompareSingle(self):
         nr = NexusReader(self.filename)
         projections_full = nr.load_projection()
         projections_part = nr.load_projection((slice(0,1), slice(0,135), slice(0,160))) 
         numpy.testing.assert_array_equal(projections_part, projections_full[0:1,:,:])
         
-    def testLoadProjectionCompareMulti(self):
+    # def testLoadProjectionCompareMulti(self):
         nr = NexusReader(self.filename)
         projections_full = nr.load_projection()
         projections_part = nr.load_projection((slice(0,3), slice(0,135), slice(0,160))) 
         numpy.testing.assert_array_equal(projections_part, projections_full[0:3,:,:])
         
-    def testLoadProjectionCompareRandom(self):
+    # def testLoadProjectionCompareRandom(self):
         nr = NexusReader(self.filename)
         projections_full = nr.load_projection()
         projections_part = nr.load_projection((slice(1,8), slice(5,10), slice(8,20))) 
         numpy.testing.assert_array_equal(projections_part, projections_full[1:8,5:10,8:20])                
         
-    def testLoadProjectionCompareFull(self):
+    # def testLoadProjectionCompareFull(self):
         nr = NexusReader(self.filename)
         projections_full = nr.load_projection()
         projections_part = nr.load_projection((slice(None,None), slice(None,None), slice(None,None))) 
         numpy.testing.assert_array_equal(projections_part, projections_full[:,:,:])
         
-    def testLoadFlatCompareFull(self):
+    # def testLoadFlatCompareFull(self):
         nr = NexusReader(self.filename)
         flats_full = nr.load_flat()
         flats_part = nr.load_flat((slice(None,None), slice(None,None), slice(None,None))) 
         numpy.testing.assert_array_equal(flats_part, flats_full[:,:,:])
         
-    def testLoadDarkCompareFull(self):
+    # def testLoadDarkCompareFull(self):
         nr = NexusReader(self.filename)
         darks_full = nr.load_dark()
         darks_part = nr.load_dark((slice(None,None), slice(None,None), slice(None,None))) 
         numpy.testing.assert_array_equal(darks_part, darks_full[:,:,:])
         
-    def testProjectionAngles(self):
+    # def testProjectionAngles(self):
         nr = NexusReader(self.filename)
         angles = nr.get_projection_angles()
         self.assertEqual(angles.shape, (91,), "Loaded projection number of angles are not correct")        
         
-    def test_get_acquisition_data_subset(self):
+    # def test_get_acquisition_data_subset(self):
         nr = NexusReader(self.filename)
         key = nr.get_image_keys()
         sl = nr.get_acquisition_data_subset(0,10)
diff --git a/Wrappers/Python/test/test_functions.py b/Wrappers/Python/test/test_functions.py
index af419c7..082548b 100644
--- a/Wrappers/Python/test/test_functions.py
+++ b/Wrappers/Python/test/test_functions.py
@@ -299,7 +299,7 @@ class TestFunction(unittest.TestCase):
         
         A = 0.5 * Identity(ig)
         old_chisq = Norm2sq(A, b, 1.0)
-        new_chisq = FunctionOperatorComposition(A, L2NormSquared(b=b))
+        new_chisq = FunctionOperatorComposition(L2NormSquared(b=b),A)
 
         yold = old_chisq(u)
         ynew = new_chisq(u)