Works with dark and flat fields with same number or one fewer dimensi…

…ons than dataset

Wrappers/Python/cil/io/NEXUSDataReader.py

@@ -299,15 +299,13 @@ def _load_with_image_id(self, image_key_id):
                     raise ValueError("Data with image key: ", image_key_id, "is not in the data. Data Path ",
                                     self.file_name)
                 else:
-                    print("Image keys: ", image_keys)
                     angle_index = self._geometry.dimension_labels.index('angle')
                     if angle_index == 0:
                         data = all_data[image_keys == image_key_id]
                     elif angle_index == 1:
                         data = all_data[:, image_keys == image_key_id]
                     elif angle_index == 2:
                         data = all_data[:,:,image_keys == image_key_id]
-
             else:
                 raise ValueError("Dark and Flat fields are not saved in the data. Data Path ",
                                     self.file_name)

Wrappers/Python/cil/io/NEXUSDataWriter.py

@@ -102,15 +102,18 @@ def write(self):
                 if self.flat_field is not None:
                     if len(self.flat_field.shape) == len(self.data.as_array().shape)-1:
                         flat_fields_len = 1
-                    else:
-                        flat_fields_len = self.flat_field.shape[angles_id]
+                        shape = list(self.flat_field.shape)
+                        shape.insert(angles_id,1)
+                        self.flat_field = self.flat_field.reshape(shape)
+                    flat_fields_len = self.flat_field.shape[angles_id]
                 else:
                     flat_fields_len = 0
                 if self.dark_field is not None:
                     if len(self.dark_field.shape) == len(self.data.as_array().shape)-1:
-                        dark_fields_len = 1
-                    else:
-                        dark_fields_len = self.dark_field.shape[angles_id]
+                        shape = list(self.dark_field.shape)
+                        shape.insert(angles_id,1)
+                        self.dark_field = self.dark_field.reshape(shape)
+                    dark_fields_len = self.dark_field.shape[angles_id]
                 else:
                     dark_fields_len = 0
 
@@ -131,6 +134,7 @@ def write(self):
                         for i in range(0, dark_fields_len):
                             dark_field = np.take(self.dark_field, i, axis=angles_id)
                             data_to_write = np.insert(data_to_write, 0, dark_field, axis=angles_id)
+
                 if self.flat_field is not None:
                     if flat_fields_len ==0:
                         data_to_write = np.insert(data_to_write, 0, self.flat_field, axis=angles_id)

Wrappers/Python/test/test_NexusReaderWriter.py

@@ -36,7 +36,8 @@ def setUp(self):
                                     .set_angles([0, 90, 180],-3.0, 'radian')\
                                     .set_panel(5, 0.2, origin='top-right')\
                                     .set_channels(6)\
-                                    .set_labels(['horizontal', 'angle'])
+                                    .set_labels(['angle', 'horizontal'])
+                                    #.set_labels(['horizontal', 'angle'])
 
         self.ad2d = self.ag2d.allocate('random_int')
 
@@ -114,14 +115,13 @@ def readAcquisitionDataAndTest(self):
 
         if self.dark_field_2d is not None:
             dark_field_2d = reader2d.load_dark()
-            numpy.testing.assert_array_equal(dark_field_2d, self.dark_field_2d, 'Dark Field Data is not correct')
+            numpy.testing.assert_array_equal(dark_field_2d, self.dark_field_2d.reshape(1,5), 'Dark Field Data is not correct')
 
         if self.flat_field_2d is not None:
             flat_field_2d = reader2d.load_flat()
-            numpy.testing.assert_array_equal(flat_field_2d, self.flat_field_2d, 'Flat Field Data is not correct')
+            numpy.testing.assert_array_equal(flat_field_2d, self.flat_field_2d.reshape(1,5), 'Flat Field Data is not correct')
 
         assert ag2d == self.ag2d
-
         reader3d = NEXUSDataReader()
         reader3d.set_up(file_name = os.path.join(self.data_dir, 'test_nexus_ad3d.nxs'))
         ad3d = reader3d.read()
@@ -150,8 +150,9 @@ def readAcquisitionDataAndTest(self):
         assert ag3d == self.ag3d
 
     def test_writeAcquisitionData_with_dark_and_flat_fields(self):
-        self.flat_field_2d = numpy.ones((5,1))
-        self.dark_field_2d = numpy.zeros((5,1))
+        print("write 2d")
+        self.flat_field_2d = numpy.ones(5)
+        self.dark_field_2d = numpy.zeros(5)
 
         writer = NEXUSDataWriter()
         writer.set_up(file_name=os.path.join(self.data_dir, 'test_nexus_ad2d.nxs'),
@@ -161,7 +162,7 @@ def test_writeAcquisitionData_with_dark_and_flat_fields(self):
         self.flat_field_3d = numpy.ones((1, 10, 5))
         self.dark_field_3d = numpy.zeros((1, 10, 5))
 
-
+        print("write 3d")
         writer = NEXUSDataWriter()
         writer.set_up(file_name=os.path.join(self.data_dir, 'test_nexus_ad3d.nxs'),
                       data=self.ad3d, flat_field=self.flat_field_3d, dark_field=self.dark_field_3d)