Partitioner - to fix error if partitions have size 1 (#1499)

fix error if partitions have size 1 Signed-off-by: Margaret Duff <[email protected]> Co-authored-by: Gemma Fardell <[email protected]> Co-authored-by: Edoardo Pasca <[email protected]>
TomographicImaging · Sep 22, 2023 · a6ce31c · a6ce31c
1 parent a3c5823
commit a6ce31c
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 10 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@
   - Add norm for CompositionOperator
   - Refactor SIRT algorithm to make it more computationally and memory efficient
   - Optimisation in L2NormSquared
+  - Added support for partitioner, when partitions have size 1
   - Fix for show_geometry bug for 2D data
   - Added warmstart capability to proximal evaluation of the CIL TotalVariation function.
   - FBP split processing bug fix - now respects panel origin

diff --git a/Wrappers/Python/cil/framework/framework.py b/Wrappers/Python/cil/framework/framework.py
@@ -126,7 +126,6 @@ def _construct_BlockGeometry_from_indices(self, indices):
             ag = self.geometry.copy()
             ag.config.angles.angle_data = numpy.take(self.geometry.angles, mask, axis=0)
             ags.append(ag)
-
         return BlockGeometry(*ags)
 
     def partition(self, num_batches, mode, seed=None):
@@ -200,8 +199,11 @@ def _partition_deterministic(self, num_batches, stagger=False, indices=None):
 
         for i in range(num_batches):
             out[i].fill(
-                numpy.take(self.array, partition_indices[i], axis=axis)
+                numpy.squeeze(
+                    numpy.take(self.array, partition_indices[i], axis=axis)
+                )
             )
+
         return out
 
     def _partition_random_permutation(self, num_batches, seed=None):
@@ -2111,7 +2113,6 @@ def shape(self):
                      AcquisitionGeometry.ANGLE: self.config.angles.num_positions,
                      AcquisitionGeometry.VERTICAL: self.config.panel.num_pixels[1],        
                      AcquisitionGeometry.HORIZONTAL: self.config.panel.num_pixels[0]}
-
         shape = []
         for label in self.dimension_labels:
             shape.append(shape_dict[label])

diff --git a/Wrappers/Python/test/test_BlockDataContainer.py b/Wrappers/Python/test/test_BlockDataContainer.py
@@ -882,30 +882,61 @@ def test_partition(self):
         # split in num_batches
         num_batches = 4
 
+        #Testing sequential partitioning 
         data = self.data.partition(num_batches, 'sequential')
         idxs = self.data._partition_indices(num_batches, indices=self.data.geometry.num_projections, stagger=False)
-
-        # ig = data.get_ImageGeometry()
         assert len(data.containers) == num_batches
         self.assertDataIsTheSame(data, idxs)
 
+        #Testing staggered partitioning 
         data = self.data.partition(num_batches, Partitioner.STAGGERED)
         idxs = self.data._partition_indices(num_batches, indices=self.data.geometry.num_projections, stagger=True)
-
-        # ig = data.get_ImageGeometry()
         assert len(data.containers) == num_batches
         self.assertDataIsTheSame(data, idxs)
 
+    def test_partition_diff_num_batches(self):
+
+        #Check what happens when the number of batches is equal to the number of projection angles 
+        num_batches=9
+        data = self.data.partition(num_batches, 'sequential')
+        idxs = self.data._partition_indices(num_batches, indices=self.data.geometry.num_projections, stagger=False)
+        assert len(data.containers) == num_batches
+        self.assertDataIsTheSame(data, idxs, msg='Failed when num_batches=number of projections')
+
+        #Check what happens when the number of batches is one, the whole set of projection angles 
+        num_batches=1
+        data = self.data.partition(num_batches, 'sequential')
+        idxs = self.data._partition_indices(num_batches, indices=self.data.geometry.num_projections, stagger=False)
+        assert len(data.containers) == num_batches
+        self.assertDataIsTheSame(data, idxs, msg="Failed when num_batches=1")
+
+        #Check what happens when the number of batches is zero
+        num_batches=0
+        with self.assertRaises(ZeroDivisionError):
+            data = self.data.partition(num_batches, 'sequential')
+
+       #Check what happens when the number of batches is greater than the number of projection angles
+        num_batches=10
+        with self.assertRaises(ValueError):
+            data = self.data.partition(num_batches, 'sequential')
+
+
+
+
 
-    def assertDataIsTheSame(self, data, idxs):
+    def assertDataIsTheSame(self, data, idxs, msg=None):
         # let's check that the data is the same
         k = 0
         wrong = 0
         for i, el in enumerate(data):
-            for j in range(el.shape[0]):
+            if len(el.shape)>1:
+                j_range=el.shape[0]
+            else:
+                j_range=1
+            for j in range(j_range):
                 idx = idxs[i][j]
                 try:
-                    np.testing.assert_array_equal(el.as_array()[j], self.data.as_array()[idx])
+                    np.testing.assert_array_equal(el.as_array()[j], self.data.as_array()[idx], err_msg=msg)
                 except AssertionError:
                     wrong += 1
                 k += 1