defines shape for a subset slice (#410)

* defines shape for a subset slice

closes #408

* fix subset for AcquisitionData

* add future imports (#411)

* override the dimension parameter, set angles to None by default (#412)

* defines shape for a subset slice

* fix subset for AcquisitionData

* added initial unittest for subset

* added .gitignore

* adds unittest for subset, fix subset for AcquisitionData

* reverted change of behaviour of as_array

* removed empty line

2 files changed, 168 insertions, 14 deletions

diff --git a/Wrappers/Python/ccpi/framework/framework.py b/Wrappers/Python/ccpi/framework/framework.py
index 627e6d2..c30c436 100755
--- a/Wrappers/Python/ccpi/framework/framework.py
+++ b/Wrappers/Python/ccpi/framework/framework.py
@@ -110,6 +110,8 @@ class ImageGeometry(object):
             if order != [i for i in range(len(dim_labels))]:
                 # resort
                 self.shape = tuple([shape[i] for i in order])
+            else:
+                self.shape = tuple(order)
             self.dimension_labels = labels
                 
     def get_order_by_label(self, dimension_labels, default_dimension_labels):
@@ -316,12 +318,12 @@ class AcquisitionGeometry(object):
                 if not reduce(lambda x,y: (y in allowed_labels) and x, labels , True):
                     raise ValueError('Requested axis are not possible. Expected {},\ngot {}'.format(
                                     allowed_labels,labels))
-            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 != [i for i in range(len(dim_labels))]:
                 # resort
                 self.shape = tuple([shape[i] for i in order])
+            else:
+                self.shape = tuple(order)
             self.dimension_labels = labels
         
     def get_order_by_label(self, dimension_labels, default_dimension_labels):
@@ -447,7 +449,6 @@ class DataContainer(object):
             raise ValueError('Unknown dimension {0}. Should be one of'.format(dimension_label,
                              self.dimension_labels.values()))
                         
-
     def as_array(self, dimensions=None):
         '''Returns the DataContainer as Numpy Array
         
@@ -1199,20 +1200,24 @@ class AcquisitionData(DataContainer):
     def subset(self, dimensions=None, **kw):
         '''returns a subset of the AcquisitionData and regenerates the geometry'''
 
+        # # Check that this is actually a resorting
+        # if dimensions is not None and \
+        #     (len(dimensions) != len(self.shape) ):
+        #     raise ValueError('Please specify the slice on the axis/axes you want to cut away, or the same amount of axes for resorting')
+
+        # requested_labels = kw.get('dimension_labels', None)
+        # if requested_labels is not None:
+        #     allowed_labels = [AcquisitionGeometry.CHANNEL,
+        #                           AcquisitionGeometry.ANGLE,
+        #                           AcquisitionGeometry.VERTICAL,
+        #                           AcquisitionGeometry.HORIZONTAL]
+        #     if not reduce(lambda x,y: (y in allowed_labels) and x, requested_labels , True):
+        #         raise ValueError('Requested axis are not possible. Expected {},\ngot {}'.format(
+        #                         allowed_labels,requested_labels))
         # Check that this is actually a resorting
         if dimensions is not None and \
             (len(dimensions) != len(self.shape) ):
             raise ValueError('Please specify the slice on the axis/axes you want to cut away, or the same amount of axes for resorting')
-
-        requested_labels = kw.get('dimension_labels', None)
-        if requested_labels is not None:
-            allowed_labels = [AcquisitionGeometry.CHANNEL,
-                                  AcquisitionGeometry.ANGLE,
-                                  AcquisitionGeometry.VERTICAL,
-                                  AcquisitionGeometry.HORIZONTAL]
-            if not reduce(lambda x,y: (y in allowed_labels) and x, requested_labels , True):
-                raise ValueError('Requested axis are not possible. Expected {},\ngot {}'.format(
-                                allowed_labels,requested_labels))
         out = super(AcquisitionData, self).subset(dimensions, **kw)
         
         if out.number_of_dimensions > 1:
@@ -1226,6 +1231,14 @@ class AcquisitionData(DataContainer):
             pixel_size_v = 1
             dist_source_center = self.geometry.dist_source_center
             dist_center_detector = self.geometry.dist_center_detector
+
+            # update the angles if necessary
+            sliceme = kw.get(AcquisitionGeometry.ANGLE, None)
+            if sliceme is not None:
+                angles = numpy.asarray([ self.geometry.angles[sliceme] ] , numpy.float32)
+            else:
+                angles = self.geometry.angles.copy()
+            
             for key in out.dimension_labels.keys():
                 if out.dimension_labels[key] == AcquisitionGeometry.CHANNEL:
                     channels = self.geometry.channels
@@ -1243,7 +1256,7 @@ class AcquisitionData(DataContainer):
             
             out.geometry = AcquisitionGeometry(geom_type=self.geometry.geom_type, 
                                     dimension=dim,
-                                    angles=self.geometry.angles,
+                                    angles=angles,
                                     pixel_num_h=pixel_num_h,
                                     pixel_size_h = pixel_size_h,
                                     pixel_num_v = pixel_num_v,
diff --git a/Wrappers/Python/test/test_subset.py b/Wrappers/Python/test/test_subset.py
new file mode 100644
index 0000000..dd017fb
--- /dev/null
+++ b/Wrappers/Python/test/test_subset.py
@@ -0,0 +1,141 @@
+# -*- coding: utf-8 -*-
+#  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.
+import sys
+import unittest
+import numpy
+from ccpi.framework import DataContainer
+from ccpi.framework import ImageData
+from ccpi.framework import AcquisitionData
+from ccpi.framework import ImageGeometry
+from ccpi.framework import AcquisitionGeometry
+from timeit import default_timer as timer
+
+class TestSubset(unittest.TestCase):
+    def setUp(self):
+        self.ig = ImageGeometry(1,2,3,channels=4)
+        angles = numpy.asarray([90.,0.,-90.], dtype=numpy.float32)
+
+        self.ag = AcquisitionGeometry('cone', 'edo', pixel_num_h=20, pixel_num_v=2, angles=angles, 
+                         dist_source_center = 312.2, 
+                         dist_center_detector = 123.,
+                         channels=4 )
+
+    def test_ImageDataAllocate1a(self):
+        data = self.ig.allocate()
+        default_dimension_labels = [ImageGeometry.CHANNEL, ImageGeometry.VERTICAL,
+                ImageGeometry.HORIZONTAL_Y, ImageGeometry.HORIZONTAL_X]
+        self.assertTrue( default_dimension_labels == list(data.dimension_labels.values()) )
+    def test_ImageDataAllocate1b(self):
+        data = self.ig.allocate()
+        default_dimension_labels = [ImageGeometry.CHANNEL, ImageGeometry.VERTICAL,
+                ImageGeometry.HORIZONTAL_Y, ImageGeometry.HORIZONTAL_X]
+        self.assertTrue( data.shape == (4,3,2,1))
+        
+    def test_ImageDataAllocate2a(self):
+        non_default_dimension_labels = [ ImageGeometry.HORIZONTAL_X, ImageGeometry.VERTICAL,
+                ImageGeometry.HORIZONTAL_Y, ImageGeometry.CHANNEL]
+        data = self.ig.allocate(dimension_labels=non_default_dimension_labels)
+        self.assertTrue( non_default_dimension_labels == list(data.dimension_labels.values()) )
+        
+    def test_ImageDataAllocate2b(self):
+        non_default_dimension_labels = [ ImageGeometry.HORIZONTAL_X, ImageGeometry.VERTICAL,
+                ImageGeometry.HORIZONTAL_Y, ImageGeometry.CHANNEL]
+        data = self.ig.allocate(dimension_labels=non_default_dimension_labels)
+        self.assertTrue( data.shape == (1,3,2,4))
+
+    def test_AcquisitionDataAllocate1a(self):
+        data = self.ag.allocate()
+        default_dimension_labels = [AcquisitionGeometry.CHANNEL ,
+                 AcquisitionGeometry.ANGLE , AcquisitionGeometry.VERTICAL ,
+                 AcquisitionGeometry.HORIZONTAL]
+        self.assertTrue(  default_dimension_labels == list(data.dimension_labels.values()) )
+
+    def test_AcquisitionDataAllocate1b(self):
+        data = self.ag.allocate()
+        default_dimension_labels = [AcquisitionGeometry.CHANNEL ,
+                 AcquisitionGeometry.ANGLE , AcquisitionGeometry.VERTICAL ,
+                 AcquisitionGeometry.HORIZONTAL]
+
+        self.assertTrue( data.shape == (4,3,2,20))
+
+    def test_AcquisitionDataAllocate2a(self):
+        non_default_dimension_labels = [AcquisitionGeometry.CHANNEL, AcquisitionGeometry.HORIZONTAL,
+         AcquisitionGeometry.VERTICAL, AcquisitionGeometry.ANGLE]
+        data = self.ag.allocate(dimension_labels=non_default_dimension_labels)
+
+
+        self.assertTrue(  non_default_dimension_labels == list(data.dimension_labels.values()) )
+        
+    def test_AcquisitionDataAllocate2b(self):
+        non_default_dimension_labels = [AcquisitionGeometry.CHANNEL, AcquisitionGeometry.HORIZONTAL,
+         AcquisitionGeometry.VERTICAL, AcquisitionGeometry.ANGLE]
+        data = self.ag.allocate(dimension_labels=non_default_dimension_labels)
+        self.assertTrue( data.shape == (4,20,2,3))
+
+    def test_AcquisitionDataSubset1a(self):
+        non_default_dimension_labels = [AcquisitionGeometry.CHANNEL, AcquisitionGeometry.HORIZONTAL,
+         AcquisitionGeometry.VERTICAL, AcquisitionGeometry.ANGLE]
+        data = self.ag.allocate(dimension_labels=non_default_dimension_labels)
+        #self.assertTrue( data.shape == (4,20,2,3))
+        sub = data.subset(vertical = 0)
+        self.assertTrue( sub.shape == (4,20,3))
+    
+    def test_AcquisitionDataSubset1b(self):
+        non_default_dimension_labels = [AcquisitionGeometry.CHANNEL, AcquisitionGeometry.HORIZONTAL,
+         AcquisitionGeometry.VERTICAL, AcquisitionGeometry.ANGLE]
+        data = self.ag.allocate(dimension_labels=non_default_dimension_labels)
+        #self.assertTrue( data.shape == (4,20,2,3))
+        sub = data.subset(channel = 0)
+        self.assertTrue( sub.shape == (20,2,3))
+    def test_AcquisitionDataSubset1c(self):
+        non_default_dimension_labels = [AcquisitionGeometry.CHANNEL, AcquisitionGeometry.HORIZONTAL,
+         AcquisitionGeometry.VERTICAL, AcquisitionGeometry.ANGLE]
+        data = self.ag.allocate(dimension_labels=non_default_dimension_labels)
+        #self.assertTrue( data.shape == (4,20,2,3))
+        sub = data.subset(horizontal = 0)
+        self.assertTrue( sub.shape == (4,2,3))
+    def test_AcquisitionDataSubset1d(self):
+        non_default_dimension_labels = [AcquisitionGeometry.CHANNEL, AcquisitionGeometry.HORIZONTAL,
+         AcquisitionGeometry.VERTICAL, AcquisitionGeometry.ANGLE]
+        data = self.ag.allocate(dimension_labels=non_default_dimension_labels)
+        #self.assertTrue( data.shape == (4,20,2,3))
+        sliceme = 1
+        sub = data.subset(angle = sliceme)
+        #print (sub.shape  , sub.dimension_labels)
+        self.assertTrue( sub.shape == (4,20,2) )
+        self.assertTrue( sub.geometry.angles[0] == data.geometry.angles[sliceme])
+    def test_AcquisitionDataSubset1e(self):
+        non_default_dimension_labels = [AcquisitionGeometry.CHANNEL, AcquisitionGeometry.HORIZONTAL,
+         AcquisitionGeometry.VERTICAL, AcquisitionGeometry.ANGLE]
+        data = self.ag.allocate(dimension_labels=non_default_dimension_labels)
+        #self.assertTrue( data.shape == (4,20,2,3))
+        sliceme = 1
+        sub = data.subset(angle = sliceme)
+        self.assertTrue( sub.geometry.angles[0] == data.geometry.angles[sliceme])
+    def test_AcquisitionDataSubset1f(self):
+        
+        data = self.ag.allocate()
+        #self.assertTrue( data.shape == (4,20,2,3))
+        sliceme = 1
+        sub = data.subset(angle = sliceme)
+        self.assertTrue( sub.geometry.angles[0] == data.geometry.angles[sliceme])
+        
+        
+            
+
+        
\ No newline at end of file