Merge pull request #13 from yoshikisd/crop

Add a function to shrink the range of translation positions in a dataset

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]

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,:])