Add wavelength info to the data saved by fancy_ptycho

src/cdtools/models/fancy_ptycho.py

@@ -12,6 +12,9 @@ from copy import copy
 
 __all__ = ['FancyPtycho']
 
+# TODO The information on bases is all wrong when the probe_fourier_crop is
+# used, this affects display but also saveout of data and maybe even stuff
+# like the probe translation simulation
 
 class FancyPtycho(CDIModel):
 
@@ -151,6 +154,7 @@ class FancyPtycho(CDIModel):
                      translation_scale=1,
                      saturation=None,
                      probe_support_radius=None,
+                     probe_fourier_crop=None,
                      propagation_distance=None,
                      scattering_mode=None,
                      oversampling=1,
@@ -230,6 +234,12 @@ class FancyPtycho(CDIModel):
         else:
             probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
 
+        if probe_fourier_crop is not None:
+            probe = tools.propagators.far_field(probe)
+            probe = probe[probe_fourier_crop:-probe_fourier_crop,
+                          probe_fourier_crop:-probe_fourier_crop]
+            probe = tools.propagators.inverse_far_field(probe)
+            
         # Now we initialize all the subdominant probe modes
         probe_max = t.max(t.abs(probe))
         probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
@@ -361,7 +371,20 @@ class FancyPtycho(CDIModel):
                                      self.J_phase[None,...]))
             prs = prs * probe_masks[...,None,:,:]
 
+        # We automatically rescale the probe to match the background size,
+        # which allows us to do stuff like let the object be super-resolution,
+        # while restricting the probe to the detector resolution but still
+        # doing an explicit real-space limitation of the probe
+        padding = [self.background.shape[-2] - prs.shape[-2],
+                   self.background.shape[-1] - prs.shape[-1]]
 
+        if any([p != 0 for p in padding]): # For probe_fourier_crop != 0.
+            padding = [padding[-1]//2, padding[-1]-padding[-1]//2,
+                       padding[-2]//2, padding[-2]-padding[-2]//2]
+            prs = tools.propagators.far_field(prs)
+            prs = t.nn.functional.pad(prs, padding)
+            prs = tools.propagators.inverse_far_field(prs)
+            
         # 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(
@@ -651,9 +674,10 @@ class FancyPtycho(CDIModel):
         background = self.background.detach().cpu().numpy()**2
         weights = self.weights.detach().cpu().numpy()
         oversampling = self.oversampling
+        wavelength = self.wavelength.cpu().numpy()
 
         return {'basis': basis, 'translation': translations,
                 'probe': probe, 'obj': obj,
                 'background': background,
                 'oversampling': oversampling,
-                'weights': weights}
+                'weights': weights, 'wavelength': wavelength}

src/cdtools/tools/analysis/analysis.py

@@ -514,7 +514,12 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
     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))
     cor_fft = f1 * t.conj(f2)
-
+    #from cdtools.tools import plotting as p
+    #from matplotlib import pyplot as plt
+    #p.plot_amplitude(t.log(t.abs(cor_fft)))
+    #p.plot_amplitude(t.log(t.abs(f1)))
+    #p.plot_phase(cor_fft)
+    #plt.show()
     F1 = t.abs(f1)**2
     F2 = t.abs(f2)**2