diff --git a/ptypy/experiment/hdf5_loader.py b/ptypy/experiment/hdf5_loader.py index 44344477f..ac54caa81 100644 --- a/ptypy/experiment/hdf5_loader.py +++ b/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' diff --git a/test/ptyscan_tests/hdf5_loader_test.py b/test/ptyscan_tests/hdf5_loader_test.py index cdda48bf0..57e1a3933 100644 --- a/test/ptyscan_tests/hdf5_loader_test.py +++ b/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.