diff --git a/CDTools/scripts/synthesize.py b/CDTools/scripts/synthesize.py index 03e87dd..6e8d2aa 100644 --- a/CDTools/scripts/synthesize.py +++ b/CDTools/scripts/synthesize.py @@ -8,66 +8,23 @@ from scipy import fftpack from CDTools.tools import cmath, plotting from CDTools.tools import image_processing as ip +from CDTools.tools.analysis import * -def standardize(probe, obj, obj_slice=None): - # First, we normalize the probe intensity to a fixed value. - # Should this be the maximum or the integrated intensity? I think - # probably the integrated intensity. We set the average per-[ixel - # intensity in the probe to be one - normalization = np.sqrt(np.sum(np.abs(probe)**2) / len(probe.ravel())) - probe = cmath.complex_to_torch(probe / normalization) - obj = cmath.complex_to_torch(obj * normalization) - - - # Default slice of the object to use for alignment, etc. - if obj_slice is None: - obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5, - (obj.shape[1]//8)*3:(obj.shape[1]//8)*5] - - - - # Now we get rid of the probe's phase ramp - # Currently disabled - #center_freq = ip.centroid_sq(cmath.fftshift(t.fft(probe,2)),comp=True) - #center_freq -= (t.tensor(probe.shape[:-1]) // 2).to(t.float32) - #center_freq /= t.tensor(probe.shape[:-1]).to(t.float32) - - - - #Is, Js = np.mgrid[:probe.shape[0],:probe.shape[1]] - #probe_phase_ramp = cmath.expi(2*np.pi * - # (center_freq[0] * t.tensor(Is).to(t.float32) + - # center_freq[1] * t.tensor(Js).to(t.float32))) - #probe = cmath.cmult(probe, cmath.cconj(probe_phase_ramp)) - #Is, Js = np.mgrid[:obj.shape[0],:obj.shape[1]] - #obj_phase_ramp = cmath.expi(2*np.pi * - # (center_freq[0] * t.tensor(Is).to(t.float32) + - # center_freq[1] * t.tensor(Js).to(t.float32))) - #obj = cmath.cmult(obj, obj_phase_ramp) - - - # Then, we set them to consistent absolute phases - probe_angle = cmath.cphase(t.sum(probe,dim=(0,1))) - obj_angle = cmath.cphase(t.sum(obj[obj_slice],dim=(0,1))) - - probe = cmath.cmult(probe, cmath.expi(-probe_angle)) - obj = cmath.cmult(obj, cmath.expi(-obj_angle)) - - return probe, obj - - def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None): if obj_slice is None: obj_slice = np.s_[(objects[0].shape[0]//8)*3:(objects[0].shape[0]//8)*5, (objects[0].shape[1]//8)*3:(objects[0].shape[1]//8)*5] + probes = [cmath.complex_to_torch(probe).to(t.float32) for probe in probes] + objects = [cmath.complex_to_torch(obj).to(t.float32) for obj in objects] synth_probe, synth_obj = standardize(probes[0], objects[0]) obj_stack = [cmath.torch_to_complex(synth_obj)] for i, (probe, obj) in enumerate(zip(probes[1:],objects[1:])): probe, obj = standardize(probe, obj) + probe = probe[0] print(i) #plt.imshow(np.angle(cmath.torch_to_complex(obj[obj_slice]))) @@ -154,19 +111,17 @@ if __name__ == '__main__': synth_probe, synth_obj, aligned_objs = synthesize_reconstructions( dataset['probe'], dataset['obj'], args.use_probe) - + freqs, prtf = calc_prtf(synth_obj, aligned_objs, dataset['basis']) print(np.linalg.norm(dataset['basis'],axis=0)) plotting.plot_phase(dataset['probe'][0][0],basis=1e6*dataset['basis']) plotting.plot_amplitude(dataset['probe'][0][0],basis=1e6*dataset['basis']) plotting.plot_colorized(dataset['probe'][0][0],basis=1e6*dataset['basis']) - #plotting.plot_amplitude(synth_obj[400:750,450:850],basis=1e6*dataset['basis']) - #plotting.plot_colorized(synth_obj[400:750,450:850],basis=1e6*dataset['basis']) - #plotting.plot_phase(synth_obj[400:750,450:850],basis=1e6*dataset['basis']) - plotting.plot_amplitude(synth_obj[::-1,::-1][450:900,325:775],basis=1e6*dataset['basis']) - plotting.plot_phase(synth_obj[::-1,::-1][450:900,325:775],basis=1e6*dataset['basis']) - plotting.plot_colorized(synth_obj[::-1,::-1][450:900,325:775],basis=1e6*dataset['basis']) + plotting.plot_amplitude(synth_obj,basis=1e6*dataset['basis']) + plotting.plot_colorized(synth_obj,basis=1e6*dataset['basis']) + plotting.plot_phase(synth_obj,basis=1e6*dataset['basis']) + plt.figure() real_translations = dataset['basis'].dot(dataset['translation'][0].transpose()) diff --git a/CDTools/tools/analysis.py b/CDTools/tools/analysis.py index 09ae5f4..bb870f0 100644 --- a/CDTools/tools/analysis.py +++ b/CDTools/tools/analysis.py @@ -3,8 +3,9 @@ from __future__ import division, print_function import torch as t import numpy as np from CDTools.tools import cmath +from CDTools.tools import image_processing as ip -__all__ = ['orthogonalize_probes'] +__all__ = ['orthogonalize_probes','standardize'] from matplotlib import pyplot as plt def orthogonalize_probes(probes): @@ -56,4 +57,86 @@ def orthogonalize_probes(probes): return cmath.complex_to_torch(np.stack(ortho_probes[::-1])) - + + +def standardize(probe, obj, obj_slice=None, correct_ramp=False): + """Standardizes a probe and object to prepare them for comparison + + There are a number of ambiguities in the definition of a ptychographic + reconstruction. This function makes an explicit choice for each ambiguity + to allow comparisons between independent reconstructions without confusing + these ambiguities for real differences between the reconstructions. + + The ambiguities and standardizations are: + * Probe and object can be scaled inversely to one another + * So we set the probe intensity to an average per-pixel value of 1 + * The probe and object can aquire equal and opposite phase ramps + * So we set the centroid of the FFT of the probe to zero frequency + * The probe and object can each acquire an arbitrary overall phase + * So we set the phase of the sum of all values of both the probe and object to 0 + + When dealing with the properties of the object, a slice is used by + default as the edges of the object often are dominated by unphysical + noise. The default slice is from 3/8 to 5/8 of the way across. + + Args: + probe (t.tensor) : tensor or numpy array storing a retrieved probe + obj (t.tensor) : tensor or numpy array storing a retrieved probe + obj_slice (slice) : optional, a slice to take from the object for calculating normalizations + correct_ramp (bool) : Default False, whether to correct for the relative phase ramps + + """ + # First, we normalize the probe intensity to a fixed value. + probe_np = False + if isinstance(probe, np.ndarray): + probe = cmath.complex_to_torch(probe).to(t.float32) + probe_np = True + obj_np = False + if isinstance(obj, np.ndarray): + obj = cmath.complex_to_torch(obj).to(t.float32) + obj_np = True + + normalization = t.sqrt(t.sum(cmath.cabssq(probe)) / (len(probe.view(-1))/2)) + probe = probe / normalization + obj = obj * normalization + + # Default slice of the object to use for alignment, etc. + if obj_slice is None: + obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5, + (obj.shape[1]//8)*3:(obj.shape[1]//8)*5] + + + if correct_ramp: + # Need to check if this is actually working and, if noy, why not + center_freq = ip.centroid_sq(cmath.fftshift(t.fft(probe,2)),comp=True) + center_freq -= (t.tensor(probe.shape[:-1]) // 2).to(t.float32) + center_freq /= t.tensor(probe.shape[:-1]).to(t.float32) + + + + Is, Js = np.mgrid[:probe.shape[0],:probe.shape[1]] + probe_phase_ramp = cmath.expi(2*np.pi * + (center_freq[0] * t.tensor(Is).to(t.float32) + + center_freq[1] * t.tensor(Js).to(t.float32))) + probe = cmath.cmult(probe, cmath.cconj(probe_phase_ramp)) + Is, Js = np.mgrid[:obj.shape[0],:obj.shape[1]] + obj_phase_ramp = cmath.expi(2*np.pi * + (center_freq[0] * t.tensor(Is).to(t.float32) + + center_freq[1] * t.tensor(Js).to(t.float32))) + obj = cmath.cmult(obj, obj_phase_ramp) + + + # Then, we set them to consistent absolute phases + probe_angle = cmath.cphase(t.sum(probe,dim=(0,1))) + obj_angle = cmath.cphase(t.sum(obj[obj_slice],dim=(0,1))) + + probe = cmath.cmult(probe, cmath.expi(-probe_angle)) + obj = cmath.cmult(obj, cmath.expi(-obj_angle)) + + if probe_np: + probe = cmath.torch_to_complex(probe.detach().cpu()) + if obj_np: + obj = cmath.torch_to_complex(obj.detach().cpu()) + + return probe, obj + diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py index 33268b1..1810f4e 100644 --- a/CDTools/tools/initializers.py +++ b/CDTools/tools/initializers.py @@ -122,7 +122,7 @@ def calc_object_setup(probe_shape, translations, padding=0): def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]): """Returns an array with a centered Gaussian - Takes in the shape, amplitude, and standard deviation of a gaussian + Takes in the shape and standard deviation of a gaussian and returns a complex torch tensor (trailing dimension is 2) with values corresponding to a two-dimensional gaussian function @@ -135,8 +135,8 @@ def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]): Args: shape (array_like) : A 1x2 array-like object specifying the dimensions of the output array in the form (i shape, j shape) - amplitude (float or int): The amplitude the gaussian to simulate sigma (array_like): A 1x2 array-like object specifying the i- and j- standard deviation of the gaussian in the form (i stdev, j stdev) + amplitude (float or int): Default 1, the amplitude the gaussian to simulate center (array_like) : Optional 1x2 array-like object specifying the location of the center of the gaussian (i center, j center) curvature (array_like) : Optional complex part to add to the gaussian coefficient diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index 9919b90..261da08 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -5,10 +5,9 @@ import numpy as np import torch as t from itertools import combinations -from CDTools.tools import analysis, cmath +from CDTools.tools import analysis, cmath, initializers -from matplotlib import pyplot as plt def test_orthogonalize_probes(): # The test strategy should be to define a few non-orthogonal probes @@ -41,3 +40,61 @@ def test_orthogonalize_probes(): ortho_probe_intensity = np.sum(np.abs(ortho_probes)**2,axis=0) assert np.allclose(probe_intensity,ortho_probe_intensity) + + +from matplotlib import pyplot as plt + +def test_standardize(): + + # Start by making a probe and object that should meet the standardization + # conditions + probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)) + probe = cmath.torch_to_complex(probe) + probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2)) + probe = probe * np.exp(-1j * np.angle(np.sum(probe))) + + assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel())) + assert np.isclose(0,np.angle(np.sum(probe))) + + obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5)) + obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5, + (obj.shape[1]//8)*3:(obj.shape[1]//8)*5] + + obj = obj * np.exp(-1j * np.angle(np.sum(obj[obj_slice]))) + assert np.isclose(0,np.angle(np.sum(obj[obj_slice]))) + + + # Then make a nonstandard version of them and standardize it + # First, don't add a phase ramp and test + test_probe = probe * 37.6 * np.exp(1j*0.35) + test_obj = obj / 37.6 * np.exp(1j*1.43) + s_probe, s_obj = analysis.standardize(test_probe, test_obj) + assert np.allclose(probe, s_probe) + assert np.allclose(obj, s_obj) + + # Test that it works on torch tensors + s_probe, s_obj = analysis.standardize(cmath.complex_to_torch(test_probe).to(t.float32), cmath.complex_to_torch(test_obj).to(t.float32)) + s_probe = cmath.torch_to_complex(s_probe) + s_obj = cmath.torch_to_complex(s_obj) + assert np.allclose(probe, s_probe) + assert np.allclose(obj, s_obj) + + + # And ensure that standardization maps back to the standard versions + phase_ramp_dir = (np.random.rand(2) - 0.5) + + probe_Xs, probe_Ys = np.mgrid[:probe.shape[0],:probe.shape[1]] + phase_ramp = np.exp(1j*probe_Ys * phase_ramp_dir[1]+ + 1j*probe_Xs * phase_ramp_dir[0]) + test_probe = test_probe * phase_ramp + + obj_Xs, obj_Ys = np.mgrid[:obj.shape[0],:obj.shape[1]] + phase_ramp = np.exp(-1j*obj_Ys * phase_ramp_dir[1]+ + -1j*obj_Xs * phase_ramp_dir[0]) + test_obj = test_obj * phase_ramp + + s_probe, s_obj = analysis.standardize(test_probe, test_obj, correct_ramp=True) + + assert np.max(s_probe - probe) / np.max(np.abs(probe)) < 1e-4 + assert np.max(s_obj - obj) / np.max(np.abs(obj)) < 1e-4 +