Refactor the PRTF calculation, fix a small bug creating a slight offset in the frequency labels, and write a test

This commit is contained in:
Abe Levitan
2019-04-25 13:32:37 -04:00
parent 9a35d44689
commit bc9c1c466c
3 changed files with 121 additions and 52 deletions
+17 -51
View File
@@ -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()
+73 -1
View File
@@ -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)
+31
View File
@@ -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))