From 22899288fc6bd6086cf8d8bcc050b4f7e089945a Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Fri, 3 May 2019 12:30:53 -0400 Subject: [PATCH] Add probe propagation with initialization as a base feature to fancy_ptycho, and fix the test for orthogonalize_probes --- CDTools/models/fancy_ptycho.py | 8 ++++---- CDTools/scripts/synthesize.py | 8 +++++--- CDTools/tools/analysis.py | 10 ++++++---- CDTools/tools/initializers.py | 34 ++++++++++++++++++++++++++++---- CDTools/tools/plotting.py | 13 +++++++++--- tests/tools/test_analysis.py | 17 +++++++++------- tests/tools/test_initializers.py | 6 +++++- 7 files changed, 70 insertions(+), 26 deletions(-) diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 4247759..f5ba210 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/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 diff --git a/CDTools/scripts/synthesize.py b/CDTools/scripts/synthesize.py index 15ada9b..d16cdc2 100644 --- a/CDTools/scripts/synthesize.py +++ b/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']) diff --git a/CDTools/tools/analysis.py b/CDTools/tools/analysis.py index 05a002a..ae8df78 100644 --- a/CDTools/tools/analysis.py +++ b/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]) diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py index 0acaed9..111778c 100644 --- a/CDTools/tools/initializers.py +++ b/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 diff --git a/CDTools/tools/plotting.py b/CDTools/tools/plotting.py index 3d6297d..f6c2639 100644 --- a/CDTools/tools/plotting.py +++ b/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 + ')') diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index 926e452..131c3e7 100644 --- a/tests/tools/test_analysis.py +++ b/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) diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index 37cf541..fe13990 100644 --- a/tests/tools/test_initializers.py +++ b/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]) + +