diff --git a/CDTools/scripts/synthesize.py b/CDTools/scripts/synthesize.py index 34bbbe1..e876f8b 100644 --- a/CDTools/scripts/synthesize.py +++ b/CDTools/scripts/synthesize.py @@ -4,50 +4,11 @@ import torch as t from matplotlib import pyplot as plt import pickle import argparse -from scipy import fftpack from CDTools.tools import cmath, plotting -from CDTools.tools import image_processing as ip from CDTools.tools.analysis import * - - -def calc_prtf(synth_obj, objects, basis, 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] - - synth_obj = cmath.complex_to_torch(synth_obj[obj_slice]) - - synth_fft = cmath.cabssq(cmath.fftshift(t.fft(synth_obj,2))).numpy() - - prtfs = [] - for obj in objects: - obj = cmath.complex_to_torch(obj[obj_slice]) - single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy() - - - di = np.linalg.norm(basis[:,0]) - dj = np.linalg.norm(basis[:,1]) - - i_freqs = fftpack.fftshift(fftpack.fftfreq(single_fft.shape[0],d=di)) - j_freqs = fftpack.fftshift(fftpack.fftfreq(single_fft.shape[1],d=dj)) - - Js,Is = np.meshgrid(j_freqs,i_freqs) - Is = Is - np.mean(Is) - Js = Js - np.mean(Js) - Rs = np.sqrt(Is**2+Js**2) - - single_ints, bins = np.histogram(Rs,bins=100,weights=single_fft) - synth_ints, bins = np.histogram(Rs,bins=100,weights=synth_fft) - - prtfs.append(synth_ints/single_ints) - - return bins[:-1], np.mean(prtfs,axis=0) - - - def make_argparser(): parser = argparse.ArgumentParser(description=__doc__) @@ -56,6 +17,7 @@ def make_argparser(): return parser + if __name__ == '__main__': args = make_argparser().parse_args() @@ -66,7 +28,7 @@ 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']) + freqs, prtf = calc_consistency_prtf(synth_obj, aligned_objs, dataset['basis']) plotting.plot_phase(dataset['probe'][0][0],basis=1e6*dataset['basis']) @@ -74,10 +36,12 @@ if __name__ == '__main__': plotting.plot_colorized(dataset['probe'][0][0],basis=1e6*dataset['basis']) - plotting.plot_phase(synth_probe[1],basis=1e6*dataset['basis']) - plotting.plot_amplitude(synth_probe[1],basis=1e6*dataset['basis']) - plotting.plot_colorized(synth_probe[1],basis=1e6*dataset['basis']) - + try: + plotting.plot_phase(synth_probe[1],basis=1e6*dataset['basis']) + plotting.plot_amplitude(synth_probe[1],basis=1e6*dataset['basis']) + plotting.plot_colorized(synth_probe[1],basis=1e6*dataset['basis']) + except: + pass plotting.plot_amplitude(synth_obj,basis=1e6*dataset['basis']) @@ -86,12 +50,14 @@ if __name__ == '__main__': plt.figure() - plt.show() - exit() - real_translations = dataset['basis'].dot(dataset['translation'][0].transpose()) - real_translations -= np.min(real_translations,axis=1)[:,None] - plt.plot(real_translations[0]*1e6,real_translations[1]*1e6,'k.') - plt.plot(real_translations[0]*1e6,real_translations[1]*1e6,'b-',linewidth=0.5) - plt.figure() + try: + real_translations = dataset['basis'].dot(dataset['translation'][0].transpose()) + real_translations -= np.min(real_translations,axis=1)[:,None] + plt.plot(real_translations[0]*1e6,real_translations[1]*1e6,'k.') + plt.plot(real_translations[0]*1e6,real_translations[1]*1e6,'b-',linewidth=0.5) + plt.figure() + except: + pass + plt.plot(freqs*1e-6, prtf) plt.show() diff --git a/CDTools/tools/analysis.py b/CDTools/tools/analysis.py index 76d0866..1e84def 100644 --- a/CDTools/tools/analysis.py +++ b/CDTools/tools/analysis.py @@ -4,8 +4,11 @@ import torch as t import numpy as np from CDTools.tools import cmath from CDTools.tools import image_processing as ip +from scipy import fftpack + +__all__ = ['orthogonalize_probes','standardize', 'synthesize_reconstructions', + 'calc_consistency_prtf'] -__all__ = ['orthogonalize_probes','standardize', 'synthesize_reconstructions'] from matplotlib import pyplot as plt def orthogonalize_probes(probes): @@ -240,3 +243,72 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None, return synth_probe/(i+2), synth_obj/(i+2), obj_stack + + +def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None): + """Calculates a PRTF between each the individual objects and an averaged one + + The consistency PRTF at any given spatial frequency is defined as the ratio + between the intensity of any given reconstruction and the intensity + of a synthesized or averaged reconstruction at that spatial frequency. + Typically, the PRTF is averaged over spatial frequencies with the same + magnitude. + + Args: + synth_obj (t.Tensor) : The synthesized object in the numerator of the PRTF + objects (list): A list of objects or diffraction patterns for the denomenator of the PRTF + basis (array_like) : The basis for the reconstruction array to allow output in physical unit + obj_slice : Optional, a slice of the objects to use for calculating the PRTF + nbinbs (int) : Optional, number of bins to use in the histogram. Defaults to a sensible value + + Returns: + (t.Tensor) : The frequencies for the PRTF + (t.Tensor) : The values of the PRTF + """ + + obj_np = False + if isinstance(objects[0], np.ndarray): + objects = [cmath.complex_to_torch(obj).to(t.float32) for obj in objects] + obj_np = True + if isinstance(synth_obj, np.ndarray): + synth_obj = cmath.complex_to_torch(synth_obj).to(t.float32) + + + 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] + + if nbins is None: + nbins = np.max(synth_obj[obj_slice].shape) // 4 + + synth_fft = cmath.cabssq(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy() + + + di = np.linalg.norm(basis[:,0]) + dj = np.linalg.norm(basis[:,1]) + + i_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[0],d=di)) + j_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[1],d=dj)) + + Js,Is = np.meshgrid(j_freqs,i_freqs) + Rs = np.sqrt(Is**2+Js**2) + + + synth_ints, bins = np.histogram(Rs,bins=nbins,weights=synth_fft) + + prtfs = [] + for obj in objects: + obj = obj[obj_slice] + single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy() + single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft) + + prtfs.append(synth_ints/single_ints) + + + if not obj_np: + bins = t.Tensor(bins) + prtfs = t.Tensor(prtfs) + + return bins[:-1], np.mean(prtfs,axis=0) + + diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index df2ef04..a1fb468 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -115,4 +115,35 @@ def test_standardize(): from matplotlib import pyplot as plt def test_synthesize_reconstructions(): + # Not really sure how to test this to be honest + + # Perhaps it's just best to test for a lack of failures? pass + + +def test_calc_consistency_prtf(): + + # Create an object with a specific structure + obj = 30 * np.random.rand(1030,1040) * np.exp(1j * (np.random.rand(1030,1040) - 0.5)) + + # + synth_obj = np.sqrt(0.7) * obj + + obj_stack = [obj]#,obj,obj,obj] + + basis = np.array([[0,2,0], + [3,0,0]]) + + freqs, prtf = analysis.calc_consistency_prtf(synth_obj, obj_stack, basis) + assert np.allclose(prtf, 0.7) + + freqs, prtf = analysis.calc_consistency_prtf(synth_obj, obj_stack, basis, nbins=30) + assert np.allclose(prtf, 0.7) + + # Check that is uses the right number of bins + assert len(prtf) == 30 + assert len(freqs) == 30 + + # Check that the maximum frequency is correct for the basis + assert np.isclose(freqs[-1]-freqs[-2] + freqs[-1], np.sqrt(1/4**2 + 1/6**2)) +