Astra nexus (#9)

* Add 3D ASTRA operator and nexus demo

* Add astra sophiabeads 3d demo

* 2D astra scaling working simple case, not sophiabeads

* Simplied bp scaling, now seems to work 2D incl sophiabeads

* MC correct astra fp and bp scaling

* Fixed also ASTRA BP scaling in 3D, works sophiabeads 3D

* Minor tidy up

3 files changed, 222 insertions, 11 deletions

diff --git a/Wrappers/Python/ccpi/astra/ops.py b/Wrappers/Python/ccpi/astra/ops.py
index cd0ef9e..2e9d6af 100755
--- a/Wrappers/Python/ccpi/astra/ops.py
+++ b/Wrappers/Python/ccpi/astra/ops.py
@@ -17,11 +17,11 @@
 
 from ccpi.optimisation.ops import Operator
 import numpy
-import astra
 from ccpi.framework import AcquisitionData, ImageData, DataContainer
 from ccpi.optimisation.ops import PowerMethodNonsquare
 from ccpi.astra.processors import AstraForwardProjector, AstraBackProjector, \
-     AstraForwardProjectorMC, AstraBackProjectorMC
+     AstraForwardProjectorMC, AstraBackProjectorMC, AstraForwardProjector3D, \
+     AstraBackProjector3D
 
 class AstraProjectorSimple(Operator):
     """ASTRA projector modified to use DataSet and geometry."""
@@ -74,6 +74,58 @@ class AstraProjectorSimple(Operator):
         return DataContainer(numpy.random.randn(inputsize[0],
                                                 inputsize[1]))
 
+class AstraProjector3DSimple(Operator):
+    """ASTRA projector modified to use DataSet and geometry."""
+    def __init__(self, geomv, geomp):
+        super(AstraProjector3DSimple, self).__init__()
+        
+        # Store volume and sinogram geometries.
+        self.sinogram_geometry = geomp
+        self.volume_geometry = geomv
+        
+        self.fp = AstraForwardProjector3D(volume_geometry=geomv,
+                                        sinogram_geometry=geomp,
+                                        output_axes_order=['vertical','angle','horizontal'])
+        
+        self.bp = AstraBackProjector3D(volume_geometry=geomv,
+                                        sinogram_geometry=geomp,
+                                        output_axes_order=['vertical','horizontal_y','horizontal_x'])
+                
+        # Initialise empty for singular value.
+        self.s1 = None
+    
+    def direct(self, IM):
+        self.fp.set_input(IM)
+        out = self.fp.get_output()
+        return out
+    
+    def adjoint(self, DATA):
+        self.bp.set_input(DATA)
+        out = self.bp.get_output()
+        return out
+    
+    #def delete(self):
+    #    astra.data2d.delete(self.proj_id)
+    
+    def get_max_sing_val(self):
+        self.s1, sall, svec = PowerMethodNonsquare(self,10)
+        return self.s1
+    
+    def size(self):
+        # Only implemented for 2D
+        return ( (self.sinogram_geometry.angles.size, \
+                  self.sinogram_geometry.pixel_num_h, \
+                  self.sinogram_geometry.pixel_num_v,), \
+                 (self.volume_geometry.voxel_num_x, \
+                  self.volume_geometry.voxel_num_y, \
+                  self.volume_geometry.voxel_num_z) )
+    
+    def create_image_data(self):
+        inputsize = self.size()[1]
+        return DataContainer(numpy.random.randn(inputsize[2],
+                                                inputsize[1],
+                                                inputsize[0]))
+
 
 class AstraProjectorMC(Operator):
     """ASTRA Multichannel projector"""
diff --git a/Wrappers/Python/ccpi/astra/processors.py b/Wrappers/Python/ccpi/astra/processors.py
index db52a6b..855f890 100644
--- a/Wrappers/Python/ccpi/astra/processors.py
+++ b/Wrappers/Python/ccpi/astra/processors.py
@@ -74,8 +74,15 @@ class AstraForwardProjector(DataProcessor):
         sinogram_id, DATA.array = astra.create_sino(IM.as_array(), 
                                                            self.proj_id)
         astra.data2d.delete(sinogram_id)
-        #return AcquisitionData(array=DATA, geometry=self.sinogram_geometry)
-        return DATA
+        
+        if self.device == 'cpu':
+            return DATA
+        else:
+            if self.sinogram_geometry.geom_type == 'cone':
+                return DATA
+            else:
+                 scaling = 1.0/self.volume_geometry.voxel_size_x
+                 return scaling*DATA
 
 class AstraBackProjector(DataProcessor):
     '''AstraBackProjector
@@ -145,7 +152,12 @@ class AstraBackProjector(DataProcessor):
         rec_id, IM.array = astra.create_backprojection(DATA.as_array(),
                             self.proj_id)
         astra.data2d.delete(rec_id)
-        return IM
+        
+        if self.device == 'cpu':
+            return IM
+        else:
+            scaling = self.volume_geometry.voxel_size_x**3  
+            return scaling*IM
 
 class AstraForwardProjectorMC(AstraForwardProjector):
     '''AstraForwardProjector Multi channel
@@ -173,7 +185,15 @@ class AstraForwardProjectorMC(AstraForwardProjector):
             sinogram_id, DATA.as_array()[k] = astra.create_sino(IM.as_array()[k], 
                                                            self.proj_id)
             astra.data2d.delete(sinogram_id)
-        return DATA
+        
+        if self.device == 'cpu':
+            return DATA
+        else:
+            if self.sinogram_geometry.geom_type == 'cone':
+                return DATA
+            else:
+                 scaling = (1.0/self.volume_geometry.voxel_size_x) 
+                 return scaling*DATA
 
 class AstraBackProjectorMC(AstraBackProjector):
     '''AstraBackProjector Multi channel
@@ -202,4 +222,142 @@ class AstraBackProjectorMC(AstraBackProjector):
                     DATA.as_array()[k], 
                     self.proj_id)
             astra.data2d.delete(rec_id)
-        return IM
+        
+        if self.device == 'cpu':
+            return IM
+        else:
+            scaling = self.volume_geometry.voxel_size_x**3  
+            return scaling*IM
+
+class AstraForwardProjector3D(DataProcessor):
+    '''AstraForwardProjector3D
+    
+    Forward project ImageData to AcquisitionData using ASTRA proj_geom and 
+    vol_geom.
+    
+    Input: ImageData
+    Parameter: proj_geom, vol_geom
+    Output: AcquisitionData
+    '''
+    
+    def __init__(self,
+                 volume_geometry=None,
+                 sinogram_geometry=None,
+                 proj_geom=None,
+                 vol_geom=None,
+                 output_axes_order=None):
+        kwargs = {
+                  'volume_geometry'  : volume_geometry,
+                  'sinogram_geometry'  : sinogram_geometry,
+                  'proj_geom'  : proj_geom,
+                  'vol_geom'  : vol_geom,
+                  'output_axes_order'  : output_axes_order
+                  }
+        
+        #DataProcessor.__init__(self, **kwargs)
+        super(AstraForwardProjector3D, self).__init__(**kwargs)
+        
+        self.set_ImageGeometry(volume_geometry)
+        self.set_AcquisitionGeometry(sinogram_geometry)
+        
+        # Set up ASTRA Volume and projection geometry, not to be stored in self
+        vol_geom, proj_geom = convert_geometry_to_astra(self.volume_geometry,
+                                                        self.sinogram_geometry)
+        
+        # Also store ASTRA geometries
+        self.vol_geom = vol_geom
+        self.proj_geom = proj_geom
+    
+    def check_input(self, dataset):
+        if dataset.number_of_dimensions == 3:
+               return True
+        else:
+            raise ValueError("Expected input dimensions 3, got {0}"\
+                             .format(dataset.number_of_dimensions))
+    
+    def set_ImageGeometry(self, volume_geometry):
+        self.volume_geometry = volume_geometry
+    
+    def set_AcquisitionGeometry(self, sinogram_geometry):
+        self.sinogram_geometry = sinogram_geometry
+    
+    def set_vol_geom(self, vol_geom):
+        self.vol_geom = vol_geom
+    
+    def set_AcquisitionGeometry(self, sinogram_geometry):
+        self.sinogram_geometry = sinogram_geometry
+    
+    def process(self):
+        IM = self.get_input()
+        DATA = AcquisitionData(geometry=self.sinogram_geometry,
+                               dimension_labels=self.output_axes_order)
+        sinogram_id, DATA.array = astra.create_sino3d_gpu(IM.as_array(), 
+                                                           self.proj_geom,
+                                                           self.vol_geom)
+        astra.data3d.delete(sinogram_id)
+        # 3D CUDA FP does not need scaling
+        return DATA
+
+class AstraBackProjector3D(DataProcessor):
+    '''AstraBackProjector3D
+    
+    Back project AcquisitionData to ImageData using ASTRA proj_geom, vol_geom.
+    
+    Input: AcquisitionData
+    Parameter: proj_geom, vol_geom
+    Output: ImageData
+    '''
+    
+    def __init__(self,
+                 volume_geometry=None,
+                 sinogram_geometry=None,
+                 proj_geom=None,
+                 vol_geom=None,
+                 output_axes_order=None):
+        kwargs = {
+                  'volume_geometry'  : volume_geometry,
+                  'sinogram_geometry'  : sinogram_geometry,
+                  'proj_geom'  : proj_geom,
+                  'vol_geom'  : vol_geom,
+                  'output_axes_order'  : output_axes_order
+                  }
+        
+        #DataProcessor.__init__(self, **kwargs)
+        super(AstraBackProjector3D, self).__init__(**kwargs)
+        
+        self.set_ImageGeometry(volume_geometry)
+        self.set_AcquisitionGeometry(sinogram_geometry)
+        
+                # Set up ASTRA Volume and projection geometry, not to be stored in self
+        vol_geom, proj_geom = convert_geometry_to_astra(self.volume_geometry,
+                                                        self.sinogram_geometry)
+        
+        # Also store ASTRA geometries
+        self.vol_geom = vol_geom
+        self.proj_geom = proj_geom
+    
+    def check_input(self, dataset):
+        if dataset.number_of_dimensions == 3:
+               return True
+        else:
+            raise ValueError("Expected input dimensions is 3, got {0}"\
+                             .format(dataset.number_of_dimensions))
+        
+    def set_ImageGeometry(self, volume_geometry):
+        self.volume_geometry = volume_geometry
+        
+    def set_AcquisitionGeometry(self, sinogram_geometry):
+        self.sinogram_geometry = sinogram_geometry
+    
+    def process(self):
+        DATA = self.get_input()
+        IM = ImageData(geometry=self.volume_geometry,
+                       dimension_labels=self.output_axes_order)
+        rec_id, IM.array = astra.create_backprojection3d_gpu(DATA.as_array(),
+                            self.proj_geom,
+                            self.vol_geom)
+        astra.data3d.delete(rec_id)
+        
+        # Scaling of 3D ASTRA backprojector, works both parallel and cone.
+        scaling = 1/self.volume_geometry.voxel_size_x**2  
+        return scaling*IM
diff --git a/Wrappers/Python/ccpi/astra/utils.py b/Wrappers/Python/ccpi/astra/utils.py
index 9f8fe46..2b19305 100755
--- a/Wrappers/Python/ccpi/astra/utils.py
+++ b/Wrappers/Python/ccpi/astra/utils.py
@@ -1,4 +1,5 @@
 import astra
+import numpy as np
 
 def convert_geometry_to_astra(volume_geometry, sinogram_geometry):
     # Set up ASTRA Volume and projection geometry, not stored
@@ -20,8 +21,8 @@ def convert_geometry_to_astra(volume_geometry, sinogram_geometry):
                                                sinogram_geometry.pixel_size_h,
                                                sinogram_geometry.pixel_num_h,
                                                sinogram_geometry.angles,
-                                               sinogram_geometry.dist_source_center,
-                                               sinogram_geometry.dist_center_detector)
+                                               np.abs(sinogram_geometry.dist_source_center),
+                                               np.abs(sinogram_geometry.dist_center_detector))
         else:
             NotImplemented
             
@@ -50,8 +51,8 @@ def convert_geometry_to_astra(volume_geometry, sinogram_geometry):
                                                sinogram_geometry.pixel_num_v,
                                                sinogram_geometry.pixel_num_h,
                                                sinogram_geometry.angles,
-                                               sinogram_geometry.dist_source_center,
-                                               sinogram_geometry.dist_center_detector)
+                                               np.abs(sinogram_geometry.dist_source_center),
+                                               np.abs(sinogram_geometry.dist_center_detector))
         else:
             NotImplemented