Add probe propagation with initialization as a base feature to fancy_ptycho, and fix the test for orthogonalize_probes

CDTools/models/fancy_ptycho.py

@@ -79,7 +79,7 @@ class FancyPtycho(CDIModel):
         
         
     @classmethod
-    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None):
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None):
         wavelength = dataset.wavelength
         det_basis = dataset.detector_geometry['basis']
         det_shape = dataset[0][1].shape
@@ -110,16 +110,16 @@ class FancyPtycho(CDIModel):
         
         obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=50)
 
-        if hasattr(dataset, 'background'):
+        if hasattr(dataset, 'background') and dataset.background is not None:
             background = t.sqrt(dataset.background)
         else:
             background = None
 
         # Finally, initialize the probe and  object using this information
         if probe_size is None:
-            probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice)
+            probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance)
         else:
-            probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size)
+            probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
 
 
         # Now we initialize all the subdominant probe modes

CDTools/scripts/synthesize.py

@@ -27,12 +27,14 @@ if __name__ == '__main__':
 
     synth_probe, synth_obj, aligned_objs = synthesize_reconstructions(
         dataset['probe'], dataset['obj'], args.use_probe)
+    print(dataset['basis'])
 
+    print(synth_probe.shape)
     freqs, prtf = calc_consistency_prtf(synth_obj, aligned_objs, dataset['basis'])
     
-    plotting.plot_phase(dataset['probe'][0][0]),basis=dataset['basis'])
-    plotting.plot_amplitude(dataset['probe'][0][0]),basis=dataset['basis'])
-    plotting.plot_colorized(dataset['probe'][0][0]),basis=dataset['basis'])
+    plotting.plot_phase(synth_probe,basis=dataset['basis'])
+    plotting.plot_amplitude(synth_probe,basis=dataset['basis'])
+    plotting.plot_colorized(synth_probe,basis=dataset['basis'])
     
     try:
         plotting.plot_phase(synth_probe[1],basis=dataset['basis'])

CDTools/tools/analysis.py

@@ -28,8 +28,9 @@ def orthogonalize_probes(probes):
 
     try:
         probes = cmath.torch_to_complex(probes.detach().cpu())
+        send_to_torch = True
     except:
-        pass
+        send_to_torch = False
 
     bases = []
     coefficients = np.zeros((probes.shape[0],probes.shape[0]), dtype=np.complex64)
@@ -50,15 +51,16 @@ def orthogonalize_probes(probes):
     ortho_probes = []
     for i in range(len(eigvals)):
         coefficients = np.sqrt(eigvals[i]) * eigvecs[:,i]
-        print(coefficients)
         probe = np.zeros(bases[0].shape, dtype=np.complex64)
         for coefficient, basis in zip(coefficients, bases):
             probe += basis * coefficient
         ortho_probes.append(probe)
 
-        
-    return cmath.complex_to_torch(np.stack(ortho_probes[::-1]))
     
+    if send_to_torch:
+        return cmath.complex_to_torch(np.stack(ortho_probes[::-1]))
+    else:
+        return np.stack(ortho_probes[::-1])
         
 
 

CDTools/tools/initializers.py

@@ -6,7 +6,7 @@ __all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
            'gaussian_probe', 'SHARP_style_probe'] 
 
 from CDTools.tools import cmath
-from CDTools.tools.propagators import inverse_far_field
+from CDTools.tools.propagators import inverse_far_field, generate_angular_spectrum_propagator, near_field
 from scipy.fftpack import next_fast_len
 import numpy as np
 
@@ -211,7 +211,7 @@ def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0):
     return avg_intensity / probe_intensity * probe
     
 
-def SHARP_style_probe(dataset, shape, det_slice):
+def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None):
     """Generates a SHARP style probe guess from a dataset
 
     What we call the "SHARP" style probe guess is to take a mean of all
@@ -233,6 +233,7 @@ def SHARP_style_probe(dataset, shape, det_slice):
         dataset (Ptycho_2D_Dataset) : The dataset to work from
         shape (torch.Size) : The size of the probe array to simulate
         det_slice (slice) : A slice or tuple of slices corresponding to the detector region in Fourier space
+        propagatioin_distance (float) : Default is no propagation, an amount to propagate the guessed probe from it's focal point
     """
 
     
@@ -242,7 +243,7 @@ def SHARP_style_probe(dataset, shape, det_slice):
     intensities /= len(dataset)
 
     # Subtract off a known background if it's stored
-    if hasattr(dataset, 'background'):
+    if hasattr(dataset, 'background') and dataset.background is not None:
         intensities[det_slice] = np.clip(intensities[det_slice] - dataset.background.cpu().numpy(), a_min=0,a_max=None)
     
     probe_fft = cmath.complex_to_torch(np.sqrt(intensities))
@@ -262,7 +263,32 @@ def SHARP_style_probe(dataset, shape, det_slice):
         probe_guess[center[0]+1, center[1]],
         probe_guess[center[0], center[1]-1],
         probe_guess[center[0], center[1]+1]])
+
+
+    probe_guess = cmath.complex_to_torch(probe_guess)
     
-    return cmath.complex_to_torch(probe_guess)
+    if propagation_distance is not None:
+        # First generate the propagation array
+
+        probe_shape = t.Tensor(tuple(shape))
+
+        # Start by recalculating the probe basis from the given information
+        det_basis = t.Tensor(dataset.detector_geometry['basis'])
+        basis_dirs = det_basis / t.norm(det_basis, dim=0)
+        distance = dataset.detector_geometry['distance']
+        probe_basis = basis_dirs * dataset.wavelength * distance / \
+            (probe_shape * t.norm(det_basis,dim=0))
+
+        # Then package everything as it's needed
+        probe_spacing = t.norm(probe_basis,dim=0).numpy()
+        probe_shape = probe_shape.numpy().astype(np.int32)
+
+        #assert 0
+        # And generate the propagator
+        AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance)
+
+        probe_guess = near_field(probe_guess,AS_prop)
+    
+    return probe_guess
 
     

CDTools/tools/plotting.py

@@ -81,6 +81,7 @@ def plot_1D(arr, fig = None, **kwargs):
         ax = fig.add_subplot(111, **kwargs)
     else:
         plt.figure(fig.number)
+        plt.gcf().clear()
 
     plt.scatter(np.arange(arr.shape[-1]), arr)
 
@@ -101,6 +102,7 @@ def plot_amplitude(im, fig = None, basis=None, units='um', **kwargs):
         ax = fig.add_subplot(111, **kwargs)
     else:
         plt.figure(fig.number)
+        plt.gcf().clear()
 
     if isinstance(im, t.Tensor):
         absolute = cmath.cabs(im).detach().cpu().numpy()
@@ -145,6 +147,7 @@ def plot_phase(im, fig=None, basis=None, units='um', **kwargs):
         ax = fig.add_subplot(111, **kwargs)
     else:
         plt.figure(fig.number)
+        plt.gcf().clear()
 
     # If the user has matplotlib >=3.0, use the preferred colormap
     if isinstance(im, t.Tensor):
@@ -198,6 +201,7 @@ def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
         ax = fig.add_subplot(111, **kwargs)
     else:
         plt.figure(fig.number)
+        plt.gcf().clear()
         
     if isinstance(im, t.Tensor):
         im = cmath.torch_to_complex(im.detach().cpu())
@@ -226,16 +230,17 @@ def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
 
 
 
-def plot_translations(translations, fig=None, units='um'):
+def plot_translations(translations, fig=None, units='um', lines=True):
     """Plots a set of probe translations in a nicely formatted way
     
     Args:
         translations: An Nx2 or Nx3 set of translations in real space
         fig : Optional, a figure to plot into
         units : Default is um, units to report in (assuming input in m)
+        lines : Whether to plot the lines indicating the path
 
     Returns:
-        
+        None
     """
     
     factor = get_units_factor(units)
@@ -245,13 +250,15 @@ def plot_translations(translations, fig=None, units='um'):
         ax = fig.add_subplot(111)
     else:
         plt.figure(fig.number)
+        plt.gcf().clear()
 
     if isinstance(translations, t.Tensor):
         translations = translations.detach().cpu().numpy()
         
     translations = translations * factor
     plt.plot(translations[:,0], translations[:,1],'k.')
-    plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
+    if lines:
+        plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
     plt.xlabel('X (' + units + ')')
     plt.ylabel('Y (' + units + ')')
 

tests/tools/test_analysis.py

@@ -25,21 +25,24 @@ def test_orthogonalize_probes():
                        3*np.exp(-probe_Rs**2 / (2 * 12**2)),
                        1*np.exp(-probe_Rs**2 / (2 * 15**2))]).astype(np.complex64)
 
-    ortho_probes = cmath.torch_to_complex(analysis.orthogonalize_probes(probes))
-    
+    # test that it works on numpy arrays
+    ortho_probes = analysis.orthogonalize_probes(probes)
+
+    # test that it also works on torch tensors
+    ortho_probes_t = cmath.torch_to_complex(analysis.orthogonalize_probes(cmath.complex_to_torch(probes)))
+        
     for p1,p2 in combinations(ortho_probes,2):
         assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6
 
-    for probe in ortho_probes:
-        print(np.sum(np.abs(probe)**2))
+    for p1,p2 in combinations(ortho_probes_t,2):
+        assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6
 
-    for probe in probes:
-        print(np.sum(np.abs(probe)**2))
-        
     probe_intensity = np.sum(np.abs(probes)**2,axis=0)
     ortho_probe_intensity = np.sum(np.abs(ortho_probes)**2,axis=0)
+    ortho_probe_t_intensity = np.sum(np.abs(ortho_probes_t)**2,axis=0)
 
     assert np.allclose(probe_intensity,ortho_probe_intensity)
+    assert np.allclose(probe_intensity,ortho_probe_t_intensity)
 
 
 

tests/tools/test_initializers.py

@@ -186,4 +186,8 @@ def test_SHARP_style_probe(ptycho_cxi_1):
 
     probe = initializers.SHARP_style_probe(dataset, shape, det_slice)
     assert probe.shape == t.Size([256,256,2])
-    
+
+    probe = initializers.SHARP_style_probe(dataset, shape, det_slice, propagation_distance=20e-6)
+    assert probe.shape == t.Size([256,256,2])
+
+