Add new option to provide arbitrary order of frames (#526)

* Add new option to provide arbitrary order of frames

* Add logging

* remove debugging traces

ptypy/experiment/hdf5_loader.py

@@ -193,7 +193,7 @@ class Hdf5Loader(PtyScan):
     default =
     type = Param
     help = Parameters for the filtering of frames
-    doc = The shape of loaded data is assumed to hvae the same dimensionality as data.shape[:-2]
+    doc = The shape of loaded data is assumed to have the same dimensionality as data.shape[:-2]
 
     [framefilter.file]
     default = None
@@ -298,6 +298,18 @@ class Hdf5Loader(PtyScan):
     help = Switch for loading data from electron ptychography experiments.
     doc = If True, the energy provided in keV will be considered as electron energy
           and converted to electron wavelengths.
+
+    [frameorder]
+    default =
+    type = Param
+    help = Parameters for the re-ordering of frames
+    doc = The shape of loaded array of indices is matching the dimensionality of the loaded intensity
+
+    [frameorder.indices]
+    default = None
+    type = list, ndarray
+    help = This is the array or list with the re-ordered indices.
+
     """
 
     def __init__(self, pars=None, swmr=False, **kwargs):
@@ -596,6 +608,16 @@ def _prepare_center(self):
             log(3, "center is %s, auto_center: %s" % (self.info.center, self.info.auto_center))
             log(3, "The loader will not do any cropping.")
 
+    def _reorder_preview_indices(self):
+        if self.p.frameorder.indices is None:
+            return
+        order = np.array(self.p.frameorder.indices, dtype=int)
+        if (order.max() > self.preview_indices.shape[-1]):
+            log(3, "Given frameorder does not match dimensionality of data, keeping the original order")
+            return
+        log(3, "Reordering indices")
+        self.preview_indices = self.preview_indices.T[order].T
+
     def load_unmapped_raster_scan(self, indices):
         intensities = {}
         positions = {}
@@ -732,6 +754,7 @@ def compute_scan_mapping_and_trajectory(self, data_shape, positions_fast_shape,
                 self.preview_indices = np.array([indices[1][::skip,::skip].flatten(), indices[0][::skip,::skip].flatten()], dtype=int)
                 if self.framefilter is not None:
                     self.preview_indices = self.preview_indices[:,self.framefilter[indices[1][::skip,::skip], indices[0][::skip,::skip]].flatten()]
+                self._reorder_preview_indices()
                 self.num_frames = len(self.preview_indices[0])
 
             else:
@@ -748,6 +771,7 @@ def compute_scan_mapping_and_trajectory(self, data_shape, positions_fast_shape,
                 self.preview_indices = indices[::skip]
                 if self.framefilter is not None:
                     self.preview_indices = self.preview_indices[self.framefilter[indices][::skip]]
+                self._reorder_preview_indices()
                 self.num_frames = len(self.preview_indices)
 
         elif ((len(positions_fast_shape)>1) and (len(positions_slow_shape)>1)) and data_shape[0] == np.prod(positions_fast_shape) == np.prod(positions_slow_shape):
@@ -776,6 +800,7 @@ def compute_scan_mapping_and_trajectory(self, data_shape, positions_fast_shape,
             self.preview_indices = np.array([indices[1][::skip,::skip].flatten(), indices[0][::skip,::skip].flatten()])
             if self.framefilter:
                 log(3, "Framefilter not supported for this case")
+            self._reorder_preview_indices()
             self.num_frames = len(self.preview_indices[0])
             self._ismapped = False
             self._scantype = 'raster'
@@ -809,6 +834,7 @@ def compute_scan_mapping_and_trajectory(self, data_shape, positions_fast_shape,
                 self.preview_indices = np.array([indices[1][::skip,::skip].flatten(), indices[0][::skip,::skip].flatten()], dtype=int)
                 if self.framefilter:
                     log(3, "Framefilter not supported for this case")
+                self._reorder_preview_indices()
                 self.num_frames = len(self.preview_indices[0])
                 self._ismapped = True
                 self._scantype = 'raster'
@@ -840,6 +866,7 @@ def compute_scan_mapping_and_trajectory(self, data_shape, positions_fast_shape,
                 self.preview_indices = np.array([indices[1][::skip,::skip].flatten(), indices[0][::skip,::skip].flatten()], dtype=int)
                 if self.framefilter:
                     log(3, "Framefilter not supported for this case")
+                self._reorder_preview_indices()
                 self.num_frames = len(self.preview_indices[0])
                 self._ismapped = False
                 self._scantype = 'raster'

test/ptyscan_tests/hdf5_loader_test.py

@@ -325,6 +325,106 @@ def test_position_data_mapping_case_1_with_framefilter(self):
         np.testing.assert_equal(out_data.shape, ground_truth.shape, err_msg="The shapes don't match for the positions for case 1 with framefilter")
         np.testing.assert_array_equal(out_data, ground_truth, err_msg='There is something up with the positions for case 1 with framefilter')
 
+    def test_position_data_mapping_case_1_with_frameorder_1(self):
+        '''
+        axis_data.shape (A, B) for data.shape (A, B, frame_size_m, frame_size_n),
+        '''
+        A = 106
+        B = 101
+        frame_size_m = 5
+        frame_size_n = 5
+
+        positions_slow = np.arange(A)
+        positions_fast = np.arange(B)
+        fast, slow = np.meshgrid(positions_fast, positions_slow) # just pretend it's a simple grid
+        fast = fast[..., np.newaxis, np.newaxis]
+        slow = slow[..., np.newaxis, np.newaxis]
+        # now chuck them in the files
+        with h5.File(self.positions_file, 'w') as f:
+            f[self.positions_slow_key] = slow
+            f[self.positions_fast_key] = fast
+
+        # make up some data ...
+        data = np.arange(A*B*frame_size_m*frame_size_n).reshape(A, B, frame_size_m, frame_size_n)
+        with h5.File(self.intensity_file, 'w') as f:
+            f[self.intensity_key] = data
+
+        # create frameorder array of indices
+        frameorder = np.arange(A*B)
+        np.random.shuffle(frameorder)
+
+        data_params = u.Param()
+        data_params.auto_center = False
+        data_params.intensities = u.Param()
+        data_params.intensities.file = self.intensity_file
+        data_params.intensities.key = self.intensity_key
+
+        data_params.positions = u.Param()
+        data_params.positions.file = self.positions_file
+        data_params.positions.slow_key = self.positions_slow_key
+        data_params.positions.fast_key = self.positions_fast_key
+
+        data_params.frameorder = u.Param()
+        data_params.frameorder.indices = frameorder
+        output = PtyscanTestRunner(Hdf5Loader, data_params, auto_frames=A*B, cleanup=False)
+
+        with h5.File(output['output_file'],'r') as f:
+            out_data = f['chunks/0/data'][...].squeeze()
+        ground_truth = data.reshape((-1, frame_size_m, frame_size_n))[frameorder]
+        np.testing.assert_equal(out_data.shape, ground_truth.shape, err_msg="The shapes don't match for the positions for case 1 with different frameorder")
+        np.testing.assert_array_equal(out_data, ground_truth, err_msg='There is something up with the positions for case 1 with different frameorder')
+
+
+    def test_position_data_mapping_case_1_with_frameorder_2(self):
+        '''
+        axis_data.shape (A, B) for data.shape (A, B, frame_size_m, frame_size_n),
+        '''
+        A = 106
+        B = 101
+        frame_size_m = 5
+        frame_size_n = 5
+
+        positions_slow = np.arange(A)
+        positions_fast = np.arange(B)
+        fast, slow = np.meshgrid(positions_fast, positions_slow) # just pretend it's a simple grid
+        fast = fast[..., np.newaxis, np.newaxis]
+        slow = slow[..., np.newaxis, np.newaxis]
+        # now chuck them in the files
+        with h5.File(self.positions_file, 'w') as f:
+            f[self.positions_slow_key] = slow
+            f[self.positions_fast_key] = fast
+
+        # make up some data ...
+        data = np.arange(A*B*frame_size_m*frame_size_n).reshape(A, B, frame_size_m, frame_size_n)
+        with h5.File(self.intensity_file, 'w') as f:
+            f[self.intensity_key] = data
+
+        # create frameorder array of indices
+        frameorder = np.hstack([np.arange(A*B), np.random.randint(A*B, size=int(0.1*A*B))])
+        np.random.shuffle(frameorder)
+
+        data_params = u.Param()
+        data_params.auto_center = False
+        data_params.intensities = u.Param()
+        data_params.intensities.file = self.intensity_file
+        data_params.intensities.key = self.intensity_key
+
+        data_params.positions = u.Param()
+        data_params.positions.file = self.positions_file
+        data_params.positions.slow_key = self.positions_slow_key
+        data_params.positions.fast_key = self.positions_fast_key
+
+        data_params.frameorder = u.Param()
+        data_params.frameorder.indices = frameorder
+        output = PtyscanTestRunner(Hdf5Loader, data_params, auto_frames=len(frameorder), cleanup=False)
+
+        with h5.File(output['output_file'],'r') as f:
+            out_data = f['chunks/0/data'][...].squeeze()
+        ground_truth = data.reshape((-1, frame_size_m, frame_size_n))[frameorder]
+        np.testing.assert_equal(out_data.shape, ground_truth.shape, err_msg="The shapes don't match for the positions for case 1 with different frameorder")
+        np.testing.assert_array_equal(out_data, ground_truth, err_msg='There is something up with the positions for case 1 with different frameorder')
+
+
     def test_darkfield_applied_case_1(self):
         '''
         Applies the darkfield and assumes it is shaped like the data
@@ -600,6 +700,57 @@ def test_position_data_mapping_case_2_with_framefilter(self):
                                       err_msg='There is something up with the positions for case 2 with framefilter')
 
 
+    def test_position_data_mapping_case_2_with_frameorder(self):
+        '''
+        axis_data.shape (k,) for data.shape (k, frame_size_m, frame_size_n)
+        '''
+        k = 12
+        frame_size_m = 5
+        frame_size_n = 5
+
+        positions_slow = np.arange(k)
+        positions_fast = np.arange(k)
+
+        # now chuck them in the files
+        with h5.File(self.positions_file, 'w') as f:
+            f[self.positions_slow_key] = positions_slow
+            f[self.positions_fast_key] = positions_fast
+
+        # make up some data ...
+        data = np.arange(k*frame_size_m*frame_size_n).reshape(k, frame_size_m, frame_size_n)
+        with h5.File(self.intensity_file, 'w') as f:
+            f[self.intensity_key] = data
+
+        # create frameorder array of indices
+        frameorder = np.hstack([np.arange(k), np.random.randint(k, size=int(0.1*k))])
+        np.random.shuffle(frameorder)
+
+        data_params = u.Param()
+        data_params.auto_center = False
+        data_params.intensities = u.Param()
+        data_params.intensities.file = self.intensity_file
+        data_params.intensities.key = self.intensity_key
+
+        data_params.positions = u.Param()
+        data_params.positions.file = self.positions_file
+        data_params.positions.slow_key = self.positions_slow_key
+        data_params.positions.fast_key = self.positions_fast_key
+
+        data_params.frameorder = u.Param()
+        data_params.frameorder.indices = frameorder
+
+        output = PtyscanTestRunner(Hdf5Loader, data_params, auto_frames=len(frameorder), cleanup=False)
+
+        with h5.File(output['output_file'], 'r') as f:
+            out_data = f['chunks/0/data'][...].squeeze()
+        ground_truth = data.reshape((-1, frame_size_m, frame_size_n))[frameorder]
+
+        np.testing.assert_equal(ground_truth.shape, out_data.shape,
+                                err_msg="The shapes don't match for the positions for case 2 with different order of frames")
+        np.testing.assert_array_equal(ground_truth, out_data,
+                                      err_msg='There is something up with the positions for case 2 with different order of frames')
+
+
     def test_flatfield_applied_case_2(self):
         '''
         Applies the flatfield and assumes it is shaped like a single frame
@@ -865,6 +1016,58 @@ def test_position_data_mapping_case_3_with_skipping(self):
         np.testing.assert_array_equal(out_data, ground_truth,
                                       err_msg='There is something up with the positions for case 4 with skipping')
 
+    def test_position_data_mapping_case_3_with_frameorder(self):
+        '''
+        axis_data.shape (C, D) for data.shape (C*D, frame_size_m, frame_size_n) ,
+        '''
+        C = 10
+        D = 11
+        frame_size_m = 5
+        frame_size_n = 5
+
+        positions_slow = np.arange(C)
+        positions_fast = np.arange(D)
+        fast, slow = np.meshgrid(positions_fast, positions_slow) # just pretend it's a simple grid
+        # now chuck them in the files
+        with h5.File(self.positions_file, 'w') as f:
+            f[self.positions_slow_key] = slow
+            f[self.positions_fast_key] = fast
+
+        # make up some data ...
+        data = np.arange(C*D*frame_size_m*frame_size_n).reshape(C*D, frame_size_m, frame_size_n)
+        with h5.File(self.intensity_file, 'w') as f:
+            f[self.intensity_key] = data
+
+        # create frameorder array of indices
+        frameorder = np.hstack([np.arange(C*D), np.random.randint(C*D, size=int(0.1*C*D))])
+        np.random.shuffle(frameorder)
+
+        data_params = u.Param()
+        data_params.auto_center = False
+        data_params.intensities = u.Param()
+        data_params.intensities.file = self.intensity_file
+        data_params.intensities.key = self.intensity_key
+
+        data_params.positions = u.Param()
+        data_params.positions.file = self.positions_file
+        data_params.positions.slow_key = self.positions_slow_key
+        data_params.positions.fast_key = self.positions_fast_key
+
+        data_params.frameorder = u.Param()
+        data_params.frameorder.indices = frameorder
+
+        output = PtyscanTestRunner(Hdf5Loader, data_params, auto_frames=len(frameorder), cleanup=False)
+
+        with h5.File(output['output_file'], 'r') as f:
+            out_data = f['chunks/0/data'][...].squeeze()
+        ground_truth = data.reshape((-1, frame_size_m, frame_size_n))[frameorder]
+
+        np.testing.assert_equal(out_data.shape, ground_truth.shape,
+                                err_msg="The shapes don't match for the positions for case 4 with different order of frames")
+        np.testing.assert_array_equal(out_data, ground_truth,
+                                      err_msg='There is something up with the positions for case 4 with different order of frames')
+
+
     def test_position_data_mapping_case_4(self):
         '''
         axis_data.shape (C,) for data.shape (C, D, frame_size_m, frame_size_n) where D is the size of the other axis,
@@ -1008,6 +1211,57 @@ def test_position_data_mapping_case_4_with_skipping(self):
         np.testing.assert_array_equal(out_data, ground_truth,
                                       err_msg='There is something up with the positions for case 4 with skipping')
 
+
+    def test_position_data_mapping_case_4_with_frameorder(self):
+        '''
+        axis_data.shape (C,) for data.shape (C, D, frame_size_m, frame_size_n) where D is the size of the other axis,
+        '''
+        C = 4
+        D = 8
+        frame_size_m = 5
+        frame_size_n = 5
+
+        slow = np.arange(C)
+        fast = np.arange(D)
+        # now chuck them in the files
+        with h5.File(self.positions_file, 'w') as f:
+            f[self.positions_slow_key] = slow
+            f[self.positions_fast_key] = fast
+
+        # make up some data ...
+        data = np.arange(C*D*frame_size_m*frame_size_n).reshape(C, D, frame_size_m, frame_size_n)
+        with h5.File(self.intensity_file, 'w') as f:
+            f[self.intensity_key] = data
+
+        # create frameorder array of indices
+        frameorder = np.hstack([np.arange(C*D), np.random.randint(C*D, size=int(0.1*C*D))])
+        np.random.shuffle(frameorder)
+
+        data_params = u.Param()
+        data_params.auto_center = False
+        data_params.intensities = u.Param()
+        data_params.intensities.file = self.intensity_file
+        data_params.intensities.key = self.intensity_key
+
+        data_params.positions = u.Param()
+        data_params.positions.file = self.positions_file
+        data_params.positions.slow_key = self.positions_slow_key
+        data_params.positions.fast_key = self.positions_fast_key
+
+        data_params.frameorder = u.Param()
+        data_params.frameorder.indices = frameorder
+
+        output = PtyscanTestRunner(Hdf5Loader, data_params, auto_frames=len(frameorder), cleanup=False)
+
+        with h5.File(output['output_file'], 'r') as f:
+            out_data = f['chunks/0/data'][...].squeeze()
+        ground_truth = data.reshape((-1, frame_size_m, frame_size_n))[frameorder]
+        np.testing.assert_equal(out_data.shape, ground_truth.shape,
+                                err_msg="The shapes don't match for the positions for case 4 with different order of frames")
+        np.testing.assert_array_equal(out_data, ground_truth,
+                                      err_msg='There is something up with the positions for case 4 with different order of frames')
+
+
     def test_position_data_mapping_case_5(self):
         '''
         axis_data.shape (C,) for data.shape (C*D, frame_size_m, frame_size_n) where D is the size of the other axis.