MC FP and BP working

4 files changed, 267 insertions, 25 deletions

diff --git a/Wrappers/Python/ccpi/astra/astra_processors.py b/Wrappers/Python/ccpi/astra/astra_processors.py
index 91612b1..650e11b 100644
--- a/Wrappers/Python/ccpi/astra/astra_processors.py
+++ b/Wrappers/Python/ccpi/astra/astra_processors.py
@@ -1,6 +1,7 @@
 from ccpi.framework import DataSetProcessor, DataSet, VolumeData, SinogramData
 from ccpi.astra.astra_utils import convert_geometry_to_astra
 import astra
+import numpy
 
 
 class AstraForwardProjector(DataSetProcessor):
@@ -127,4 +128,146 @@ class AstraBackProjector(DataSetProcessor):
         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)
+
+class AstraForwardProjectorMC(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(AstraForwardProjectorMC, 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:
+            return True
+            #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()
+        DATA = numpy.zeros((self.sinogram_geometry.angles.size,
+                                self.sinogram_geometry.pixel_num_h,
+                                1,
+                                self.sinogram_geometry.channels),
+                                'float32')
+        for k in range(self.sinogram_geometry.channels):
+            sinogram_id, DATA[:,:,0,k] = astra.create_sino(IM.as_array()[:,:,0,k], 
+                                                           self.proj_id)
+            astra.data2d.delete(sinogram_id)
+        return SinogramData(DATA,geometry=self.sinogram_geometry)
+
+class AstraBackProjectorMC(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(AstraBackProjectorMC, 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:
+            return True
+            #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()
+        IM = numpy.zeros((self.volume_geometry.voxel_num_x,
+                                self.volume_geometry.voxel_num_y,
+                                1,
+                                self.volume_geometry.channels),
+                                'float32')
+        for k in range(self.volume_geometry.channels):
+            rec_id, IM[:,:,0,k] = astra.create_backprojection(DATA.as_array()[:,:,0,k], 
+                                  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/framework.py b/Wrappers/Python/ccpi/framework.py
index 5ac17ee..5ee8279 100644
--- a/Wrappers/Python/ccpi/framework.py
+++ b/Wrappers/Python/ccpi/framework.py
@@ -428,20 +428,20 @@ class VolumeData(DataSet):
         
         if type(array) == DataSet:
             # if the array is a DataSet get the info from there
-            if not ( array.number_of_dimensions == 2 or \
-                     array.number_of_dimensions == 3 ):
-                raise ValueError('Number of dimensions are not 2 or 3: {0}'\
-                                 .format(array.number_of_dimensions))
+            #if not ( array.number_of_dimensions == 2 or \
+            #         array.number_of_dimensions == 3 ):
+            #    raise ValueError('Number of dimensions are not 2 or 3: {0}'\
+            #                     .format(array.number_of_dimensions))
             
             #DataSet.__init__(self, array.as_array(), deep_copy,
             #                 array.dimension_labels, **kwargs)
             super(VolumeData, self).__init__(array.as_array(), deep_copy,
                              array.dimension_labels, **kwargs)
         elif type(array) == numpy.ndarray:
-            if not ( array.ndim == 3 or array.ndim == 2 ):
-                raise ValueError(
-                        'Number of dimensions are not 3 or 2 : {0}'\
-                        .format(array.ndim))
+            #if not ( array.ndim == 3 or array.ndim == 2 ):
+            #    raise ValueError(
+            #            'Number of dimensions are not 3 or 2 : {0}'\
+            #            .format(array.ndim))
                 
             if dimension_labels is None:
                 if array.ndim == 3:
@@ -476,17 +476,17 @@ class SinogramData(DataSet):
         
         if type(array) == DataSet:
             # if the array is a DataSet get the info from there
-            if not ( array.number_of_dimensions == 2 or \
-                     array.number_of_dimensions == 3 ):
-                raise ValueError('Number of dimensions are not 2 or 3: {0}'\
-                                 .format(array.number_of_dimensions))
+            #if not ( array.number_of_dimensions == 2 or \
+            #         array.number_of_dimensions == 3 ):
+            #    raise ValueError('Number of dimensions are not 2 or 3: {0}'\
+            #                     .format(array.number_of_dimensions))
             
             DataSet.__init__(self, array.as_array(), deep_copy,
                              array.dimension_labels, **kwargs)
         elif type(array) == numpy.ndarray:
-            if not ( array.ndim == 3 or array.ndim == 2 ):
-                raise ValueError('Number of dimensions are != 3: {0}'\
-                                 .format(array.ndim))
+            #if not ( array.ndim == 3 or array.ndim == 2 ):
+            #    raise ValueError('Number of dimensions are != 3: {0}'\
+            #                     .format(array.ndim))
             if dimension_labels is None:
                 if array.ndim == 3:
                     dimension_labels = ['angle' , 
diff --git a/Wrappers/Python/ccpi/reconstruction/astra_ops.py b/Wrappers/Python/ccpi/reconstruction/astra_ops.py
index 6af7af7..cf138b4 100644
--- a/Wrappers/Python/ccpi/reconstruction/astra_ops.py
+++ b/Wrappers/Python/ccpi/reconstruction/astra_ops.py
@@ -19,7 +19,7 @@ from ccpi.reconstruction.ops import Operator
 import numpy
 from ccpi.framework import SinogramData, VolumeData
 from ccpi.reconstruction.ops import PowerMethodNonsquare
-from ccpi.astra.astra_processors import AstraForwardProjector, AstraBackProjector
+from ccpi.astra.astra_processors import *
 
 class AstraProjectorSimple(Operator):
     """ASTRA projector modified to use DataSet and geometry."""
@@ -66,6 +66,52 @@ class AstraProjectorSimple(Operator):
                   self.sinogram_geometry.pixel_num_h), \
                  (self.volume_geometry.voxel_num_x, \
                   self.volume_geometry.voxel_num_y) )
+
+class AstraProjectorMC(Operator):
+    """ASTRA Multichannel projector"""
+    def __init__(self, geomv, geomp, device):
+        super(AstraProjectorMC, self).__init__()
+        
+        # Store volume and sinogram geometries.
+        self.sinogram_geometry = geomp
+        self.volume_geometry = geomv
+        
+        self.fp = AstraForwardProjectorMC(volume_geometry=geomv,
+                                        sinogram_geometry=geomp,
+                                        proj_id=None,
+                                        device=device)
+        
+        self.bp = AstraBackProjectorMC(volume_geometry=geomv,
+                                        sinogram_geometry=geomp,
+                                        proj_id=None,
+                                        device=device)
+                
+        # Initialise empty for singular value.
+        self.s1 = None
+    
+    def direct(self, IM):
+        self.fp.setInput(IM)
+        out = self.fp.getOutput()
+        return out
+    
+    def adjoint(self, DATA):
+        self.bp.setInput(DATA)
+        out = self.bp.getOutput()
+        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.volume_geometry.voxel_num_x, \
+                  self.volume_geometry.voxel_num_y) )
     
 
 class AstraProjector(Operator):
diff --git a/Wrappers/Python/test/simple_mc_demo.py b/Wrappers/Python/test/simple_mc_demo.py
index ca6a89e..1888740 100644
--- a/Wrappers/Python/test/simple_mc_demo.py
+++ b/Wrappers/Python/test/simple_mc_demo.py
@@ -21,23 +21,76 @@ N = 128
 
 numchannels = 3
 
-x = np.zeros((N,N,numchannels))
+x = np.zeros((N,N,1,numchannels))
 
-x[round(N/4):round(3*N/4),round(N/4):round(3*N/4),0] = 1.0
-x[round(N/8):round(7*N/8),round(3*N/8):round(5*N/8),0] = 2.0
+x[round(N/4):round(3*N/4),round(N/4):round(3*N/4),:,0] = 1.0
+x[round(N/8):round(7*N/8),round(3*N/8):round(5*N/8),:,0] = 2.0
 
-x[round(N/4):round(3*N/4),round(N/4):round(3*N/4),1] = 0.7
-x[round(N/8):round(7*N/8),round(3*N/8):round(5*N/8),1] = 1.2
+x[round(N/4):round(3*N/4),round(N/4):round(3*N/4),:,1] = 0.7
+x[round(N/8):round(7*N/8),round(3*N/8):round(5*N/8),:,1] = 1.2
 
-x[round(N/4):round(3*N/4),round(N/4):round(3*N/4),2] = 1.5
-x[round(N/8):round(7*N/8),round(3*N/8):round(5*N/8),2] = 2.2
+x[round(N/4):round(3*N/4),round(N/4):round(3*N/4),:,2] = 1.5
+x[round(N/8):round(7*N/8),round(3*N/8):round(5*N/8),:,2] = 2.2
 
 f, axarr = plt.subplots(1,numchannels)
 for k in numpy.arange(3):
-    axarr[k].imshow(x[:,:,k],vmin=0,vmax=2.5)
+    axarr[k].imshow(x[:,:,0,k],vmin=0,vmax=2.5)
 plt.show()
 
 vg = VolumeGeometry(N,N,None, 1,1,None,channels=numchannels)
 
 
-Phantom = VolumeData(x,geometry=vg)
\ No newline at end of file
+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,
+                          channels=numchannels)
+elif test_case==2:
+    pg = SinogramGeometry('cone',
+                          '2D',
+                          angles,
+                          det_num,
+                          det_w,
+                          dist_source_center=SourceOrig, 
+                          dist_center_detector=OrigDetec,
+                          channels=numchannels)
+
+# ASTRA operator using volume and sinogram geometries
+Aop = AstraProjectorMC(vg, pg, 'gpu')
+
+
+# Try forward and backprojection
+b = Aop.direct(Phantom)
+
+fb, axarrb = plt.subplots(1,numchannels)
+for k in numpy.arange(3):
+    axarrb[k].imshow(b.as_array()[:,:,0,k],vmin=0,vmax=250)
+plt.show()
+
+out2 = Aop.adjoint(b)
+
+fo, axarro = plt.subplots(1,numchannels)
+for k in range(3):
+    axarro[k].imshow(out2.as_array()[:,:,0,k],vmin=0,vmax=3500)
+plt.show()
+