From ac1499f511ee494a2a32d2673df0f82630819c82 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 7 Apr 2025 11:28:15 +0200 Subject: [PATCH] get it running in principle, but it is not good --- src/cdtools/models/fancy_ptycho.py | 68 +++++++++++++------ .../tools/measurements/measurements.py | 47 +++++++++---- 2 files changed, 83 insertions(+), 32 deletions(-) diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 1be40c6..9459855 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -28,6 +28,7 @@ class FancyPtycho(CDIModel): probe_fourier_shifts=None, mask=None, weights=None, + qe_mask=None, translation_scale=1, saturation=None, probe_support=None, @@ -87,7 +88,14 @@ class FancyPtycho(CDIModel): else: self.register_buffer('mask', t.as_tensor(mask, dtype=t.bool)) - + + + if qe_mask is None: + self.qe_mask = None + else: + self.qe_mask = t.nn.Parameter( + t.as_tensor(qe_mask, dtype=dtype)) + probe_guess = t.as_tensor(probe_guess, dtype=t.complex64) obj_guess = t.as_tensor(obj_guess, dtype=t.complex64) @@ -202,6 +210,7 @@ class FancyPtycho(CDIModel): dm_rank=None, translation_scale=1, saturation=None, + use_qe_mask=False, probe_support_radius=None, probe_fourier_crop=None, propagation_distance=None, @@ -376,6 +385,11 @@ class FancyPtycho(CDIModel): else: mask = None + if use_qe_mask: + qe_mask = t.ones(mask.shape, dtype=t.float32) + else: + qe_mask = None + if probe_support_radius is not None: probe_support = t.zeros(probe[0].shape, dtype=t.bool) xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]] @@ -389,24 +403,34 @@ class FancyPtycho(CDIModel): else: probe_support = None - return cls(wavelength, det_geo, obj_basis, probe, obj, - surface_normal=surface_normal, - min_translation=min_translation, - translation_offsets=translation_offsets, - weights=Ws, mask=mask, background=background, - translation_scale=translation_scale, - saturation=saturation, - probe_basis=probe_basis, - probe_support=probe_support, - fourier_probe=fourier_probe, - oversampling=oversampling, - loss=loss, units=units, - probe_fourier_shifts=probe_fourier_shifts, - simulate_probe_translation=simulate_probe_translation, - simulate_finite_pixels=simulate_finite_pixels, - phase_only=phase_only, - exponentiate_obj=exponentiate_obj, - obj_view_crop=obj_view_crop) + return cls( + wavelength, + det_geo, + obj_basis, + probe, + obj, + surface_normal=surface_normal, + min_translation=min_translation, + translation_offsets=translation_offsets, + weights=Ws, + mask=mask, + background=background, + qe_mask=qe_mask, + translation_scale=translation_scale, + saturation=saturation, + probe_basis=probe_basis, + probe_support=probe_support, + fourier_probe=fourier_probe, + oversampling=oversampling, + loss=loss, + units=units, + probe_fourier_shifts=probe_fourier_shifts, + simulate_probe_translation=simulate_probe_translation, + simulate_finite_pixels=simulate_finite_pixels, + phase_only=phase_only, + exponentiate_obj=exponentiate_obj, + obj_view_crop=obj_view_crop + ) def interaction(self, index, translations, *args): @@ -521,6 +545,7 @@ class FancyPtycho(CDIModel): wavefields, self.background, measurement=tools.measurements.incoherent_sum, + qe_mask=self.qe_mask, saturation=self.saturation, oversampling=self.oversampling, simulate_finite_pixels=self.simulate_finite_pixels, @@ -840,7 +865,10 @@ class FancyPtycho(CDIModel): ('Corrected Translations', lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)), ('Background', - lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)) + lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)), + ('Quantum Efficiency Mask', + lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig), + lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None)) ] diff --git a/src/cdtools/tools/measurements/measurements.py b/src/cdtools/tools/measurements/measurements.py index e726149..8ce41ae 100644 --- a/src/cdtools/tools/measurements/measurements.py +++ b/src/cdtools/tools/measurements/measurements.py @@ -155,7 +155,18 @@ 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, simulate_finite_pixels=False): +def quadratic_background( + wavefield, + background, + *args, + detector_slice=None, + measurement=intensity, + epsilon=1e-7, + qe_mask=None, + 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 @@ -173,6 +184,8 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas Optional, a slice or tuple of slices defining a section of the simulation to return measurement : function Default is measurements.intensity, the measurement function to use. + qe_mask : torch.Tensor + A tensor storing the per-pixel quantum efficiency (up to an unknown global scaling factor) saturation : float Optional, a maximum saturation value to clamp the resulting intensities to oversampling : int @@ -183,18 +196,28 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas sim_patterns : torch.Tensor A real MxN array storing the wavefield's intensities """ - - if detector_slice is None: - output = measurement(wavefield, *args, epsilon=epsilon, - oversampling=oversampling, - simulate_finite_pixels=simulate_finite_pixels) \ - + background**2 - else: - output = measurement(wavefield, *args, detector_slice=detector_slice, - epsilon=epsilon, oversampling=oversampling, - simulate_finite_pixels=simulate_finite_pixels) \ - + background**2 + if detector_slice is None: + raw_intensity = measurement( + wavefield, + *args, + epsilon=epsilon, + oversampling=oversampling, + simulate_finite_pixels=simulate_finite_pixels) + else: + raw_intensity = measurement( + wavefield, + *args, + detector_slice=detector_slice, + epsilon=epsilon, + oversampling=oversampling, + simulate_finite_pixels=simulate_finite_pixels) + + if qe_mask is None: + output = raw_intensity + background**2 + else: + output = (qe_mask * raw_intensity) + background**2 + # This has to be done after the background is added, hence we replicate # it here if saturation is None: