diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py index 46bec83..bdf4e60 100644 --- a/src/cdtools/datasets/ptycho_2d_dataset.py +++ b/src/cdtools/datasets/ptycho_2d_dataset.py @@ -398,3 +398,52 @@ class Ptycho2DDataset(CDataset): self.background.unsqueeze(0).unsqueeze(0), factor, divisor_override=1)[0,0] + + + def crop_translations(self, roi): + """Shrinks the range of translation positions that are analyzed + + This deletes all diffraction patterns associated with x- and + y-translations that lie outside of a specified rectangular + region of interest. In essence, this operation crops the "relative + displacement map" (shown in self.inspect()) down to the region of + interest. + + Parameters: + ---------- + roi : tuple(float, float, float, float) + The translation-x and -y coordinates that define the rectangular + region of interest as (in units of meters) + (left, right, bottom, top). The definition of these bounds are + based on how an image is normally displayed with matplotlib's + imshow. The order in which these elements are defined in roi + do not matter as long as roi[:2] and roi[2:] correspond with + the x and y coordinates, respectively. + """ + + # Pull out the bounds of the ROI, ensuring that left < right and + # top < bottom + x_left, x_right = sorted(roi[:2]) + y_top, y_bottom = sorted(roi[2:]) + + # Create pointers to the x- and y-translation positions in + # self.translations + x = self.translations[:, 0] + y = self.translations[:, 1] + + # Go look for all translation values that lie inside of the roi + # and store their indices. + inside_roi = (x >= x_left) & (x <= x_right) & (y <= y_bottom) & (y >= y_top) + + # Throw a value error if inside_roi is empty + if not t.any(inside_roi): + raise ValueError('The roi does not contain any positions from the dataset ' + '(i.e., patterns and translations will be empty).' + ' Please redefine the bounds of the roi.') + + # Update patterns and translations + self.patterns = self.patterns[inside_roi] + self.translations = self.translations[inside_roi] + + if hasattr(self, 'intensities') and self.intensities is not None: + self.intensities = self.intensities[inside_roi] \ No newline at end of file diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 1b135a1..9e72fb7 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -5,7 +5,8 @@ import torch as t import h5py import datetime from copy import deepcopy - +import pytest +import itertools # # We start by testing the CDataset base class @@ -340,3 +341,72 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis): assert np.allclose(np.array(dataset.background.shape) // factor, np.array(copied_dataset.background.shape)) + +def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1): + # Grab dataset + cxi, expected = ptycho_cxi_1 + dataset = Ptycho2DDataset.from_cxi(cxi) + copied_dataset = deepcopy(dataset) + + # Test 1: Complain when the the bounds of an ROI are correctly defined, + # but it does not contain any sample positions inside of it. The + # translations in ptycho_cxi_1 ranges from 0m to -300m in both x and y + # (it looks like a line scan). We select an ROI at (0, -300) which should + # not contain any translation positions. + with pytest.raises(ValueError) as excinfo: + copied_dataset.crop_translations(roi=(0,-5,-295,-300)) + assert('(i.e., patterns and translations will be empty)') in str(excinfo.value) + + # Test 2: Draw an ROI that's centered in the middle of the x/y translation range + # and make sure that the first and last x/y elements in dataset.translate + # contain x_left, x_right, y_top, and y_bottom + + # Make tuples that will store the x and y positions. This will be used for + # making permutations of the x/y positions in the ROI later. + x_left, y_top = dataset.translations[10,:2] + x_right, y_bottom = dataset.translations[-11,:2] + + x_permutations = ((x_left, x_right), (x_right, x_left)) + y_permutations = ((y_top, y_bottom), (y_bottom, y_top)) + roi_permutations = tuple((x1, x2, y1, y2) for (x1, x2), (y1, y2) in + itertools.product(x_permutations, y_permutations)) + + # Get the dataset + copied_dataset.crop_translations(roi=roi_permutations[0]) + + # Execute the actual test + assert (copied_dataset.translations[0,0] in x_permutations[0]) and \ + (copied_dataset.translations[-1,0] in x_permutations[0]) + + assert (copied_dataset.translations[0,1] in y_permutations[0]) and \ + (copied_dataset.translations[-1,1] in y_permutations[0]) + + # Test 3: Check if the shape of dataset.patterns and dataset.translate is correct + # (designed to be 20 fewer rows here) + # In the future, this should include a check for dataset.intensities once + # an appropriate cxi file is set up for conftest. + expected_patterns_shape = np.concatenate([[dataset.patterns.shape[0] - 20], + dataset.patterns.shape[-2:]]) + + expected_translations_shape = np.concatenate([[dataset.translations.shape[0] - 20], + dataset.translations.shape[-1:]]) + + assert np.allclose(np.array(copied_dataset.patterns.shape), expected_patterns_shape) + + assert np.allclose(np.array(copied_dataset.translations.shape), expected_translations_shape) + + # Test 4: Make sure that we always get the same result no matter what order we + # define the left/right and bottom/top values in roi, provided that roi[:2] + # and roi[2:] correspond with the x and y coordinates, respectively. + + # Check each permutation + for roi in roi_permutations: + # Copy the dataset again; dataset will be modified each time crop_translation + # is successfully executed. + copied_dataset = deepcopy(dataset) + copied_dataset.crop_translations(roi=roi) + + # Check if the contents of dataset.patterns and dataset.translate is correct + assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10,:]) + + assert t.allclose(copied_dataset.translations, dataset.translations[10:-10,:])