From fdd100cd1c2877a67eaf6be957c84740e54c41f1 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 4 Apr 2019 17:03:23 -0400 Subject: [PATCH] Finish building and testing the gaussian probe guess from a dataset --- CDTools/tools/initializers.py | 30 +++++++-------- tests/tools/test_initializers.py | 63 ++++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 17 deletions(-) diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py index 24a6ca6..9e8aa5e 100644 --- a/CDTools/tools/initializers.py +++ b/CDTools/tools/initializers.py @@ -2,7 +2,8 @@ from __future__ import division, print_function, absolute_import import numpy as np import torch as t -__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian'] +__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian', + 'gaussian_probe'] from CDTools.tools import cmath from scipy.fftpack import next_fast_len @@ -10,7 +11,7 @@ import numpy as np def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0): - """Returns an exit wave basis and a detector slice for the given detector geometry + """Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry It takes in the parameters for a given detector - the basis defining the pixel pitch and the shape, as well as the wavelength and propagation @@ -137,7 +138,8 @@ def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]): return cmath.complex_to_torch(amplitude*result) -def gaussian_initialization(dataset, basis, shape, sigma, propagation_distance=0): + +def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0): """Initializes a gaussian probe based on experimental parameters This function generates a gaussian probe initialization which has a @@ -169,32 +171,26 @@ def gaussian_initialization(dataset, basis, shape, sigma, propagation_distance=0 wavelength = dataset.wavelength z = propagation_distance # for shorthand sigma = np.array(sigma) - curvature = np.array(curvature) k = 2 * np.pi / wavelength zr = k * sigma**2 sigmaz = sigma * np.sqrt(1 + (z / zr)**2) - curvature = k * z / (z**2 + zr**2) + curvature = -k * z / (z**2 + zr**2) - # So both sigmaz and curvature can be either scalars or tensors here - # We make them consistent - if len(sigmaz.shape) == 0: - sigmaz = np.array([sigmaz, sigmaz]) - if len(curvature.shape) == 0: - curvature = np.array([curvature, curvature]) - # The conversion must then be done to pixel space sigma_pix = sigmaz / np.array([np.linalg.norm(basis[:,0]), np.linalg.norm(basis[:,1])]) - curvature_pix = sigmaz * np.array([np.linalg.norm(basis[:,0]), - np.linalg.norm(basis[:,1])])**2 - + curvature_pix = curvature * np.array([np.linalg.norm(basis[:,0]), + np.linalg.norm(basis[:,1])])**2 + # Then we can generate the gaussian array probe = gaussian(shape, sigma=sigma_pix, curvature=curvature_pix) # Finally, we should calculate the average pattern intensity from the # dataset and normalize the gaussian probe. This should be done by - avg_intensity = 1 - #probe_intensity = + avg_intensities = [t.sum(dataset[idx][1]) for idx in range(len(dataset))] + avg_intensity = t.mean(t.Tensor(avg_intensities)) + probe_intensity = t.sum(cmath.cabssq(probe)) + return avg_intensity / probe_intensity * probe diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index 9d06d86..69a3eca 100644 --- a/tests/tools/test_initializers.py +++ b/tests/tools/test_initializers.py @@ -2,6 +2,7 @@ from __future__ import division, print_function, absolute_import from CDTools.tools import initializers from CDTools.tools import cmath +from CDTools.datasets import Ptycho_2D_Dataset import numpy as np import torch as t @@ -108,3 +109,65 @@ def test_gaussian(): center=center, curvature=curvature, amplitude=10)) assert np.allclose(init_result, np_result) + +def test_gaussian_probe(ptycho_cxi_1): + + dataset = Ptycho_2D_Dataset.from_cxi(ptycho_cxi_1[0]) + + det_basis = t.Tensor(dataset.detector_geometry['basis']) + det_shape = t.Size(dataset.patterns.shape[-2:]) + wavelength = dataset.wavelength + distance = dataset.detector_geometry['distance'] + + basis, shape, s = initializers.exit_wave_geometry(det_basis, + det_shape, + wavelength, + distance) + # Basis is around 60nm in the i(y) direction, 85nm in the j(x) direction + # Full window is therefore about 15 um in i(y) and 20 um in the j(x) dir + + # Come up with a roughly matching set of probe parameters + sigma = 5e-7 + + # Build a stage explicitly with numpy to compare against + x = (np.arange(256) - 127.5) * (-basis[0,1]).numpy() + y = (np.arange(256) - 127.5) * (-basis[1,0]).numpy() + Xs,Ys = np.meshgrid(x,y) + Rs = np.sqrt(Xs**2+Ys**2) + + + + # Now we first test the non-propagated probe + np_probe = np.exp(-1/(2*sigma**2) * Rs**2) + + normalization = 0 + for params, im in dataset: + normalization += np.sum(im.cpu().numpy()) + normalization /= len(dataset) + + normalization_1 = normalization / np.sum(np.abs(np_probe)**2) + + probe = initializers.gaussian_probe(dataset, basis, shape, sigma) + probe = cmath.torch_to_complex(probe) + assert np.allclose(probe, normalization_1*np_probe) + + # And then a propagated probe + z = 1e-4 #nm + k = 2 * np.pi / wavelength + w0 = np.sqrt(2)*sigma + zr = np.pi * w0**2 / wavelength + wz = w0 * np.sqrt(1 + (z / zr)**2) + Rz = z * (1 + (zr / z)**2) + np_probe = np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz)) + + normalization_2 = normalization / np.sum(np.abs(np_probe)**2) + + probe = initializers.gaussian_probe(dataset, basis, shape, sigma, + propagation_distance=z) + probe = cmath.torch_to_complex(probe) + + assert np.allclose(probe, normalization_2*np_probe) + + + +