Add a function to downsample a dataset by an integer factor

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 <factor> x <factor> 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]

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))
+