diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py index e270714..46bec83 100644 --- a/src/cdtools/datasets/ptycho_2d_dataset.py +++ b/src/cdtools/datasets/ptycho_2d_dataset.py @@ -361,3 +361,40 @@ class Ptycho2DDataset(CDataset): self.background = t.nn.functional.pad(self.background, to_pad) + def downsample(self, factor=2): + """Downsamples all diffraction patterns by the specified factor + + This is an easy way to shrink the amount of data you need to work with + if the speckle size is much larger than the detector pixel size. + + The downsampling factor must be an integer. The size of the output + patterns are reduced by the specified factor, with each output pixel + equal to the sum of a x region of pixels in the + input pattern. This summation is done by pytorch.functional.avg_pool2d. + + Any mask and background data which is stored with the dataset is + downsampled with the data. The background is downsampled using the same + method as the data. The mask is expanded so that any output pixel + containing a masked pixel will be masked. + + Parameters + ---------- + factor : int + Default 2, the factor to downsample by + + """ + self.patterns = t.nn.functional.avg_pool2d( + self.patterns.unsqueeze(0), factor, divisor_override=1)[0] + self.mask = t.logical_not(t.nn.functional.max_pool2d( + (1-self.mask.to(dtype=t.uint8)).unsqueeze(0).unsqueeze(0), + factor + )[0,0].to(dtype=t.bool)) + + self.detector_geometry['basis'] = \ + self.detector_geometry['basis'] * factor + + if self.background is not None: + self.background = t.nn.functional.avg_pool2d( + self.background.unsqueeze(0).unsqueeze(0), + factor, + divisor_override=1)[0,0] diff --git a/tests/test_datasets.py b/tests/test_datasets.py index d1b5f82..1b135a1 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -4,6 +4,7 @@ import numpy as np import torch as t import h5py import datetime +from copy import deepcopy # @@ -276,3 +277,66 @@ def test_Ptycho2DDataset_get_as(ptycho_cxi_1): assert t.allclose(pattern.to(device='cpu'), t.tensor(expected['data'][3,:,:])) + +def test_Ptycho2DDataset_downsample(test_ptycho_cxis): + for cxi, expected in test_ptycho_cxis: + dataset = Ptycho2DDataset.from_cxi(cxi) + + # First we test the case of downsampling by 2 against some explicit + # calculations + copied_dataset = deepcopy(dataset) + copied_dataset.downsample(2) + + # May start failing if the test datasets are changed to include + # a dataset with any dimension not even. That's a problem with the + # test, not the code. Sorry! -Abe + assert t.allclose( + copied_dataset.patterns, + dataset.patterns[:,::2,::2] + + dataset.patterns[:,1::2,::2] + + dataset.patterns[:,::2,1::2] + + dataset.patterns[:,1::2,1::2] + ) + + assert t.allclose( + copied_dataset.mask, + t.logical_and( + t.logical_and(dataset.mask[::2,::2], + dataset.mask[1::2,::2]), + t.logical_and(dataset.mask[::2,1::2], + dataset.mask[1::2,1::2]), + ) + ) + + + + if dataset.background is not None: + assert t.allclose( + copied_dataset.background, + dataset.background[::2,::2] + + dataset.background[1::2,::2] + + dataset.background[::2,1::2] + + dataset.background[1::2,1::2] + ) + + # And then we just test the shape for a few factors, and check that + # it doesn't fail on edge cases (e.g. factor=1) + for factor in [1, 2, 3]: + copied_dataset = deepcopy(dataset) + copied_dataset.downsample(factor=factor) + + expected_pattern_shape = np.concatenate( + [[dataset.patterns.shape[0]], + np.array(dataset.patterns.shape[-2:]) // factor] + ) + + assert np.allclose(expected_pattern_shape, + np.array(copied_dataset.patterns.shape)) + + assert np.allclose(np.array(dataset.mask.shape) // factor, + np.array(copied_dataset.mask.shape)) + + if dataset.background is not None: + assert np.allclose(np.array(dataset.background.shape) // factor, + np.array(copied_dataset.background.shape)) +