[core][test][dimension] Add tests for find_shift_along_dimension and...

[core][test][dimension] Add tests for find_shift_along_dimension and apply_shift_along_dimension methods

darfix/core/test/test_dimension.py

@@ -229,6 +229,46 @@ class TestDimension(unittest.TestCase):
         self.assertEqual(new_dataset.data.urls[0, 0, 0], dataset.data.urls[0, 0, 0])
         self.assertNotEqual(new_dataset.data.urls[0, 0, 1], dataset.data.urls[0, 0, 1])
 
+    def test_find_shift_along_dimension(self):
+        """ Tests the shift detection along a dimension"""
+
+        # In memory
+        self.in_memory_dataset.find_dimensions(POSITIONER_METADATA)
+        dataset = self.in_memory_dataset.reshape_data()
+        indices = numpy.arange(10)
+        shift = dataset.find_shift_along_dimension(dimension=[1], indices=indices)
+        self.assertEqual(shift.shape, (2, 2, 5))
+        shift = dataset.find_shift_along_dimension(dimension=[0], indices=indices)
+        self.assertEqual(shift.shape, (5, 2, 2))
+
+        # In disk
+        self.in_disk_dataset.find_dimensions(POSITIONER_METADATA)
+        dataset = self.in_disk_dataset.reshape_data()
+        indices = numpy.arange(10)
+        shift = dataset.find_shift_along_dimension(dimension=[1], indices=indices)
+        self.assertEqual(shift.shape, (2, 2, 5))
+        shift = dataset.find_shift_along_dimension(dimension=[0], indices=indices)
+        self.assertEqual(shift.shape, (5, 2, 2))
+
+    def test_apply_shift_along_dimension(self):
+        """ Tests the shift correction with dimensions and indices"""
+
+        # In memory
+        self.in_memory_dataset.find_dimensions(POSITIONER_METADATA)
+        dataset = self.in_memory_dataset.reshape_data()
+        shift = numpy.random.random((4, 2, 2))
+        new_dataset = dataset.apply_shift_along_dimension(shift=shift, dimension=[1], indices=[1, 2, 3, 4])
+        self.assertEqual(new_dataset.data.urls[0, 0, 0], dataset.data.urls[0, 0, 0])
+        self.assertNotEqual(new_dataset.data.urls[0, 0, 1], dataset.data.urls[0, 0, 1])
+        #  In disk
+        self.in_disk_dataset.find_dimensions(POSITIONER_METADATA)
+        dataset = self.in_disk_dataset.reshape_data()
+        shift = numpy.random.random((4, 2, 2))
+        new_dataset = dataset.apply_shift_along_dimension(shift=shift, dimension=[1], indices=[1, 2, 3, 4])
+
+        self.assertEqual(new_dataset.data.urls[0, 0, 0], dataset.data.urls[0, 0, 0])
+        self.assertNotEqual(new_dataset.data.urls[0, 0, 1], dataset.data.urls[0, 0, 1])
+
     def test_zsum(self):
         """ Tests the shift detection with dimensions and indices"""