From 6b2358df36b36bd74e35a9b8e11e80a2c37e42e4 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 19 Sep 2024 18:09:37 +0200 Subject: [PATCH] Add a fourier translation to fancy ptycho --- src/cdtools/models/fancy_ptycho.py | 30 ++++++++-- src/cdtools/tools/analysis/analysis.py | 81 ++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 4 deletions(-) diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 611fc72..c3cc4e6 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -25,6 +25,7 @@ class FancyPtycho(CDIModel): background=None, probe_basis=None, translation_offsets=None, + probe_fourier_shifts=None, mask=None, weights=None, translation_scale=1, @@ -134,6 +135,13 @@ class FancyPtycho(CDIModel): t_o = t.as_tensor(translation_offsets, dtype=t.float32) t_o = t_o / translation_scale self.translation_offsets = t.nn.Parameter(t_o) + + if probe_fourier_shifts is None: + self.probe_fourier_shifts = None + else: + self.probe_fourier_shifts = t.nn.Parameter( + t.as_tensor(translation_offsets, dtype=t.float32) + ) self.register_buffer('translation_scale', t.as_tensor(translation_scale, dtype=dtype)) @@ -152,7 +160,7 @@ class FancyPtycho(CDIModel): t.as_tensor(simulate_probe_translation, dtype=bool) ) - if simulate_probe_translation: + if simulate_probe_translation or (self.probe_fourier_shifts is not None): Is = t.arange(self.probe.shape[-2], dtype=dtype) Js = t.arange(self.probe.shape[-1], dtype=dtype) Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js)) @@ -195,6 +203,7 @@ class FancyPtycho(CDIModel): fourier_probe=False, loss='amplitude mse', units='um', + allow_probe_fourier_shifts=False, simulate_probe_translation=False, simulate_finite_pixels=False, exponentiate_obj=False, @@ -327,6 +336,11 @@ class FancyPtycho(CDIModel): translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5) + if allow_probe_fourier_shifts: + probe_fourier_shifts = t.zeros((len(dataset), 2), dtype=t.float32) + else: + probe_fourier_shifts = None + if dm_rank is not None and dm_rank != 0: if dm_rank > n_modes: raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.') @@ -380,6 +394,7 @@ class FancyPtycho(CDIModel): 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, @@ -431,12 +446,19 @@ class FancyPtycho(CDIModel): # Maybe this can be done with a matmul now? prs = t.sum(Ws[..., None, None] * basis_prs, axis=-3) - if self.simulate_probe_translation: - det_pix_trans = tools.interactions.translations_to_pixel( + if self.simulate_probe_translation or (self.probe_fourier_shifts is not None): + if self.probe_fourier_shifts is not None: + det_pix_trans = self.probe_fourier_shifts[index] + else: + det_pix_trans = t.zeros_like(translations) + + if self.simulate_probe_translation: + det_pix_trans = det_pix_trans + tools.interactions.translations_to_pixel( self.det_basis, translations, surface_normal=self.surface_normal) - + + probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] * self.I_phase[None,...] + det_pix_trans[:,1,None,None] * diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py index 4c7c980..d6f2f78 100644 --- a/src/cdtools/tools/analysis/analysis.py +++ b/src/cdtools/tools/analysis/analysis.py @@ -13,6 +13,7 @@ import cdtools from scipy import linalg as sla from scipy import special from scipy import optimize as opt +from scipy import spatial __all__ = [ 'product_svd', @@ -31,6 +32,7 @@ __all__ = [ 'remove_amplitude_exponent', 'standardize_reconstruction_set', 'standardize_reconstruction_pair', + 'calc_spectral_info', ] @@ -1432,3 +1434,82 @@ def standardize_reconstruction_pair( return results + +def calc_spectral_info(dataset, nbins=50): + """Makes a properly normalized sum diffraction pattern + + This returns a scaled version of sum of all the diffraction patterns + within the dataset. The scaling is defined so that the total intensity + in the final image is equal to the intensity arising from a region of + the scan pattern whose area matches one detector conjugate field of + view. + + Parameters + ---------- + dataset : Ptycho2DDataset + A ptychography dataset to use + nbins : int + The number of bins to use for the SNR curve + + Returns + ------- + spectrum : t.tensor + An image of the spectral signal rate + freqs : t.tensor + The frequencies at which the SSNR is estimated + SSNR : t.tensor + The estimated SSNR + + """ + + scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy()) + + scan_area = scan_hull.volume + + ewg = cdtools.tools.initializers.exit_wave_geometry + obj_basis = ewg( + dataset.detector_geometry['basis'], + dataset[0][1].shape, + dataset.wavelength, + dataset.detector_geometry['distance'], + ) + + det_conj_fov_area = np.linalg.norm( + np.cross(obj_basis[:,0]*dataset.patterns.shape[-2], + obj_basis[:,1]*dataset.patterns.shape[-1]) + ) + scale_factor = det_conj_fov_area / scan_area + + + mask = dataset.mask.cpu().numpy().astype(int) + sum_pattern = dataset.mask * t.sum(dataset.patterns, dim=0) * scale_factor + sum_pattern = sum_pattern.cpu().numpy() + + # TODO this assumes orthogonal axes + pix_sizes = np.linalg.norm(obj_basis, axis=0) + i_freqs = np.fft.fftshift(np.fft.fftfreq( + sum_pattern.shape[0],d=pix_sizes[0])) + j_freqs = np.fft.fftshift(np.fft.fftfreq( + sum_pattern.shape[1],d=pix_sizes[1])) + + + Js,Is = np.meshgrid(j_freqs,i_freqs) + Rs = np.sqrt(Is**2+Js**2) + max_i = np.max(i_freqs) + max_j = np.max(j_freqs) + frc_range = [0, max(max_i,max_j)] + + sum_spectrum, frc_bins = np.histogram(Rs, bins=nbins, range=frc_range, + weights=sum_pattern) + sum_spectrum_sq, frc_bins = np.histogram(Rs, bins=nbins, range=frc_range, + weights=sum_pattern**2) + + + n_pix, frc_bins = np.histogram(Rs, bins=nbins, range=frc_range, + weights=mask) + mean_spectrum = sum_spectrum / n_pix + + pattern_snr = sum_spectrum_sq / sum_spectrum + + return sum_pattern, frc_bins[:-1], mean_spectrum +