From b3b19554d07e8568f990566b7cde82d2b229758a Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Fri, 30 Aug 2024 11:29:51 +0200 Subject: [PATCH] Add a phase only constraint to RPI and fancy_ptycho, misc small changes --- src/cdtools/models/fancy_ptycho.py | 70 +++++- src/cdtools/models/rpi.py | 103 +++++++-- src/cdtools/tools/analysis/analysis.py | 205 ++++++++++++++++-- .../tools/interactions/interactions.py | 15 +- .../tools/measurements/measurements.py | 2 +- 5 files changed, 340 insertions(+), 55 deletions(-) diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 09e2718..7b4a540 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -36,6 +36,8 @@ class FancyPtycho(CDIModel): units='um', simulate_probe_translation=False, simulate_finite_pixels=False, + exponentiate_obj=False, + phase_only=False, dtype=t.float32, obj_view_crop=0 ): @@ -70,6 +72,12 @@ class FancyPtycho(CDIModel): self.register_buffer('fourier_probe', t.tensor(fourier_probe, dtype=bool)) + self.register_buffer('exponentiate_obj', + t.tensor(exponentiate_obj, dtype=bool)) + + self.register_buffer('phase_only', + t.tensor(phase_only, dtype=bool)) + # Not sure how to make this a buffer... self.units = units @@ -82,6 +90,7 @@ class FancyPtycho(CDIModel): probe_guess = t.tensor(probe_guess, dtype=t.complex64) obj_guess = t.tensor(obj_guess, dtype=t.complex64) + # We rescale the probe here so it learns at the same rate as the # object if probe_guess.dim() > 2: @@ -186,6 +195,8 @@ class FancyPtycho(CDIModel): units='um', simulate_probe_translation=False, simulate_finite_pixels=False, + exponentiate_obj=False, + phase_only=False, obj_view_crop=None, obj_padding=200, ): @@ -292,10 +303,16 @@ class FancyPtycho(CDIModel): probe = t.stack([probe, ] + probe_stack) - obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) + obj = (randomize_ang * (t.rand(obj_size)-0.5)).to(dtype=t.complex64) + if not exponentiate_obj: + obj = t.exp(1j * obj) + if n_obj_modes != 1: obj = t.stack([obj,] + [0.05*t.ones_like(obj),]*(n_obj_modes-1)) + if phase_only: + obj.imag[:] = 0 + pfc = (probe_fourier_crop if probe_fourier_crop else 0) if obj_view_crop is None: obj_view_crop = min(probe.shape[-2], probe.shape[-1]) // 2 + pfc @@ -363,6 +380,8 @@ class FancyPtycho(CDIModel): loss=loss, units=units, 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) @@ -437,12 +456,23 @@ class FancyPtycho(CDIModel): prs = t.nn.functional.pad(prs, padding) prs = tools.propagators.inverse_far_field(prs) + + if self.exponentiate_obj: + if self.phase_only: + obj = t.exp(1j*self.obj.real) + else: + obj = t.exp(1j*self.obj) + else: + obj = self.obj + # Now we actually do the interaction, using the sinc subpixel # translation model as per usual exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( - prs, self.obj, pix_trans, - shift_probe=True, multiple_modes=True) + prs, obj, pix_trans, + shift_probe=True, + multiple_modes=True, + probe_support=self.probe_support) return exit_waves @@ -707,12 +737,12 @@ class FancyPtycho(CDIModel): ('Basis Probe Fourier Space Amplitudes', lambda self, fig: p.plot_amplitude( (self.probe if self.fourier_probe - else tools.propagators.inverse_far_field(self.probe)), + else tools.propagators.far_field(self.probe)), fig=fig)), - ('Basis Probe Fourier Space Phases', - lambda self, fig: p.plot_phase( + ('Basis Probe Fourier Space Colorized', + lambda self, fig: p.plot_colorized( (self.probe if self.fourier_probe - else tools.propagators.inverse_far_field(self.probe)) + else tools.propagators.far_field(self.probe)) , fig=fig)), ('Basis Probe Real Space Amplitudes', lambda self, fig: p.plot_amplitude( @@ -721,8 +751,8 @@ class FancyPtycho(CDIModel): fig=fig, basis=self.probe_basis, units=self.units)), - ('Basis Probe Real Space Phases', - lambda self, fig: p.plot_phase( + ('Basis Probe Real Space Colorized', + lambda self, fig: p.plot_colorized( (self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe)), fig=fig, @@ -750,13 +780,31 @@ class FancyPtycho(CDIModel): self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, - units=self.units)), + units=self.units), + lambda self: not self.exponentiate_obj), ('Object Phase', lambda self, fig: p.plot_phase( self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, - units=self.units)), + units=self.units), + lambda self: not self.exponentiate_obj), + ('Real Part of T', + lambda self, fig: p.plot_real( + self.obj[self.obj_view_slice], + fig=fig, + basis=self.obj_basis, + units=self.units, + cmap='cividis'), + lambda self: self.exponentiate_obj), + ('Imaginary Part of T', + lambda self, fig: p.plot_imag( + self.obj[self.obj_view_slice], + fig=fig, + basis=self.obj_basis, + units=self.units), + lambda self: self.exponentiate_obj), + ('Corrected Translations', lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)), ('Background', diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index d8ae029..9d8a9ff 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -53,6 +53,8 @@ class RPI(CDIModel): obj_support=None, oversampling=1, weight_matrix=False, + exponentiate_obj=False, + phase_only=False, propagation_distance=0, units='um', dtype=t.float32, @@ -97,15 +99,22 @@ class RPI(CDIModel): self.register_buffer('probe', t.tensor(probe, dtype=complex_dtype)) + + self.register_buffer('exponentiate_obj', + t.tensor(exponentiate_obj, dtype=bool)) + + self.register_buffer('phase_only', + t.tensor(phase_only, dtype=bool)) + # We always use multi-modes to store the object, so we convert it # if we just get a single 2D array as an input if obj_guess.dim() == 2: obj_guess = obj_guess[None, :, :] self.obj = t.nn.Parameter(t.tensor(obj_guess, dtype=complex_dtype)) - + self.weights = t.nn.Parameter( - t.eye(probe.shape[0], dtype=complex_dtype)* 10) + t.eye(probe.shape[0], dtype=complex_dtype)) if not weight_matrix: self.weights.requires_grad=False @@ -165,6 +174,8 @@ class RPI(CDIModel): oversampling=1, initialization='random', weight_matrix=False, + exponentiate_obj=False, + phase_only=False, probe_threshold=0, dtype=t.float32, ): @@ -253,7 +264,10 @@ class RPI(CDIModel): probe, dummy_init_obj, background=background, mask=mask, saturation=saturation, - obj_support=obj_support, oversampling=oversampling, + obj_support=obj_support, + oversampling=oversampling, + exponentiate_obj=exponentiate_obj, + phase_only=phase_only, weight_matrix=weight_matrix) # I don't love this pattern, where I do the "real" obj initialization @@ -263,6 +277,13 @@ class RPI(CDIModel): # is a better pattern for doing this. rpi_object.init_obj(initialization, pattern=dataset.patterns[0]) + + if exponentiate_obj: + rpi_object.obj.data = -1j * t.log(rpi_object.obj.data) + + if phase_only: + rpi_object.obj.data.imag[:] = 0 + return rpi_object @@ -273,6 +294,8 @@ class RPI(CDIModel): obj_size=None, n_modes=1, saturation=None, # TODO can we get this from the calibration? + exponentiate_obj=False, + phase_only=False, initialization='random', dtype=t.float32 ): @@ -316,10 +339,18 @@ class RPI(CDIModel): dummy_init_obj, background=background, mask=mask, + exponentiate_obj=exponentiate_obj, + phase_only=phase_only, ) rpi_object.init_obj(initialization) + if exponentiate_obj: + rpi_object.obj.data = -1j * t.log(rpi_object.obj.data) + + if phase_only: + rpi_object.obj.data.imag[:] = 0 + return rpi_object @@ -443,20 +474,28 @@ class RPI(CDIModel): # Mix the probes with the weight matrix prs = t.sum(self.weights[..., None, None] * self.probe, axis=-3) + + if self.exponentiate_obj: + if self.phase_only: + obj = t.exp(1j*self.obj.real) + else: + obj = t.exp(1j*self.obj) + else: + obj = self.obj + for i in range(self.probe.shape[0]): pr = prs[i] # Here we have a 3D probe (one single mode) # and a 4D object (multiple modes mixing incoherently) exit_waves = RPI_interaction(pr, - self.obj_support * self.obj) + self.obj_support * obj) all_exit_waves.append(exit_waves) # This creates a bunch of modes generated from all possible combos # of the probe and object modes all strung out along the first index output = t.cat(all_exit_waves) - # If we have multiple indexes input, we unsqueeze and repeat the stack # of wavefields enough times to simulate each requested index. This # seems silly, but it enables (for example) one to do a reconstruction @@ -484,11 +523,12 @@ class RPI(CDIModel): # Here I'm taking advantage of an undocumented feature in the # incoherent_sum measurement function where it will work with # a 4D wavefield array as well as a 5D array. - m = tools.measurements.quadratic_background(wavefields, - self.background, - measurement=tools.measurements.incoherent_sum, - saturation=self.saturation, - oversampling=self.oversampling) + m = tools.measurements.quadratic_background( + wavefields, + self.background, + measurement=tools.measurements.incoherent_sum, + saturation=self.saturation, + oversampling=self.oversampling) return m def loss(self, sim_data, real_data, mask=None): @@ -510,20 +550,35 @@ class RPI(CDIModel): lambda self, fig: p.plot_amplitude( np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)), fig=fig, basis=self.probe_basis)), - ('Dominant Object Amplitude', - lambda self, fig: p.plot_amplitude(self.obj[0], fig=fig, - basis=self.obj_basis)), - ('Dominant Object Phase', - lambda self, fig: p.plot_phase(self.obj[0], fig=fig, - basis=self.obj_basis)), - ('Subdominant Object Amplitude', - lambda self, fig: p.plot_amplitude(self.obj[1], fig=fig, - basis=self.obj_basis), - lambda self: len(self.obj) >=2), - ('Subdominant Object Phase', - lambda self, fig: p.plot_phase(self.obj[1], fig=fig, - basis=self.obj_basis), - lambda self: len(self.obj) >=2) + ('Object Amplitude', + lambda self, fig: p.plot_amplitude( + self.obj, + fig=fig, + basis=self.obj_basis, + units=self.units), + lambda self: not self.exponentiate_obj), + ('Object Phase', + lambda self, fig: p.plot_phase( + self.obj, + fig=fig, + basis=self.obj_basis, + units=self.units), + lambda self: not self.exponentiate_obj), + ('Real Part of T', + lambda self, fig: p.plot_real( + self.obj, + fig=fig, + basis=self.obj_basis, + units=self.units, + cmap='cividis'), + lambda self: self.exponentiate_obj), + ('Imaginary Part of T', + lambda self, fig: p.plot_imag( + self.obj, + fig=fig, + basis=self.obj_basis, + units=self.units), + lambda self: self.exponentiate_obj), ] diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py index 5da67e5..91e1597 100644 --- a/src/cdtools/tools/analysis/analysis.py +++ b/src/cdtools/tools/analysis/analysis.py @@ -28,7 +28,8 @@ __all__ = [ 'calc_generalized_rms_error', 'remove_phase_ramp', 'remove_amplitude_exponent', - 'standardize_reconstruction_set' + 'standardize_reconstruction_set', + 'standardize_reconstruction_pair', ] @@ -569,6 +570,8 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'): cor_fft = f1 * t.conj(f2) #from cdtools.tools import plotting as p #from matplotlib import pyplot as plt + #p.plot_phase(im1[im_slice], cmap='cividis') + #p.plot_phase(im2[im_slice], cmap='cividis') #p.plot_amplitude(t.log(t.abs(cor_fft))) #p.plot_amplitude(t.log(t.abs(f1))) #p.plot_phase(cor_fft) @@ -576,7 +579,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'): F1 = t.abs(f1)**2 F2 = t.abs(f2)**2 - + # TODO this is still incorrect if the two bases arent equal di = np.linalg.norm(basis[:,0]) dj = np.linalg.norm(basis[:,1]) @@ -590,6 +593,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'): max_i = np.max(i_freqs) max_j = np.max(j_freqs) frc_range = [0, max(max_i,max_j)] + elif limit.lower().strip() == 'corner': frc_range = [0, np.max(Rs)] else: @@ -607,7 +611,13 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'): #n_pix = n_pix / 4 # This is for an apodized image, apodized with a hann window frc = numerator / np.sqrt(denominator_F1*denominator_F2) - + #plt.figure() + #plt.plot(np.real(frc)) + #plt.title('real') + #plt.figure() + #plt.plot(np.imag(frc)) + #plt.title('imag') + #plt.show() # This moves from combined-image SNR to single-image SNR snr /= 2 @@ -765,6 +775,8 @@ def calc_rms_error(field_1, field_2, align_phases=True, normalize=False, The RMS error, or tensor of RMS errors, depending on the dim argument """ + fields_1 = t.as_tensor(fields_1) + fields_2 = t.as_tensor(fields_2) sumdims = tuple(d - dims for d in range(dims)) @@ -928,7 +940,7 @@ def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2): return t.sqrt(result) -def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, snr=1.): +def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, snr=1., limit='side'): """Calculates a Fourier ring correlation between two images This function requires an input of a basis to allow for FRC calculations @@ -965,7 +977,7 @@ def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, s im_np = False if isinstance(fields_1, np.ndarray): - fields_1 = t.as_tensor(fields_) + fields_1 = t.as_tensor(fields_1) im_np = True if isinstance(fields_2, np.ndarray): fields_2 = t.as_tensor(fields_2) @@ -975,17 +987,15 @@ def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, s basis = t.tensor(basis) if im_slice is None: - im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5, - (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] + im_slice = np.s_[...,:,:] if nbins is None: nbins = np.max(fields_1[...,im_slice].shape[-2:]) // 4 - f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) - f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)) + f1 = t.fft.fftshift(t.fft.fft2(fields_1[im_slice]),dim=(-1,-2)) + f2 = t.fft.fftshift(t.fft.fft2(fields_2[im_slice]),dim=(-1,-2)) cor_fft = f1 * t.conj(f2) - F1 = t.abs(f1)**2 F2 = t.abs(f2)**2 @@ -994,18 +1004,30 @@ def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, s di = np.linalg.norm(basis[:,0]) dj = np.linalg.norm(basis[:,1]) - i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[0],d=di)) - j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[1],d=dj)) + i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[-2],d=di)) + j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[-1],d=dj)) Js,Is = np.meshgrid(j_freqs,i_freqs) Rs = np.sqrt(Is**2+Js**2) + if limit.lower().strip() == 'side': + max_i = np.max(i_freqs) + max_j = np.max(j_freqs) + frc_range = [0, max(max_i,max_j)] + + elif limit.lower().strip() == 'corner': + frc_range = [0, np.max(Rs)] + else: + raise ValueError('Invalid FRC limit: choose "side" or "corner"') + # This line is used to get a set of bins that matches the logic # used by np.histogram, so that this function will match the choices # of bin edges that comes from the non-generalized version. This also # gets us the count on the number of pixels per bin so we can calculate # the threshold curve - n_pix, bins = np.histogram(Rs,bins=nbins) + n_pix, bins = np.histogram(Rs,bins=nbins, range=frc_range) + bins = t.as_tensor(bins) + Rs = t.as_tensor(Rs) frc = [] for i in range(len(bins)-1): @@ -1202,8 +1224,8 @@ def standardize_reconstruction_set( if correct_phase_offset: obj_1 = np.exp(-1j* np.angle(np.sum(obj_1[window]))) * obj_1 - obj_2 = np.exp(-1j* np.angle(np.sum(obj_2))) * obj_2 - obj = np.exp(-1j* np.angle(np.sum(obj))) * obj + obj_2 = np.exp(-1j* np.angle(np.sum(obj_2[window]))) * obj_2 + obj = np.exp(-1j* np.angle(np.sum(obj[window]))) * obj # Todo update the translations to account for the determined shift @@ -1255,3 +1277,156 @@ def standardize_reconstruction_set( return results + +def standardize_reconstruction_pair( + half_1, + half_2, + correct_phase_offset=True, + correct_phase_ramp=True, + correct_amplitude_exponent=False, + window=np.s_[:,:], + nbins=50, + probe_nbins=50, + frc_limit='side', +): + """Standardizes and analyses a set of two repeat + + It's very common to run two subsequent ptycho reconstructions, so that the + effect of sample damage during the first reconstruction can be used in the + estimate of thefinal quality. The difference between the two datasets + datasets can be used to estimate the quality and resolution of each one. + But to do that analysis, first the reconstructions need to be aligned + with respect to each other and normalized in a few ways. + + This function takes the results (as output by model.save_results) of + a pair of reconstructions and: + + - Aligns the object reconstructions with one another + - Corrects for the global phase offset (by default) + - Sets a sensible value for the object/probe phase ramp (by default) + - Sets a sensible value for the object/probe exponential decay (off by + default) + - Calculates the FRC and derived SSNR. + + Then, these results are packaged into an output dictionary. The output + does not retain all the information from the inputs, so if full traceability + is desired, do not delete the files containing the individual + reconstructions + + Parameters + ---------- + half_1 : dict + The result of the first half dataset, as returned by model.save_results + half_2 : dict + The result of the second half dataset, as returned by model.save_results + + Returns + ------- + results : dict + A dictionary containing the synthesized results + """ + # We get the two half-data reconstructions + obj_1, probe_1, weights_1 = half_1['obj'],half_1['probe'],half_1['weights'] + obj_2, probe_2, weights_2 = half_2['obj'],half_2['probe'],half_2['weights'] + + + if correct_phase_ramp: + obj_1, probe_1 = remove_phase_ramp( + half_1['obj'], window, probe=half_1['probe']) + obj_2, probe_2 = remove_phase_ramp( + half_2['obj'], window, probe=half_2['probe']) + + if correct_amplitude_exponent: + obj_1, probe_1, weights_1 = remove_amplitude_exponent( + obj_1, window, probe=probe_1, + weights=half_1['weights'], + basis=half_1['basis'], + translations=half_1['translations']) + obj_2, probe_2, weights_2 = remove_amplitude_exponent( + obj_2, window, probe=probe_2, + weights=half_2['weights'], + basis=half_2['basis'], + translations=half_2['translations']) + + + if correct_phase_offset: + obj_1 = np.exp(-1j* np.angle(np.sum(obj_1[window]))) * obj_1 + obj_2 = np.exp(-1j* np.angle(np.sum(obj_2[window]))) * obj_2 + + + # Todo update the translations to account for the determined shift + shift = ip.find_shift( + t.as_tensor(ip.hann_window(obj_1[window])), + t.as_tensor(ip.hann_window(obj_2[window]))) + obj_2 = ip.sinc_subpixel_shift( + t.as_tensor(obj_2), shift).numpy() + + probe_shift = ip.find_shift( + t.as_tensor(probe_1[0]), + t.as_tensor(probe_2[0]), + ) + for idx in range(probe_2.shape[0]): + probe_2[idx] = ip.sinc_subpixel_shift( + t.as_tensor(probe_2[idx]), probe_shift).numpy() + + # TODO I'm not sure if the default threshold treats this case right tbh + freqs, frc, threshold = calc_frc( + ip.hann_window(obj_1[window]), + ip.hann_window(obj_2[window]), + half_1['obj_basis'], + nbins=nbins, + limit=frc_limit, + ) + + + #probe_freqs, probe_frc, probe_frc_threshold = calc_generalized_frc( + # probe_1, + # probe_2, + # half_1['probe_basis'], + # nbins=probe_nbins, + # limit=frc_limit, + #) + probe_freqs, probe_frc, probe_frc_threshold = calc_generalized_frc( + probe_1, + probe_2, + half_1['probe_basis'], + nbins=probe_nbins, + limit=frc_limit, + ) + + probe_nrms_error = calc_generalized_rms_error( + probe_1, + probe_2, + normalize=True + ) + + # The correct formulation when the final output is one of the two + # reconstructions + ssnr = np.abs(frc) / (1 - np.abs(frc)) + + results = { + 'obj_1': obj_1, + 'probe_1': probe_1, + 'weights_1': weights_1, + 'translations_1': half_1['translations'], + 'background_1': half_1['background'], + 'obj_2': obj_2, + 'probe_2': probe_2, + 'weights_2': weights_2, + 'translations_2': half_2['translations'], + 'background_2': half_2['background'], + 'wavelength': half_1['wavelength'], + 'obj_basis': half_1['obj_basis'], + 'probe_basis': half_1['probe_basis'], + 'frc_freqs': freqs, + 'frc': frc, + 'frc_threshold': threshold, + 'ssnr': ssnr, + 'probe_freqs': probe_freqs, + 'probe_frc': probe_frc, + 'probe_frc_threshold': probe_frc_threshold, + 'probe_nrms_error': probe_nrms_error, + } + + return results + diff --git a/src/cdtools/tools/interactions/interactions.py b/src/cdtools/tools/interactions/interactions.py index 0bff167..3cd244c 100644 --- a/src/cdtools/tools/interactions/interactions.py +++ b/src/cdtools/tools/interactions/interactions.py @@ -389,7 +389,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True): return t.stack(exit_waves) -def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, polarized=False, polarizer=None, analyzer=None): +def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, probe_support=None, polarized=False, polarizer=None, analyzer=None): """Returns a stack of exit waves accounting for subpixel shifts This function returns a collection of exit waves, with the first @@ -473,6 +473,10 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi shifted_fft_probe = fft_probe * phase_masks shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe, dim=(-1,-2))) + + if probe_support is not None: + shifted_probe = shifted_probe * probe_support[..., :, :] + if not polarized: # TODO This is a kludge, I will fix this. I need to handle # multiple incoherently mixing polarized objects @@ -620,6 +624,7 @@ def RPI_interaction(probe, obj): # The far-field propagator is just a 2D FFT but with an fftshift fftobj = propagators.far_field(obj) + fftobj_npix = fftobj.shape[-2] * fftobj.shape[-1] # We calculate the padding that we need to do the upsampling # This is carefully set up to keep the zero-frequency pixel in the correct # location as the overall shape changes. Don't mess with this without @@ -630,10 +635,12 @@ def RPI_interaction(probe, obj): pad1r = probe.shape[-1] - obj.shape[-1] - pad1l fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad2l, pad2r)) - + fftobj_npix_new = fftobj.shape[-2] * fftobj.shape[-1] + scale_factor =np.sqrt(fftobj_npix_new / fftobj_npix) + # Again, just an inverse FFT but with an fftshift - upsampled_obj = propagators.inverse_far_field(fftobj) - + upsampled_obj = scale_factor * propagators.inverse_far_field(fftobj) + if obj.dim() >= 3: return probe[None,...] * upsampled_obj else: diff --git a/src/cdtools/tools/measurements/measurements.py b/src/cdtools/tools/measurements/measurements.py index 87446f7..e726149 100644 --- a/src/cdtools/tools/measurements/measurements.py +++ b/src/cdtools/tools/measurements/measurements.py @@ -140,7 +140,7 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non if wavefields.dim() == 3: output = avg_pool2d(output.unsqueeze(0), oversampling)[0] else: - output = avg_pool2d(output, oversampling) + output = avg_pool2d(output, int(oversampling)) # Then we grab the detector slice if detector_slice is not None: