Make it possible to add a quantum efficiency mask to a dataset, allow fancy_ptycho to load that and use it in the forward model

src/cdtools/datasets/base.py

@@ -34,10 +34,16 @@ class CDataset(torchdata.Dataset):
     needed to allow for easy mixing of data on the CPU and GPU.
     """
 
-    def __init__(self, entry_info=None, sample_info=None,
-                 wavelength=None,
-                 detector_geometry=None, mask=None,
-                 background=None):
+    def __init__(
+            self,
+            entry_info=None,
+            sample_info=None,
+            wavelength=None,
+            detector_geometry=None,
+            mask=None,
+            qe_mask=None,
+            background=None,
+    ):
 
         """The __init__ function allows construction from python objects.
 
@@ -73,6 +79,12 @@ class CDataset(torchdata.Dataset):
             self.mask = t.tensor(mask, dtype=t.bool)
         else:
             self.mask = None
+
+        if qe_mask is not None:
+            self.qe_mask = t.as_tensor(qe_mask, dtype=t.float32)
+        else:
+            self.qe_mask = None
+
         if background is not None:
             self.background = t.tensor(background, dtype=t.float32)
         else:

src/cdtools/datasets/ptycho_2d_dataset.py

@@ -77,6 +77,7 @@ class Ptycho2DDataset(CDataset):
             self.intensities = t.as_tensor(intensities, dtype=t.float32)
         else:
             self.intensities = None
+
             
     def __len__(self):
         return self.patterns.shape[0]
@@ -372,10 +373,19 @@ class Ptycho2DDataset(CDataset):
         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.
+        Any mask, quantum efficiency, 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.
+
+        If there is no quantum efficiency mask, then the mask is downsapled so
+        that any output pixel containing a masked pixel will be masked. If there
+        is a quantum efficiency mask, then the quantum efficiency mask is
+        downsampled using the same method as the data, and the mask is
+        downsampled to include any pixels for which there is at least one valid
+        pixel.
+
+        To avoid leakage of data from masked pixels, the data is first
+        multiplied by the mask before downsampling.
         
         Parameters
         ----------
@@ -383,17 +393,41 @@ class Ptycho2DDataset(CDataset):
             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))
+        if hasattr(self, 'mask') and self.mask is not None:
+            self.patterns = t.nn.functional.avg_pool2d(
+                (self.mask * self.patterns).unsqueeze(0),
+                factor, divisor_override=1)[0]
+        else:
+            self.patterns = t.nn.functional.avg_pool2d(
+                self.patterns.unsqueeze(0),
+                factor, divisor_override=1)[0]
+            
+
+        # If we have a QE mask, we want to include all pixels for which at
+        # least one of the input pixels was unmasked, because we can account
+        # for the masked pixels through quantum efficiency
+        if hasattr(self, 'qe_mask') and self.qe_mask is not None:
+            self.qe_mask = t.nn.functional.avg_pool2d(
+                (self.mask * self.qe_mask).unsqueeze(0).unsqueeze(0),
+                factor)[0,0]
+            self.mask = t.nn.functional.max_pool2d(
+                self.mask.to(dtype=t.uint8).unsqueeze(0).unsqueeze(0),
+                factor)[0,0].to(dtype=t.bool)
+            
+        # But if there is no QE mask, we need to only preserve pixels for
+        # which all input pixels were unmasked
+        elif hasattr(self, 'mask') and self.mask is not None:
+            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:
+
+        
+        if hasattr(self, 'background') and self.background is not None:
             self.background = t.nn.functional.avg_pool2d(
                 self.background.unsqueeze(0).unsqueeze(0),
                 factor,

src/cdtools/models/fancy_ptycho.py

@@ -95,6 +95,10 @@ class FancyPtycho(CDIModel):
         else:
             self.qe_mask = t.nn.Parameter(
                 t.as_tensor(qe_mask, dtype=dtype))
+            # I want the ability to optimize over this, but experience shows
+            # that it is wildly unstable, so I think it's best to keep
+            # gradients turned off by default
+            self.qe_mask.requires_grad=False
             
         probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
         obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
@@ -386,7 +390,10 @@ class FancyPtycho(CDIModel):
             mask = None
 
         if use_qe_mask:
-            qe_mask = t.ones(mask.shape, dtype=t.float32)
+            if hasattr(dataset, 'qe_mask') and dataset.qe_mask is not None:
+                qe_mask = t.as_tensor(dataset.qe_mask, dtype=t.float32)
+            else:
+                qe_mask = t.ones(dataset.patterns.shape[-2:], dtype=t.float32)
         else:
             qe_mask = None