mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-27 20:52:08 +02:00
Lots of small changes
This commit is contained in:
Abe Levitan
@@ -138,7 +138,7 @@ class CDIModel(t.nn.Module):
|
||||
|
||||
# Define the scheduler
|
||||
if schedule:
|
||||
scheduler = t.optim.ReduceLROnPlateau(optimizer, factor=0.2)
|
||||
scheduler = t.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2,threshold=1e-9)
|
||||
else:
|
||||
scheduler = None
|
||||
|
||||
|
||||
@@ -15,13 +15,14 @@ from copy import copy
|
||||
class FancyPtycho(CDIModel):
|
||||
|
||||
def __init__(self, wavelength, detector_geometry,
|
||||
probe_basis, detector_slice,
|
||||
probe_basis,
|
||||
probe_guess, obj_guess,
|
||||
detector_slice=None,
|
||||
surface_normal=np.array([0.,0.,1.]),
|
||||
min_translation = t.Tensor([0,0]),
|
||||
background = None, translation_offsets=None, mask=None,
|
||||
weights = None, translation_scale = 1, saturation=None,
|
||||
probe_support = None, obj_support=None):
|
||||
probe_support = None, obj_support=None, oversampling=1):
|
||||
|
||||
super(FancyPtycho,self).__init__()
|
||||
self.wavelength = t.Tensor([wavelength])
|
||||
@@ -60,7 +61,10 @@ class FancyPtycho(CDIModel):
|
||||
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
|
||||
|
||||
if background is None:
|
||||
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
|
||||
if detector_slice is not None:
|
||||
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
|
||||
else:
|
||||
background = 1e-6 * t.ones(self.probe[0].shape[:-1])
|
||||
|
||||
|
||||
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
|
||||
@@ -87,10 +91,12 @@ class FancyPtycho(CDIModel):
|
||||
self.obj.data = self.obj * obj_support
|
||||
else:
|
||||
self.obj_support = t.ones_like(self.obj)
|
||||
|
||||
self.oversampling = oversampling
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None):
|
||||
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1):
|
||||
|
||||
wavelength = dataset.wavelength
|
||||
det_basis = dataset.detector_geometry['basis']
|
||||
@@ -114,7 +120,8 @@ class FancyPtycho(CDIModel):
|
||||
distance,
|
||||
center=center,
|
||||
padding=padding,
|
||||
opt_for_fft=False)
|
||||
opt_for_fft=False,
|
||||
oversampling=oversampling)
|
||||
|
||||
|
||||
if hasattr(dataset, 'sample_info') and \
|
||||
@@ -139,8 +146,8 @@ class FancyPtycho(CDIModel):
|
||||
# Next generate the object geometry from the probe geometry and
|
||||
# the translations
|
||||
pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
|
||||
|
||||
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=50)
|
||||
|
||||
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
|
||||
|
||||
if hasattr(dataset, 'background') and dataset.background is not None:
|
||||
background = t.sqrt(dataset.background)
|
||||
@@ -149,7 +156,7 @@ class FancyPtycho(CDIModel):
|
||||
|
||||
# Finally, initialize the probe and object using this information
|
||||
if probe_size is None:
|
||||
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance)
|
||||
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
|
||||
else:
|
||||
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
|
||||
|
||||
@@ -191,7 +198,8 @@ class FancyPtycho(CDIModel):
|
||||
else:
|
||||
obj_support = None
|
||||
|
||||
return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj,
|
||||
return cls(wavelength, det_geo, probe_basis, probe, obj,
|
||||
detector_slice=det_slice,
|
||||
surface_normal=surface_normal,
|
||||
min_translation=min_translation,
|
||||
translation_offsets = translation_offsets,
|
||||
@@ -199,7 +207,8 @@ class FancyPtycho(CDIModel):
|
||||
translation_scale=translation_scale,
|
||||
saturation=saturation,
|
||||
probe_support=probe_support,
|
||||
obj_support=obj_support)
|
||||
obj_support=obj_support,
|
||||
oversampling=oversampling)
|
||||
|
||||
|
||||
def interaction(self, index, translations):
|
||||
@@ -225,9 +234,10 @@ class FancyPtycho(CDIModel):
|
||||
exit_waves = self.weights[index][:,None,None,None] * exit_waves
|
||||
else:
|
||||
exit_waves = self.weights[index] * exit_waves
|
||||
|
||||
|
||||
all_exit_waves.append(exit_waves)
|
||||
|
||||
|
||||
return t.stack(all_exit_waves)
|
||||
|
||||
|
||||
@@ -244,7 +254,8 @@ class FancyPtycho(CDIModel):
|
||||
self.background,
|
||||
detector_slice=self.detector_slice,
|
||||
measurement=tools.measurements.incoherent_sum,
|
||||
saturation=self.saturation )
|
||||
saturation=self.saturation,
|
||||
oversampling=self.oversampling)
|
||||
|
||||
|
||||
def loss(self, sim_data, real_data, mask=None):
|
||||
|
||||
@@ -35,12 +35,17 @@ if __name__ == '__main__':
|
||||
# If it's a length-one reconstruction
|
||||
dataset = {key: [dataset[key]] for key in dataset}
|
||||
calc_prtf = False
|
||||
|
||||
|
||||
# Orthogonalize the probes
|
||||
dataset['probe'] = [orthogonalize_probes(p) for p in dataset['probe']]
|
||||
|
||||
print('hi')
|
||||
|
||||
synth_probe, synth_obj, aligned_objs = synthesize_reconstructions(
|
||||
dataset['probe'], dataset['obj'], args.use_probe)
|
||||
|
||||
print('hey')
|
||||
if calc_prtf:
|
||||
freqs, prtf = calc_consistency_prtf(synth_obj, aligned_objs, dataset['basis'][0])
|
||||
freqs, prtf = calc_consistency_prtf(synth_obj, aligned_objs, dataset['basis'][0], nbins=200)
|
||||
|
||||
# Either plot the only probe, or plot the dominant probe
|
||||
if len(synth_probe.shape) == 2:
|
||||
@@ -48,36 +53,45 @@ if __name__ == '__main__':
|
||||
plotting.plot_amplitude(synth_probe,basis=dataset['basis'][0])
|
||||
plotting.plot_colorized(synth_probe,basis=dataset['basis'][0])
|
||||
else:
|
||||
plotting.plot_phase(synth_probe[0],basis=dataset['basis'][0])
|
||||
plotting.plot_amplitude(synth_probe[0],basis=dataset['basis'][0])
|
||||
plotting.plot_colorized(synth_probe[0],basis=dataset['basis'][0])
|
||||
# Plot as many probes as exist
|
||||
try:
|
||||
for i in range(0,50):
|
||||
plotting.plot_phase(synth_probe[i],basis=dataset['basis'][0])
|
||||
plt.title('Probe ' + str(i+1) + 'Phase')
|
||||
plotting.plot_amplitude(synth_probe[i],basis=dataset['basis'][0])
|
||||
plt.title('Probe ' + str(i+1) + ' Amplitude')
|
||||
plotting.plot_colorized(synth_probe[i],basis=dataset['basis'][0])
|
||||
plt.title('Probe ' + str(i+1) + ' Colorized')
|
||||
except IndexError:
|
||||
pass
|
||||
|
||||
plotting.plot_amplitude(aligned_objs[0][300:-300,300:-300],basis=dataset['basis'][0])
|
||||
plotting.plot_colorized(aligned_objs[0][300:-300,300:-300],basis=dataset['basis'][0])
|
||||
plotting.plot_phase(aligned_objs[0][300:-300,300:-300],basis=dataset['basis'][0])
|
||||
|
||||
# plot the subdominant probe if it exists
|
||||
plotting.plot_amplitude(synth_obj[300:-300,300:-300],basis=dataset['basis'][0])
|
||||
plotting.plot_colorized(synth_obj[300:-300,300:-300],basis=dataset['basis'][0])
|
||||
plotting.plot_phase(synth_obj[300:-300,300:-300],basis=dataset['basis'][0])
|
||||
|
||||
|
||||
try:
|
||||
plotting.plot_phase(synth_probe[1],basis=dataset['basis'][0])
|
||||
plotting.plot_amplitude(synth_probe[1],basis=dataset['basis'][0])
|
||||
plotting.plot_colorized(synth_probe[1],basis=dataset['basis'][0])
|
||||
real_translations = dataset['translation'][0]
|
||||
real_translations -= np.min(real_translations,axis=0)[None,:]
|
||||
real_translations = real_translations
|
||||
plotting.plot_translations(real_translations)
|
||||
|
||||
plotting.plot_nanomap(real_translations,dataset['weights'][0])
|
||||
except:
|
||||
pass
|
||||
|
||||
plotting.plot_amplitude(synth_obj,basis=dataset['basis'][0])
|
||||
plotting.plot_colorized(synth_obj,basis=dataset['basis'][0])
|
||||
plotting.plot_phase(synth_obj,basis=dataset['basis'][0])
|
||||
|
||||
|
||||
plt.figure()
|
||||
try:
|
||||
real_translations = dataset['basis'][0].dot(dataset['translation'][0].transpose())
|
||||
real_translations -= np.min(real_translations,axis=1)[:,None]
|
||||
real_translations = real_translations.transpose()
|
||||
plotting.plot_translations(real_translations)
|
||||
plt.figure()
|
||||
except:
|
||||
pass
|
||||
plt.imshow(np.sqrt(dataset['background'][0]))
|
||||
|
||||
if calc_prtf:
|
||||
plt.figure()
|
||||
plt.plot(freqs*1e-6, prtf)
|
||||
plt.xlabel('Spatial Frequency (cycles/um)')
|
||||
plt.ylabel('Consistency Based PRTF')
|
||||
plt.grid()
|
||||
|
||||
plt.show()
|
||||
|
||||
@@ -199,15 +199,19 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
|
||||
objects = [cmath.complex_to_torch(obj).to(t.float32) for obj in objects]
|
||||
obj_np = True
|
||||
|
||||
obj_shape = np.min(np.array([obj.shape[:-1] for obj in objects]),axis=0)
|
||||
objects = [obj[:obj_shape[0],:obj_shape[1]] for obj in objects]
|
||||
|
||||
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].clone(), objects[0].clone(), obj_slice=obj_slice,correct_ramp=correct_ramp)
|
||||
obj_stack = [synth_obj]
|
||||
|
||||
obj_stack = [synth_obj]
|
||||
|
||||
for i, (probe, obj) in enumerate(zip(probes[1:],objects[1:])):
|
||||
probe, obj = standardize(probe.clone(), obj.clone(), obj_slice=obj_slice,correct_ramp=correct_ramp)
|
||||
if use_probe:
|
||||
@@ -218,6 +222,7 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
|
||||
|
||||
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)
|
||||
|
||||
@@ -74,6 +74,7 @@ def sinc_subpixel_shift(im, shift):
|
||||
Returns:
|
||||
(torch.Tensor) : The subpixel shifted tensor
|
||||
"""
|
||||
|
||||
i = t.arange(im.shape[0]) - im.shape[0]//2
|
||||
j = t.arange(im.shape[1]) - im.shape[1]//2
|
||||
I,J = t.meshgrid(i,j)
|
||||
@@ -231,23 +232,29 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True):
|
||||
"""
|
||||
|
||||
complex_things = 2
|
||||
im_complex = True
|
||||
if image.shape[-1] != 2:
|
||||
image = t.stack((image,t.zeros_like(image)),dim=-1)
|
||||
complex_things -= 1
|
||||
im_complex = False
|
||||
|
||||
if kernel.shape[-1] != 2:
|
||||
kernel = t.stack((kernel,t.zeros_like(kernel)),dim=-1)
|
||||
complex_things -= 1
|
||||
|
||||
# Take a correlation
|
||||
if fftshift_kernel:
|
||||
kernel = cmath.ifftshift(kernel)
|
||||
|
||||
# If the image wasn't originally complex, and the dimension
|
||||
# was passed with the nexative-indexing convention
|
||||
if not im_complex and dim < 0:
|
||||
dim = dim-1
|
||||
|
||||
# We have to transpose the relevant dimension to -2 before using the fft,
|
||||
# which expects to operate on the final non-complex dimension
|
||||
trans_im = t.transpose(image, dim, -2)
|
||||
|
||||
|
||||
# Take a correlation
|
||||
fft_im = t.fft(trans_im, 1)
|
||||
fft_kernel = t.fft(kernel, 1)
|
||||
trans_conv = t.ifft(cmath.cmult(fft_im,fft_kernel), 1)
|
||||
|
||||
@@ -11,7 +11,7 @@ from scipy.fftpack import next_fast_len
|
||||
import numpy as np
|
||||
|
||||
|
||||
def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0):
|
||||
def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0, oversampling=1):
|
||||
"""Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry
|
||||
|
||||
It takes in the parameters for a given detector - the basis defining
|
||||
@@ -28,7 +28,8 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
|
||||
distance (float) : The sample-detector distance, in m
|
||||
center (torch.Tensor) : If defined, the location of the zero frequency pixel
|
||||
opt_for_fft (bool) : Default is true, whether to increase detector size to improve fft performance
|
||||
padding (int) : Default is 0, an extra border to allow for subpixel shifting later
|
||||
padding (int) : Default is 0, the size of an extra border of nonphysical pixels around the detector
|
||||
oversampling (int) : Default is 1, the amount to multiply the exit wave shape by.
|
||||
|
||||
Returns:
|
||||
torch.Tensor : The exit wave basis
|
||||
@@ -44,7 +45,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
|
||||
center = det_shape // 2
|
||||
else:
|
||||
center = t.Tensor(center).to(t.int32)
|
||||
|
||||
|
||||
# Then, calculate the required detector size from the centering
|
||||
# This is a bit opaque but was worth doing accurately
|
||||
min_left = center * 2
|
||||
@@ -76,7 +77,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
|
||||
(full_shape.to(t.float32) * t.norm(det_basis,dim=0))
|
||||
|
||||
# Finally, convert the shape back to a torch.Size
|
||||
full_shape = t.Size([dim for dim in full_shape])
|
||||
full_shape = t.Size([dim * oversampling for dim in full_shape])
|
||||
|
||||
|
||||
return real_space_basis, full_shape, det_slice
|
||||
@@ -101,7 +102,7 @@ def calc_object_setup(probe_shape, translations, padding=0):
|
||||
torch.Size : required size of object array
|
||||
torch.Tensor : minimum pixel-valued translation
|
||||
"""
|
||||
|
||||
|
||||
# First we look at the translations to find the minimum translation
|
||||
# and the range of translations
|
||||
min_translation = t.min(translations, dim=0)[0]
|
||||
@@ -211,7 +212,7 @@ def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0):
|
||||
return avg_intensity / probe_intensity * probe
|
||||
|
||||
|
||||
def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None):
|
||||
def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, oversampling=1):
|
||||
"""Generates a SHARP style probe guess from a dataset
|
||||
|
||||
What we call the "SHARP" style probe guess is to take a mean of all
|
||||
@@ -233,12 +234,13 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None):
|
||||
dataset (Ptycho_2D_Dataset) : The dataset to work from
|
||||
shape (torch.Size) : The size of the probe array to simulate
|
||||
det_slice (slice) : A slice or tuple of slices corresponding to the detector region in Fourier space
|
||||
propagatioin_distance (float) : Default is no propagation, an amount to propagate the guessed probe from it's focal point
|
||||
propagation_distance (float) : Default is no propagation, an amount to propagate the guessed probe from it's focal point
|
||||
oversampling (int) : Default 1, the width of the region of pixels in the wavefield to bin into a single detector pixel
|
||||
"""
|
||||
|
||||
|
||||
# to use the mask or not?
|
||||
intensities = np.zeros(shape)
|
||||
intensities = np.zeros([dim // oversampling for dim in shape])
|
||||
for params, im in dataset:
|
||||
if hasattr(dataset,'mask') and dataset.mask is not None:
|
||||
intensities[det_slice] += dataset.mask.cpu().numpy() * im.cpu().numpy()
|
||||
@@ -255,9 +257,7 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None):
|
||||
probe_guess = cmath.torch_to_complex(inverse_far_field(probe_fft))
|
||||
|
||||
# Now we remove the central pixel
|
||||
|
||||
center = np.array(probe_guess.shape) // 2
|
||||
|
||||
|
||||
# I'm always divided on whether to use this modification:
|
||||
|
||||
@@ -267,13 +267,12 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None):
|
||||
probe_guess[center[0], center[1]-1],
|
||||
probe_guess[center[0], center[1]+1]])
|
||||
|
||||
|
||||
probe_guess = cmath.complex_to_torch(probe_guess)
|
||||
|
||||
if propagation_distance is not None:
|
||||
# First generate the propagation array
|
||||
|
||||
probe_shape = t.Tensor(tuple(shape))
|
||||
probe_shape = t.Tensor(tuple(probe_guess.shape))[:-1]
|
||||
|
||||
# Start by recalculating the probe basis from the given information
|
||||
det_basis = t.Tensor(dataset.detector_geometry['basis'])
|
||||
@@ -291,7 +290,15 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None):
|
||||
AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance)
|
||||
|
||||
probe_guess = near_field(probe_guess,AS_prop)
|
||||
|
||||
|
||||
return probe_guess
|
||||
# Finally, place this probe in a full-sized array if there is oversampling
|
||||
final_probe = t.zeros([dim for dim in shape] + [2])
|
||||
left = shape[0]//2 - probe_guess.shape[0] // 2
|
||||
top = shape[1]//2 - probe_guess.shape[1] // 2
|
||||
final_probe[left:left+probe_guess.shape[0],
|
||||
top:top+probe_guess.shape[1],:] = probe_guess
|
||||
|
||||
return final_probe
|
||||
|
||||
|
||||
|
||||
@@ -37,12 +37,13 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
|
||||
|
||||
projection_1 = t.Tensor([[1,0,0],
|
||||
[0,1,0],
|
||||
[0,0,0]]).to(device=basis.device,dtype=basis.dtype)
|
||||
[0,0,0]]).to(device=translations.device,dtype=translations.dtype)
|
||||
projection_2 = t.inverse(t.Tensor([[1,0,0],
|
||||
[0,1,0],
|
||||
-surface_normal/
|
||||
surface_normal[2]])).to(device=basis.device,dtype=basis.dtype)
|
||||
basis_vectors_inv = t.pinverse(basis)
|
||||
surface_normal[2]])).to(device=translations.device,dtype=translations.dtype)
|
||||
basis_vectors_inv = t.pinverse(basis).to(device=translations.device,
|
||||
dtype=translations.dtype)
|
||||
projection = t.mm(basis_vectors_inv,
|
||||
t.mm(projection_2,projection_1))
|
||||
projection = projection.t()
|
||||
|
||||
@@ -3,6 +3,7 @@ from __future__ import division, print_function, absolute_import
|
||||
from CDTools.tools import cmath
|
||||
import torch as t
|
||||
import numpy as np
|
||||
from torch.nn.functional import avg_pool2d
|
||||
|
||||
#
|
||||
# This file will host tools to turn a propagated wavefield into a measured
|
||||
@@ -12,7 +13,7 @@ import numpy as np
|
||||
__all__ = ['intensity', 'incoherent sum', 'quadratic_background']
|
||||
|
||||
|
||||
def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None):
|
||||
def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
|
||||
"""Returns the intensity of a wavefield
|
||||
|
||||
The intensity is defined as the magnitude squared of the
|
||||
@@ -23,24 +24,41 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None):
|
||||
wavefield (torch.Tensor) : A JxMxNx2 stack of complex wavefields
|
||||
detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
|
||||
saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
|
||||
oversampling (int) : Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
|
||||
|
||||
Returns:
|
||||
torch.Tensor : A real MxN array storing the wavefield's intensities
|
||||
"""
|
||||
if detector_slice is None:
|
||||
output = cmath.cabssq(wavefield) + epsilon
|
||||
else:
|
||||
output = cmath.cabssq(wavefield) + epsilon
|
||||
|
||||
# Now we apply oversampling
|
||||
if oversampling != 1:
|
||||
dim = output.dim()
|
||||
if dim == 2:
|
||||
output = output[None,None,:,:]
|
||||
if dim == 3:
|
||||
output = output[None,:,:,:]
|
||||
output = avg_pool2d(output, 2, 2)
|
||||
if dim == 2:
|
||||
output = output[0,0,:,:]
|
||||
if dim == 3:
|
||||
output = output[0,:,:,:]
|
||||
|
||||
# Then we grab the detector slice
|
||||
if detector_slice is not None:
|
||||
if wavefield.dim() == 3:
|
||||
output = cmath.cabssq(wavefield[detector_slice]) + epsilon
|
||||
output = output[detector_slice]
|
||||
else:
|
||||
output = cmath.cabssq(wavefield[(np.s_[:],) + detector_slice]) + epsilon
|
||||
output = output[(np.s_[:],) + detector_slice]
|
||||
|
||||
# And now saturation
|
||||
if saturation is None:
|
||||
return output
|
||||
else:
|
||||
return t.clamp(output,0,saturation)
|
||||
|
||||
|
||||
def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None):
|
||||
def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
|
||||
"""Returns the incoherent sum of the intensities of the wavefields
|
||||
|
||||
The intensity is defined as the sum of the magnitudes squared of
|
||||
@@ -56,26 +74,41 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
|
||||
wavefields (torch.Tensor) : An LxJxMxNx2 stack of complex wavefields
|
||||
detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
|
||||
saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
|
||||
|
||||
oversampling (int) : Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
|
||||
Returns:
|
||||
torch.Tensor : A real JXMxN array storing the incoherently summed intensities
|
||||
"""
|
||||
# This syntax just adds an axis to the slice to preserve the J direction
|
||||
if detector_slice is None:
|
||||
output = t.sum(cmath.cabssq(wavefields),dim=0) + epsilon
|
||||
else:
|
||||
|
||||
output = t.sum(cmath.cabssq(wavefields),dim=0) + epsilon
|
||||
|
||||
# Now we apply oversampling
|
||||
if oversampling != 1:
|
||||
dim = output.dim()
|
||||
if dim == 2:
|
||||
output = output[None,None,:,:]
|
||||
if dim == 3:
|
||||
output = output[None,:,:,:]
|
||||
output = avg_pool2d(output, 2, 2)
|
||||
if dim == 2:
|
||||
output = output[0,0,:,:]
|
||||
if dim == 3:
|
||||
output = output[0,:,:,:]
|
||||
|
||||
# Then we grab the detector slice
|
||||
if detector_slice is not None:
|
||||
if wavefields.dim() == 4:
|
||||
output = t.sum(cmath.cabssq(wavefields[(np.s_[:],)+detector_slice]),dim=0) + epsilon
|
||||
output = output[detector_slice]
|
||||
else:
|
||||
output = t.sum(cmath.cabssq(wavefields[(np.s_[:],np.s_[:])+detector_slice]),dim=0) + epsilon
|
||||
|
||||
output = output[(np.s_[:],) + detector_slice]
|
||||
|
||||
if saturation is None:
|
||||
return output
|
||||
else:
|
||||
return t.clamp(output,0,saturation)
|
||||
|
||||
|
||||
def quadratic_background(wavefield, background, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None):
|
||||
def quadratic_background(wavefield, background, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1):
|
||||
"""Returns the intensity of a wavefield plus a background
|
||||
|
||||
The intensity is calculated via the given measurment function
|
||||
@@ -89,16 +122,21 @@ def quadratic_background(wavefield, background, detector_slice=None, measurement
|
||||
detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
|
||||
measurement (function) : Optional, the measurement function to use. The default is measurements.intensity
|
||||
saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
|
||||
|
||||
oversampling (int) : Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
|
||||
Returns:
|
||||
torch.Tensor : A real MxN array storing the wavefield's intensities
|
||||
"""
|
||||
|
||||
if detector_slice is None:
|
||||
output = measurement(wavefield, epsilon=epsilon) + background**2
|
||||
output = measurement(wavefield, epsilon=epsilon,
|
||||
oversampling=oversampling) + background**2
|
||||
else:
|
||||
output = measurement(wavefield, detector_slice, epsilon=epsilon) \
|
||||
+ background**2
|
||||
output = measurement(wavefield, detector_slice,
|
||||
epsilon=epsilon, oversampling=oversampling) \
|
||||
+ background**2
|
||||
|
||||
# This has to be done after the background is added, hence we replicate
|
||||
# it here
|
||||
if saturation is None:
|
||||
return output
|
||||
else:
|
||||
|
||||
@@ -48,21 +48,23 @@ dataset.get_as(device='cuda')
|
||||
|
||||
# We can run the first phase of phase retrieval while leaving
|
||||
# the probe positions fixed (whether this is good is debatable)
|
||||
model.translation_offsets.requires_grad = False
|
||||
# model.translation_offsets.requires_grad = False
|
||||
|
||||
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15)):
|
||||
for i, loss in enumerate(model.Adam_optimize(15, dataset, batch_size=15)):
|
||||
print(i,loss)
|
||||
model.inspect(dataset)
|
||||
|
||||
# And we turn it on for the second phase, as we also lower the learning rate
|
||||
model.translation_offsets.requires_grad = True
|
||||
# model.translation_offsets.requires_grad = True
|
||||
|
||||
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15, lr=0.0005)):
|
||||
for i, loss in enumerate(model.Adam_optimize(15, dataset, batch_size=15, lr=0.0005)):
|
||||
print(i,loss)
|
||||
model.inspect(dataset)
|
||||
|
||||
# The third phase lowers the rate further
|
||||
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15, lr=0.00005)):
|
||||
print(i,loss)
|
||||
|
||||
model.inspect(dataset)
|
||||
|
||||
model.inspect(dataset)
|
||||
model.compare(dataset)
|
||||
|
||||
Binary file not shown.
@@ -12,7 +12,7 @@ with h5py.File(filename,'r') as f:
|
||||
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f)
|
||||
|
||||
|
||||
model = CDTools.models.FancyPtycho.from_dataset(dataset,n_modes=3,randomize_ang=0.1*np.pi)
|
||||
model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2)
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user