diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py index bdf4e60..8b6f6d8 100644 --- a/src/cdtools/datasets/ptycho_2d_dataset.py +++ b/src/cdtools/datasets/ptycho_2d_dataset.py @@ -400,9 +400,40 @@ class Ptycho2DDataset(CDataset): divisor_override=1)[0,0] + def remove_translations_mask(self, mask_remove): + """Removes one or more translation positions, and their associated + properties, from the dataset using logical indexing. + + This takes a 1D mask (boolean torch tensor) with the length + self.translations.shape[0] (i.e., the number of individual + translated points). Patterns, translations, and intensities + associated with indices that are "True" will be removed. + + Parameters: + ---------- + mask_remove : 1D torch.tensor(dtype=torch.bool) + The boolean mask indicating which elements are to be removed from + the dataset. True indicates that the corresponding element will be + removed. + """ + + # Check that the mask is the right size + if mask_remove.shape != t.Size([self.translations.shape[0]]): + raise ValueError( + 'The mask must have the same length as the number of translations in the dataset.' + ) + + # Update patterns, translations, and intensities + self.patterns = self.patterns[~mask_remove] + self.translations = self.translations[~mask_remove] + + if hasattr(self, 'intensities') and self.intensities is not None: + self.intensities = self.intensities[~mask_remove] + + 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 @@ -420,7 +451,7 @@ class Ptycho2DDataset(CDataset): 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]) @@ -441,9 +472,5 @@ class Ptycho2DDataset(CDataset): '(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 + # Remove translations outside the ROI + self.remove_translations_mask(~inside_roi) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 9e72fb7..6897991 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -1,4 +1,4 @@ -from cdtools.datasets import * +from cdtools.datasets import CDataset, Ptycho2DDataset from cdtools.tools import data as cdtdata import numpy as np import torch as t @@ -340,7 +340,31 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis): if dataset.background is not None: assert np.allclose(np.array(dataset.background.shape) // factor, np.array(copied_dataset.background.shape)) - + + +def test_Ptycho2DDataset_remove_translations_mask(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 mask is not the same shape as the pattern + # length + with pytest.raises(ValueError) as excinfo: + copied_dataset.remove_translations_mask(mask_remove=t.zeros(10)) + assert ('The mask must have the same length') in str(excinfo.value) + + # Test 2: Remove the mask from the dataset + mask_success = t.zeros(len(copied_dataset.patterns)) + mask_success[1] = 1 + mask_success[10] = 1 + mask_success[-1] = 1 + mask_success = mask_success.bool() + copied_dataset.remove_translations_mask(mask_remove=mask_success) + + # test if the mask is removed and patterns length is correct + assert len(copied_dataset.patterns) == len(mask_success) - 3 + def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1): # Grab dataset