mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
Update the standardization to work with incoherent probes and move over the synthesis
This commit is contained in:
Abe Levitan
@@ -11,52 +11,6 @@ from CDTools.tools import image_processing as ip
|
||||
from CDTools.tools.analysis import *
|
||||
|
||||
|
||||
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])))
|
||||
#plt.show()
|
||||
|
||||
if use_probe:
|
||||
shift = ip.find_shift(synth_probe,probe, resolution=50)
|
||||
else:
|
||||
shift = ip.find_shift(synth_obj[obj_slice],obj[obj_slice], resolution=50)
|
||||
|
||||
|
||||
obj = ip.sinc_subpixel_shift(obj,np.array(shift))
|
||||
probe = ip.sinc_subpixel_shift(probe,tuple(shift))
|
||||
#obj = t.roll(obj,tuple(int(s) for s in shift),dims=(0,1))
|
||||
#probe = t.roll(probe,tuple(int(s) for s in shift),dims=(0,1))
|
||||
|
||||
synth_probe += probe
|
||||
synth_obj += obj
|
||||
obj_stack.append(cmath.torch_to_complex(obj))
|
||||
|
||||
|
||||
|
||||
# If there only was one image
|
||||
try:
|
||||
i
|
||||
except:
|
||||
i = -1
|
||||
|
||||
synth_probe = cmath.torch_to_complex(synth_probe)
|
||||
synth_obj = cmath.torch_to_complex(synth_obj)
|
||||
return synth_probe/(i+2), synth_obj/(i+2), obj_stack
|
||||
|
||||
|
||||
|
||||
def calc_prtf(synth_obj, objects, basis, obj_slice=None):
|
||||
@@ -114,16 +68,26 @@ if __name__ == '__main__':
|
||||
|
||||
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_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'])
|
||||
|
||||
|
||||
|
||||
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()
|
||||
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.')
|
||||
|
||||
+111
-11
@@ -5,7 +5,7 @@ import numpy as np
|
||||
from CDTools.tools import cmath
|
||||
from CDTools.tools import image_processing as ip
|
||||
|
||||
__all__ = ['orthogonalize_probes','standardize']
|
||||
__all__ = ['orthogonalize_probes','standardize', 'synthesize_reconstructions']
|
||||
|
||||
from matplotlib import pyplot as plt
|
||||
def orthogonalize_probes(probes):
|
||||
@@ -77,14 +77,21 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
|
||||
|
||||
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.
|
||||
noise. The default slice is from 3/8 to 5/8 of the way across. If the
|
||||
probe is actually a stack of incoherently mixing probes, then the
|
||||
dominant probe mode (assumed to be the first in the list) is used, but
|
||||
all the probes are updated with the same factors.
|
||||
|
||||
Args:
|
||||
probe (t.tensor) : tensor or numpy array storing a retrieved probe
|
||||
probe (t.tensor) : tensor or numpy array storing a retrieved probe or stack of incoherently mixed probes
|
||||
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
|
||||
|
||||
Returns:
|
||||
(t.tensor) : The standardized probe
|
||||
(t.tensor) : The standardized object
|
||||
|
||||
"""
|
||||
# First, we normalize the probe intensity to a fixed value.
|
||||
probe_np = False
|
||||
@@ -95,8 +102,16 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
|
||||
if isinstance(obj, np.ndarray):
|
||||
obj = cmath.complex_to_torch(obj).to(t.float32)
|
||||
obj_np = True
|
||||
|
||||
# If this is a single probe and not a stack of probes
|
||||
if len(probe.shape) == 3:
|
||||
probe = probe[None,...]
|
||||
single_probe = True
|
||||
else:
|
||||
single_probe = False
|
||||
|
||||
|
||||
normalization = t.sqrt(t.sum(cmath.cabssq(probe)) / (len(probe.view(-1))/2))
|
||||
normalization = t.sqrt(t.sum(cmath.cabssq(probe[0])) / (len(probe[0].view(-1))/2))
|
||||
probe = probe / normalization
|
||||
obj = obj * normalization
|
||||
|
||||
@@ -108,13 +123,13 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
|
||||
|
||||
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)
|
||||
center_freq = ip.centroid_sq(cmath.fftshift(t.fft(probe[0],2)),comp=True)
|
||||
center_freq -= (t.tensor(probe[0].shape[:-1]) // 2).to(t.float32)
|
||||
center_freq /= t.tensor(probe[0].shape[:-1]).to(t.float32)
|
||||
|
||||
|
||||
|
||||
Is, Js = np.mgrid[:probe.shape[0],:probe.shape[1]]
|
||||
Is, Js = np.mgrid[:probe[0].shape[0],:probe[0].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)))
|
||||
@@ -127,12 +142,17 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
|
||||
|
||||
|
||||
# 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))
|
||||
|
||||
for i in range(probe.shape[0]):
|
||||
probe_angle = cmath.cphase(t.sum(probe[i],dim=(0,1)))
|
||||
probe[i] = cmath.cmult(probe[i], cmath.expi(-probe_angle))
|
||||
|
||||
if single_probe:
|
||||
probe = probe[0]
|
||||
|
||||
if probe_np:
|
||||
probe = cmath.torch_to_complex(probe.detach().cpu())
|
||||
if obj_np:
|
||||
@@ -140,3 +160,83 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
|
||||
|
||||
return probe, obj
|
||||
|
||||
|
||||
|
||||
def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None, correct_ramp=False):
|
||||
"""Takes a collection of reconstructions and outputs a single synthesized probe and object
|
||||
|
||||
The function first standardizes the sets of probes and objects using the
|
||||
standardize function, passing through the relevant options. Then it
|
||||
calculates the closest overlap of subsequent frames to subpixel
|
||||
precision and uses a sinc interpolation to shift all the probes and objects
|
||||
to a common frame. Then the images are summed.
|
||||
|
||||
Args:
|
||||
probes (list) : A list of probes or stacks of probe modes
|
||||
objects (list) : A list of objects
|
||||
use_probe (bool) : Default False, whether to use the probe or object for alignment
|
||||
obj_slice (slice) : Optional, A slice of the object to use for alignment and normalization
|
||||
correct_ramp (bool) : Default False, whether to correct for a relative phase ramp in the probe and object
|
||||
|
||||
Returns:
|
||||
(array_like) : The synthesized probe
|
||||
(array_like) : The synthesized object
|
||||
(list) : a list of standardized objects, for further processing
|
||||
|
||||
"""
|
||||
|
||||
probe_np = False
|
||||
if isinstance(probes[0], np.ndarray):
|
||||
probes = [cmath.complex_to_torch(probe).to(t.float32) for probe in probes]
|
||||
probe_np = True
|
||||
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 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_probe, synth_obj = standardize(probes[0], objects[0], obj_slice=obj_slice,correct_ramp=correct_ramp)
|
||||
obj_stack = [synth_obj]
|
||||
|
||||
for i, (probe, obj) in enumerate(zip(probes[1:],objects[1:])):
|
||||
probe, obj = standardize(probe, obj, obj_slice=obj_slice,correct_ramp=correct_ramp)
|
||||
|
||||
|
||||
if use_probe:
|
||||
shift = ip.find_shift(synth_probe[0],probe[0], resolution=50)
|
||||
else:
|
||||
shift = ip.find_shift(synth_obj[obj_slice],obj[obj_slice], resolution=50)
|
||||
|
||||
|
||||
obj = ip.sinc_subpixel_shift(obj,np.array(shift))
|
||||
|
||||
if len(probe.shape) == 4:
|
||||
probe = t.stack([ip.sinc_subpixel_shift(p,tuple(shift))
|
||||
for p in probe],dim=0)
|
||||
else:
|
||||
probe = ip.sinc_subpixel_shift(probe,tuple(shift))
|
||||
|
||||
synth_probe += probe
|
||||
synth_obj += obj
|
||||
obj_stack.append(obj)
|
||||
|
||||
|
||||
# If there only was one image
|
||||
try:
|
||||
i
|
||||
except:
|
||||
i = -1
|
||||
|
||||
if probe_np:
|
||||
synth_probe = cmath.torch_to_complex(synth_probe)
|
||||
if obj_np:
|
||||
synth_obj = cmath.torch_to_complex(synth_obj)
|
||||
obj_stack = [cmath.torch_to_complex(obj) for obj in obj_stack]
|
||||
|
||||
return synth_probe/(i+2), synth_obj/(i+2), obj_stack
|
||||
|
||||
|
||||
@@ -42,7 +42,6 @@ def test_orthogonalize_probes():
|
||||
assert np.allclose(probe_intensity,ortho_probe_intensity)
|
||||
|
||||
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
def test_standardize():
|
||||
|
||||
@@ -80,7 +79,7 @@ def test_standardize():
|
||||
assert np.allclose(obj, s_obj)
|
||||
|
||||
|
||||
# And ensure that standardization maps back to the standard versions
|
||||
# Then do one with a phase ramp
|
||||
phase_ramp_dir = (np.random.rand(2) - 0.5)
|
||||
|
||||
probe_Xs, probe_Ys = np.mgrid[:probe.shape[0],:probe.shape[1]]
|
||||
@@ -89,12 +88,31 @@ def test_standardize():
|
||||
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]+
|
||||
obj_phase_ramp = np.exp(-1j*obj_Ys * phase_ramp_dir[1]+
|
||||
-1j*obj_Xs * phase_ramp_dir[0])
|
||||
test_obj = test_obj * phase_ramp
|
||||
test_obj = test_obj * 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
|
||||
|
||||
# Finally a test with the phase ramp and multiple probes
|
||||
subdominant_probe = 0.1*np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
|
||||
subdominant_probe = subdominant_probe * np.exp(-1j * np.angle(np.sum(subdominant_probe)))
|
||||
test_subdominant_probe = subdominant_probe * 37.6
|
||||
test_subdominant_probe = test_subdominant_probe * phase_ramp
|
||||
|
||||
incoh_probe = np.array([test_probe,test_subdominant_probe])
|
||||
|
||||
s_probe, s_obj = analysis.standardize(incoh_probe, test_obj, correct_ramp=True)
|
||||
|
||||
assert np.max(s_probe[0] - probe) / np.max(np.abs(probe)) < 1e-4
|
||||
assert np.max(s_obj - obj) / np.max(np.abs(obj)) < 1e-4
|
||||
assert np.max(s_probe[1] - subdominant_probe) / np.max(np.abs(subdominant_probe)) < 1e-4
|
||||
|
||||
|
||||
|
||||
from matplotlib import pyplot as plt
|
||||
def test_synthesize_reconstructions():
|
||||
pass
|
||||
|
||||
Reference in New Issue
Block a user