ASTRA fp and bp as datasetprocessors used by operator (#39)

* Remove astra vol_geom and proj_geom from astra_op

* Moved astra fp and bp from op into processor

* Moved all astra from op into datasetproc

* Moved geometry conversion to new astra_utils, clean up

* two minor fixes

3 files changed, 216 insertions, 87 deletions

diff --git a/Wrappers/Python/ccpi/astra/astra_processors.py b/Wrappers/Python/ccpi/astra/astra_processors.py
new file mode 100644
index 0000000..91612b1
--- /dev/null
+++ b/Wrappers/Python/ccpi/astra/astra_processors.py
@@ -0,0 +1,130 @@
+from ccpi.framework import DataSetProcessor, DataSet, VolumeData, SinogramData
+from ccpi.astra.astra_utils import convert_geometry_to_astra
+import astra
+
+
+class AstraForwardProjector(DataSetProcessor):
+    '''AstraForwardProjector
+    
+    Forward project VolumeDataSet to SinogramDataSet using ASTRA proj_id.
+    
+    Input: VolumeDataSet
+    Parameter: proj_id
+    Output: SinogramDataSet
+    '''
+    
+    def __init__(self,
+                 volume_geometry=None,
+                 sinogram_geometry=None,
+                 proj_id=None,
+                 device='cpu'):
+        kwargs = {
+                  'volume_geometry'  : volume_geometry,
+                  'sinogram_geometry'  : sinogram_geometry,
+                  'proj_id'  : proj_id,
+                  'device'  : device
+                  }
+        
+        #DataSetProcessor.__init__(self, **kwargs)
+        super(AstraForwardProjector, self).__init__(**kwargs)
+        
+        self.setVolumeGeometry(volume_geometry)
+        self.setSinogramGeometry(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)
+        
+        # ASTRA projector, to be stored
+        if device == 'cpu':
+            # Note that 'line' is only for parallel (2D) and only one option
+            self.setProjector(astra.create_projector('line', proj_geom, vol_geom) )
+        elif device == 'gpu':
+            self.setProjector(astra.create_projector('cuda', proj_geom, vol_geom) )
+        else:
+            NotImplemented
+    
+    def checkInput(self, dataset):
+        if dataset.number_of_dimensions == 3 or dataset.number_of_dimensions == 2:
+               return True
+        else:
+            raise ValueError("Expected input dimensions is 2 or 3, got {0}"\
+                             .format(dataset.number_of_dimensions))
+    
+    def setProjector(self, proj_id):
+        self.proj_id = proj_id
+    
+    def setVolumeGeometry(self, volume_geometry):
+        self.volume_geometry = volume_geometry
+    
+    def setSinogramGeometry(self, sinogram_geometry):
+        self.sinogram_geometry = sinogram_geometry
+    
+    def process(self):
+        IM = self.getInput()
+        sinogram_id, DATA = astra.create_sino(IM.as_array(), self.proj_id)
+        astra.data2d.delete(sinogram_id)
+        return SinogramData(DATA,geometry=self.sinogram_geometry)
+
+class AstraBackProjector(DataSetProcessor):
+    '''AstraBackProjector
+    
+    Back project SinogramDataSet to VolumeDataSet using ASTRA proj_id.
+    
+    Input: SinogramDataSet
+    Parameter: proj_id
+    Output: VolumeDataSet 
+    '''
+    
+    def __init__(self,
+                 volume_geometry=None,
+                 sinogram_geometry=None,
+                 proj_id=None,
+                 device='cpu'):
+        kwargs = {
+                  'volume_geometry'  : volume_geometry,
+                  'sinogram_geometry'  : sinogram_geometry,
+                  'proj_id'  : proj_id,
+                  'device'  : device
+                  }
+        
+        #DataSetProcessor.__init__(self, **kwargs)
+        super(AstraBackProjector, self).__init__(**kwargs)
+        
+        self.setVolumeGeometry(volume_geometry)
+        self.setSinogramGeometry(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)
+        
+        # ASTRA projector, to be stored
+        if device == 'cpu':
+            # Note that 'line' is only for parallel (2D) and only one option
+            self.setProjector(astra.create_projector('line', proj_geom, vol_geom) )
+        elif device == 'gpu':
+            self.setProjector(astra.create_projector('cuda', proj_geom, vol_geom) )
+        else:
+            NotImplemented
+    
+    def checkInput(self, dataset):
+        if dataset.number_of_dimensions == 3 or dataset.number_of_dimensions == 2:
+               return True
+        else:
+            raise ValueError("Expected input dimensions is 2 or 3, got {0}"\
+                             .format(dataset.number_of_dimensions))
+    
+    def setProjector(self, proj_id):
+        self.proj_id = proj_id
+        
+    def setVolumeGeometry(self, volume_geometry):
+        self.volume_geometry = volume_geometry
+        
+    def setSinogramGeometry(self, sinogram_geometry):
+        self.sinogram_geometry = sinogram_geometry
+    
+    def process(self):
+        DATA = self.getInput()
+        rec_id, IM = astra.create_backprojection(DATA.as_array(), self.proj_id)
+        astra.data2d.delete(rec_id)
+        return VolumeData(IM,geometry=self.volume_geometry)
\ No newline at end of file
diff --git a/Wrappers/Python/ccpi/astra/astra_utils.py b/Wrappers/Python/ccpi/astra/astra_utils.py
new file mode 100644
index 0000000..8b5043a
--- /dev/null
+++ b/Wrappers/Python/ccpi/astra/astra_utils.py
@@ -0,0 +1,61 @@
+import astra
+
+def convert_geometry_to_astra(volume_geometry, sinogram_geometry):
+    # Set up ASTRA Volume and projection geometry, not stored
+    if sinogram_geometry.dimension == '2D':
+        vol_geom = astra.create_vol_geom(volume_geometry.voxel_num_x, 
+                                         volume_geometry.voxel_num_y, 
+                                         volume_geometry.getMinX(), 
+                                         volume_geometry.getMaxX(), 
+                                         volume_geometry.getMinY(), 
+                                         volume_geometry.getMaxY())
+        
+        if sinogram_geometry.geom_type == 'parallel':
+            proj_geom = astra.create_proj_geom('parallel',
+                                               sinogram_geometry.pixel_size_h,
+                                               sinogram_geometry.pixel_num_h,
+                                               sinogram_geometry.angles)
+        elif sinogram_geometry.geom_type == 'cone':
+            proj_geom = astra.create_proj_geom('fanflat',
+                                               sinogram_geometry.pixel_size_h,
+                                               sinogram_geometry.pixel_num_h,
+                                               sinogram_geometry.angles,
+                                               sinogram_geometry.dist_source_center,
+                                               sinogram_geometry.dist_center_detector)
+        else:
+            NotImplemented
+            
+    elif sinogram_geometry.dimension == '3D':
+        vol_geom = astra.create_vol_geom(volume_geometry.voxel_num_x, 
+                                         volume_geometry.voxel_num_y, 
+                                         volume_geometry.voxel_num_z, 
+                                         volume_geometry.getMinX(), 
+                                         volume_geometry.getMaxX(), 
+                                         volume_geometry.getMinY(), 
+                                         volume_geometry.getMaxY(), 
+                                         volume_geometry.getMinZ(), 
+                                         volume_geometry.getMaxZ())
+        
+        if sinogram_geometry.proj_geom == 'parallel':
+            proj_geom = astra.create_proj_geom('parallel3d',
+                                               sinogram_geometry.pixel_size_h,
+                                               sinogram_geometry.pixel_size_v,
+                                               sinogram_geometry.pixel_num_v,
+                                               sinogram_geometry.pixel_num_h,
+                                               sinogram_geometry.angles)
+        elif sinogram_geometry.geom_type == 'cone':
+            proj_geom = astra.create_proj_geom('cone',
+                                               sinogram_geometry.pixel_size_h,
+                                               sinogram_geometry.pixel_size_v,
+                                               sinogram_geometry.pixel_num_v,
+                                               sinogram_geometry.pixel_num_h,
+                                               sinogram_geometry.angles,
+                                               sinogram_geometry.dist_source_center,
+                                               sinogram_geometry.dist_center_detector)
+        else:
+            NotImplemented
+            
+    else:
+        NotImplemented
+    
+    return vol_geom, proj_geom
\ No newline at end of file
diff --git a/Wrappers/Python/ccpi/reconstruction/astra_ops.py b/Wrappers/Python/ccpi/reconstruction/astra_ops.py
index 452c86a..6af7af7 100644
--- a/Wrappers/Python/ccpi/reconstruction/astra_ops.py
+++ b/Wrappers/Python/ccpi/reconstruction/astra_ops.py
@@ -16,118 +16,56 @@
 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from ccpi.reconstruction.ops import Operator
-import astra
 import numpy
 from ccpi.framework import SinogramData, VolumeData
 from ccpi.reconstruction.ops import PowerMethodNonsquare
+from ccpi.astra.astra_processors import AstraForwardProjector, AstraBackProjector
 
 class AstraProjectorSimple(Operator):
     """ASTRA projector modified to use DataSet and geometry."""
     def __init__(self, geomv, geomp, device):
         super(AstraProjectorSimple, self).__init__()
         
-        # Store our volume and sinogram geometries. Redundant with also 
-        # storing in ASTRA format below but needed to assign to 
-        # SinogramData in "direct" method and VolumeData in "adjoint" method
+        # Store volume and sinogram geometries.
         self.sinogram_geometry = geomp
         self.volume_geometry = geomv
         
-        # ASTRA Volume geometry
-        if geomp.dimension == '2D':
-            self.vol_geom = astra.create_vol_geom(geomv.voxel_num_x, 
-                                                  geomv.voxel_num_y, 
-                                                  geomv.getMinX(), 
-                                                  geomv.getMaxX(), 
-                                                  geomv.getMinY(), 
-                                                  geomv.getMaxY())
-        elif geomp.dimension == '3D':
-            self.vol_geom = astra.create_vol_geom(geomv.voxel_num_x, 
-                                                  geomv.voxel_num_y, 
-                                                  geomv.voxel_num_z, 
-                                                  geomv.getMinX(), 
-                                                  geomv.getMaxX(), 
-                                                  geomv.getMinY(), 
-                                                  geomv.getMaxY(), 
-                                                  geomv.getMinZ(), 
-                                                  geomv.getMaxZ())
-        else:
-            NotImplemented
-            
-        
-        # ASTRA Projections geometry
-        if geomp.dimension == '2D':
-            if geomp.geom_type == 'parallel':
-                self.proj_geom = astra.create_proj_geom('parallel',
-                                                        geomp.pixel_size_h,
-                                                        geomp.pixel_num_h,
-                                                        geomp.angles)
-            elif geomp.geom_type == 'cone':
-                self.proj_geom = astra.create_proj_geom('fanflat',
-                                                        geomp.pixel_size_h,
-                                                        geomp.pixel_num_h,
-                                                        geomp.angles,
-                                                        geomp.dist_source_center,
-                                                        geomp.dist_center_detector)
-            else:
-                NotImplemented
-        elif geomp.dimension == '3D':
-            if geomp.proj_geom == 'parallel':
-                self.proj_geom = astra.create_proj_geom('parallel3d',
-                                                        geomp.pixel_size_h,
-                                                        geomp.pixel_size_v,
-                                                        geomp.pixel_num_v,
-                                                        geomp.pixel_num_h,
-                                                        geomp.angles)
-            elif geomp.geom_type == 'cone':
-                self.proj_geom = astra.create_proj_geom('cone',
-                                                        geomp.pixel_size_h,
-                                                        geomp.pixel_size_v,
-                                                        geomp.pixel_num_v,
-                                                        geomp.pixel_num_h,
-                                                        geomp.angles,
-                                                        geomp.dist_source_center,
-                                                        geomp.dist_center_detector)
-            else:
-                NotImplemented
-        else:
-            NotImplemented
-        
-        # ASTRA projector
-        if device == 'cpu':
-            # Note that 'line' is only for parallel (2D) and only one option
-            self.proj_id = astra.create_projector('line', self.proj_geom, 
-                                                  self.vol_geom) # for CPU
-        elif device == 'gpu':
-            self.proj_id = astra.create_projector('cuda', self.proj_geom, 
-                                                  self.vol_geom) # for GPU
-        else:
-            NotImplemented
+        self.fp = AstraForwardProjector(volume_geometry=geomv,
+                                        sinogram_geometry=geomp,
+                                        proj_id=None,
+                                        device=device)
         
+        self.bp = AstraBackProjector(volume_geometry=geomv,
+                                        sinogram_geometry=geomp,
+                                        proj_id=None,
+                                        device=device)
+                
+        # Initialise empty for singular value.
         self.s1 = None
     
     def direct(self, IM):
-        
-        sinogram_id, DATA = astra.create_sino(IM.as_array(), self.proj_id)
-        astra.data2d.delete(sinogram_id)
-        return SinogramData(DATA,geometry=self.sinogram_geometry)
+        self.fp.setInput(IM)
+        out = self.fp.getOutput()
+        return out
     
     def adjoint(self, DATA):
-        rec_id, IM = astra.create_backprojection(DATA.as_array(), self.proj_id)
-        astra.data2d.delete(rec_id)
-        return VolumeData(IM,geometry=self.volume_geometry)
+        self.bp.setInput(DATA)
+        out = self.bp.getOutput()
+        return out
     
-    def delete(self):
-        astra.data2d.delete(self.proj_id)
+    #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):
-        return ( (self.proj_geom['ProjectionAngles'].size, \
-                  self.proj_geom['DetectorCount']), \
-                 (self.vol_geom['GridColCount'], \
-                  self.vol_geom['GridRowCount']) )
+        # Only implemented for 2D
+        return ( (self.sinogram_geometry.angles.size, \
+                  self.sinogram_geometry.pixel_num_h), \
+                 (self.volume_geometry.voxel_num_x, \
+                  self.volume_geometry.voxel_num_y) )
     
 
 class AstraProjector(Operator):