From c86e505ae5e5871e8da88e2aa1339d872eae126c Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 13 May 2025 14:06:02 +0200 Subject: [PATCH] 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 | 20 ++++++-- src/cdtools/datasets/ptycho_2d_dataset.py | 56 ++++++++++++++++++----- src/cdtools/models/fancy_ptycho.py | 9 +++- 3 files changed, 69 insertions(+), 16 deletions(-) diff --git a/src/cdtools/datasets/base.py b/src/cdtools/datasets/base.py index 9dc4801..85ec854 100644 --- a/src/cdtools/datasets/base.py +++ b/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: diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py index 8b6f6d8..f5b92f9 100644 --- a/src/cdtools/datasets/ptycho_2d_dataset.py +++ b/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 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. + 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, diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 9459855..0f0323e 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/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