Add a phase only constraint to RPI and fancy_ptycho, misc small changes

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',

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),
     ]
 
 

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
+

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:

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: