mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
Move the standardizing tool over to analysis and write a test for it
This commit is contained in:
Abe Levitan
@@ -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())
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user