From 978c9d32068bfb55041302bb9f502219d0cc8a42 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 29 Sep 2022 14:14:43 -0700 Subject: [PATCH] Add an option to simulate finite pixel sizes --- src/cdtools/models/fancy_ptycho.py | 15 +++-- src/cdtools/models/rpi.py | 23 +++++-- .../tools/measurements/measurements.py | 60 ++++++++++++++++--- 3 files changed, 80 insertions(+), 18 deletions(-) diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 92afc65..7812b4c 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -33,7 +33,8 @@ class FancyPtycho(CDIModel): fourier_probe=False, loss='amplitude mse', units='um', - simulate_probe_translation=False + simulate_probe_translation=False, + simulate_finite_pixels=False, ): super(FancyPtycho, self).__init__() @@ -124,6 +125,8 @@ class FancyPtycho(CDIModel): Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js)) self.I_phase = 2 * np.pi* Is self.J_phase = 2 * np.pi* Js + + self.simulate_finite_pixels = simulate_finite_pixels # Here we set the appropriate loss function if (loss.lower().strip() == 'amplitude mse' @@ -155,7 +158,8 @@ class FancyPtycho(CDIModel): fourier_probe=False, loss='amplitude mse', units='um', - simulate_probe_translation=False + simulate_probe_translation=False, + simulate_finite_pixels=False, ): wavelength = dataset.wavelength @@ -209,6 +213,7 @@ class FancyPtycho(CDIModel): # Next generate the object geometry from the probe geometry and # the translations + pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal) obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200) @@ -293,7 +298,8 @@ class FancyPtycho(CDIModel): fourier_probe=fourier_probe, oversampling=oversampling, loss=loss, units=units, - simulate_probe_translation=simulate_probe_translation) + simulate_probe_translation=simulate_probe_translation, + simulate_finite_pixels=simulate_finite_pixels) def interaction(self, index, translations, *args): @@ -376,7 +382,8 @@ class FancyPtycho(CDIModel): detector_slice=self.detector_slice, measurement=tools.measurements.incoherent_sum, saturation=self.saturation, - oversampling=self.oversampling) + oversampling=self.oversampling, + simulate_finite_pixels=self.simulate_finite_pixels) # Note: No "loss" function is defined here, because it is added diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index 0bf9981..0d588a6 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -62,7 +62,8 @@ class RPI(CDIModel): def __init__(self, wavelength, detector_geometry, probe_basis, probe, obj_guess, detector_slice=None, background=None, mask=None, saturation=None, - obj_support=None, oversampling=1, weight_matrix=False): + obj_support=None, oversampling=1, weight_matrix=False, + simulate_finite_pixels=False): super(RPI, self).__init__() @@ -134,9 +135,11 @@ class RPI(CDIModel): self.oversampling = oversampling + self.simulate_finite_pixels=simulate_finite_pixels + @classmethod - def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0): + def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0, simulate_finite_pixels=False): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -213,7 +216,13 @@ class RPI(CDIModel): # on a coarser grid needs to be accounted for here that is not # accounted for yet scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2) - obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size)) + obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+list(obj_size))) + elif initialization.lower().strip() == 'uniform': + # I think something to do with the fact that the object is defined + # on a coarser grid needs to be accounted for here that is not + # accounted for yet + scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2) + obj_guess = scale * t.ones([n_modes,]+list(obj_size), dtype=probe.dtype) elif initialization.lower().strip() == 'spectral': if background is not None: obj_guess = initializers.RPI_spectral_init( @@ -226,7 +235,7 @@ class RPI(CDIModel): else: raise KeyError('Initialization "' + str(initialization) + \ - '" invalid - use "spectral" or "random"') + '" invalid - use "spectral", "uniform", or "random"') probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0)) probe_fft = tools.propagators.far_field(probe_intensity) @@ -244,7 +253,8 @@ class RPI(CDIModel): probe, obj_guess, detector_slice=det_slice, background=background, mask=mask, saturation=saturation, obj_support=obj_support, oversampling=oversampling, - weight_matrix=weight_matrix) + weight_matrix=weight_matrix, + simulate_finite_pixels=simulate_finite_pixels) @classmethod def from_calibration(cls, calibration, obj_size=None, n_modes=1, saturation=None): @@ -353,7 +363,8 @@ class RPI(CDIModel): detector_slice=self.detector_slice, measurement=tools.measurements.incoherent_sum, saturation=self.saturation, - oversampling=self.oversampling) + oversampling=self.oversampling, + simulate_finite_pixels=self.simulate_finite_pixels) return m def loss(self, sim_data, real_data, mask=None): diff --git a/src/cdtools/tools/measurements/measurements.py b/src/cdtools/tools/measurements/measurements.py index bbd7084..f49e04b 100644 --- a/src/cdtools/tools/measurements/measurements.py +++ b/src/cdtools/tools/measurements/measurements.py @@ -17,7 +17,7 @@ from torch.nn.functional import avg_pool2d __all__ = ['intensity', 'incoherent_sum', 'quadratic_background'] -def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1): +def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1, simulate_finite_pixels=False): """Returns the intensity of a wavefield The intensity is defined as the magnitude squared of the @@ -40,8 +40,30 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove sim_patterns : torch.Tensor A real MxN array storing the wavefield's intensities """ - output = t.abs(wavefield)**2 + if simulate_finite_pixels: + inverse_fft = t.fft.fftshift(t.fft.ifft2(wavefield), dim=(-2,-1)) + pad1l = wavefield.shape[-2]//2 + pad1r = wavefield.shape[-2] - pad1l + pad2l = wavefield.shape[-1]//2 + pad2r = wavefield.shape[-1] - pad2l + padded = t.nn.functional.pad(inverse_fft, (pad1l, pad1r, pad2l, pad2r)) + upsampled_field = t.fft.fft2(t.fft.ifftshift(padded, dim=(-2,-2))) + upsampled_intensity = t.abs(upsampled_field)**2 + ifft_intensity = t.fft.fftshift(t.fft.ifft2(upsampled_intensity), dim=(-2,-1)) + # Now we take a sinc function + xs = t.arange(ifft_intensity.shape[-2]) + xs = (xs / t.max(xs)) * 2 - 1 + ys = t.arange(ifft_intensity.shape[-1]) + ys = (ys / t.max(ys)) * 2 - 1 + Xs, Ys = t.meshgrid(xs, ys, indexing='ij') + mask = (t.special.sinc(Xs) * t.special.sinc(Ys)).to(device=ifft_intensity.device) + blurred_intensity = t.fft.fft2(t.fft.ifftshift(mask * ifft_intensity, dim=(-2,-2))) + output = blurred_intensity[...,::2,::2] + else: + output = t.abs(wavefield)**2 + + # Now we apply oversampling if oversampling != 1: if wavefield.dim() == 2: @@ -63,7 +85,7 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove return t.clamp(output + epsilon,0,saturation) -def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1): +def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1, simulate_finite_pixels=False): """Returns the incoherent sum of the intensities of the wavefields The intensity is defined as the sum of the magnitudes squared of @@ -91,8 +113,27 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non sim_patterns : torch.Tensor A real LXMxN array storing the incoherently summed intensities """ - - output = t.sum(t.abs(wavefields)**2,dim=-3) + if simulate_finite_pixels: + inverse_fft = t.fft.fftshift(t.fft.ifft2(wavefields), dim=(-2,-1)) + pad1l = wavefields.shape[-2]//2 + pad1r = wavefields.shape[-2] - pad1l + pad2l = wavefields.shape[-1]//2 + pad2r = wavefields.shape[-1] - pad2l + padded = t.nn.functional.pad(inverse_fft, (pad1l, pad1r, pad2l, pad2r)) + upsampled_field = t.fft.fft2(t.fft.ifftshift(padded, dim=(-2,-2))) + upsampled_intensity = t.sum(t.abs(upsampled_field)**2, dim=-3) + ifft_intensity = t.fft.fftshift(t.fft.ifft2(upsampled_intensity), dim=(-2,-1)) + # Now we take a sinc function + xs = t.arange(ifft_intensity.shape[-2]) + xs = (xs / t.max(xs)) * 2 - 1 + ys = t.arange(ifft_intensity.shape[-1]) + ys = (ys / t.max(ys)) * 2 - 1 + Xs, Ys = t.meshgrid(xs, ys, indexing='ij') + mask = (t.special.sinc(Xs) * t.special.sinc(Ys)).to(device=ifft_intensity.device) + blurred_intensity = t.fft.fft2(t.fft.ifftshift(mask * ifft_intensity, dim=(-2,-2))) + output = t.abs(blurred_intensity[...,::2,::2]) + else: + output = t.sum(t.abs(wavefields)**2,dim=-3) # Now we apply oversampling if oversampling != 1: @@ -114,7 +155,7 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non return t.clamp(output + epsilon,0,saturation) -def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1): +def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1, simulate_finite_pixels=False): """Returns the intensity of a wavefield plus a background The intensity is calculated via the given measurment function @@ -145,10 +186,13 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas if detector_slice is None: output = measurement(wavefield, *args, epsilon=epsilon, - oversampling=oversampling) + background**2 + oversampling=oversampling, + simulate_finite_pixels=simulate_finite_pixels) \ + + background**2 else: output = measurement(wavefield, *args, detector_slice=detector_slice, - epsilon=epsilon, oversampling=oversampling) \ + epsilon=epsilon, oversampling=oversampling, + simulate_finite_pixels=simulate_finite_pixels) \ + background**2 # This has to be done after the background is added, hence we replicate