cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-15 23:52:38 +02:00
Code Issues Packages Projects Releases Wiki Activity

Make tweaks to the tilt correction to work with the grazing reflection data, add basic atom-modeling for electron ptychography, and update the multislice ptycho to include a bandlimit to avoid aliasing

Browse Source
This commit is contained in:
Abe Levitan
2020-12-21 12:22:04 -05:00
parent 38345da454
commit 75011eaccf
7 changed files with 1535 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/bragg_2d_ptycho.py
+2 -5
View File
@@ -384,12 +384,9 @@ class Bragg2DPtycho(CDIModel):
exit_waves = []
for j in range(translations.size()[0]):
if self.propagate_probe:
#propagator = ggasp(pr.shape, self.probe_basis, self.wavelength,
# t.Tensor([0,0,props[j]]),
# propagation_vector=self.prop_dir,
# dtype=pr.dtype,device=pr.device, propagate_along_offset=True)
# Minus sign is empirical
propagator = cmath.expi((-props[j]*(2*np.pi)/self.wavelength)
propagator = cmath.expi(
(-1*props[j]*(2*np.pi)/self.wavelength)
* self.universal_propagator)
prop_pr = tools.propagators.near_field(pr, propagator)
#plt.close('all')
CDTools/models/fancy_ptycho.py
+6 -5
View File
@@ -143,7 +143,7 @@ class FancyPtycho(CDIModel):
outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
outgoing_dir /= np.linalg.norm(outgoing_dir)
surface_normal = outgoing_dir + np.array([0.,0.,1.])
surface_normal /= np.linalg.norm(outgoing_dir)
surface_normal /= np.linalg.norm(surface_normal)
# Next generate the object geometry from the probe geometry and
@@ -237,10 +237,11 @@ class FancyPtycho(CDIModel):
exit_waves = exit_waves * self.probe_support[...,:,:]
if exit_waves.dim() == 4:
exit_waves = self.weights[index][:,None,None,None] * exit_waves
else:
exit_waves = self.weights[index] * exit_waves
if hasattr(self,'weights') and self.weights is not None:
if exit_waves.dim() == 4:
exit_waves = self.weights[index][:,None,None,None] * exit_waves
else:
exit_waves = self.weights[index] * exit_waves
all_exit_waves.append(exit_waves)
CDTools/models/multislice_2d_ptycho.py
+63 -27
View File
@@ -22,7 +22,10 @@ class Multislice2DPtycho(CDIModel):
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, oversampling=1):
#probe_support = None,
probe_fourier_support=None,
obj_support=None, oversampling=1,
bandlimit=4/5):
super(Multislice2DPtycho,self).__init__()
self.wavelength = t.Tensor([wavelength])
@@ -83,10 +86,11 @@ class Multislice2DPtycho(CDIModel):
self.translation_scale = translation_scale
if probe_support is not None:
self.probe_support = probe_support
else:
self.probe_support = t.ones_like(self.probe[0])
self.probe_fourier_support = t.Tensor(probe_fourier_support).to(t.float32)
#if probe_support is not None:
# self.probe_support = probe_support
#else:
# self.probe_support = t.ones_like(self.probe[0])
if obj_support is not None:
self.obj_support = obj_support
@@ -99,11 +103,13 @@ class Multislice2DPtycho(CDIModel):
spacing = np.linalg.norm(self.probe_basis,axis=0)
shape = np.array(self.probe.shape[1:-1])
self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz)
self.bandlimit = bandlimit
self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=self.bandlimit)
@classmethod
def from_dataset(cls, dataset, dz, nz, 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, auto_center=True):
def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, 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, auto_center=True, bandlimit=4/5):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -170,7 +176,9 @@ class Multislice2DPtycho(CDIModel):
else:
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
# For a Fourier space probe
probe = tools.propagators.inverse_far_field(probe)
# Now we initialize all the subdominant probe modes
probe_max = t.max(cmath.cabs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
@@ -209,6 +217,16 @@ class Multislice2DPtycho(CDIModel):
else:
obj_support = None
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
xs, ys = np.mgrid[:probe.shape[-3],:probe.shape[-2]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
Rs = np.sqrt(xs**2 + ys**2)
probe_support[Rs<probe_convergence_radius] = 1
probe = probe * probe_support[None,:,:]
return cls(wavelength, det_geo, probe_basis, probe, obj, dz, nz,
detector_slice=det_slice,
surface_normal=surface_normal,
@@ -217,9 +235,11 @@ class Multislice2DPtycho(CDIModel):
weights=weights, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
#probe_support=probe_support,
probe_fourier_support=probe_support,
obj_support=obj_support,
oversampling=oversampling)
oversampling=oversampling,
bandlimit=bandlimit)
def interaction(self, index, translations):
@@ -233,7 +253,11 @@ class Multislice2DPtycho(CDIModel):
all_exit_waves = []
for i in range(self.probe.shape[0]):
pr = self.probe[i] * self.probe_support
# For a Fourier-space probe
pr = tools.propagators.inverse_far_field(self.probe[i] * self.probe_fourier_support)
# For a real-space probe
#pr = self.probe[i] * self.probe_support
#exit_waves = pr
#print(self.probe_norm)
#for i in range(self.nz):
@@ -244,21 +268,28 @@ class Multislice2DPtycho(CDIModel):
exit_wave = self.probe_norm * pr
for i in range(self.nz-1):
#exit_wave = tools.interactions.ptycho_2D_sinc(exit_wave,
# self.obj_support * self.obj,
# trans,
# shift_probe=True)
exit_wave = tools.interactions.ptycho_2D_round(exit_wave,
self.obj_support * cmath.cexpi(self.obj.data/self.nz),
trans)
exit_wave = tools.interactions.ptycho_2D_sinc(exit_wave,
self.obj_support * cmath.cexpi(self.obj/self.nz),#self.obj.data,
trans,
shift_probe=True)
#exit_wave = tools.interactions.ptycho_2D_round(exit_wave,
# self.obj_support * cmath.cexpi(self.obj.data/self.nz),
# trans)
exit_wave = tools.propagators.near_field(exit_wave,self.as_prop)
#tools.plotting.plot_amplitude(exit_wave)
#plt.show()
# only final layer gets a derivative
exit_wave = tools.interactions.ptycho_2D_round(exit_wave,
self.obj_support * cmath.cexpi(0.1*self.obj),
trans)
exit_wave = tools.interactions.ptycho_2D_sinc(exit_wave,
self.obj_support * cmath.cexpi(self.obj/self.nz),#self.obj,
trans,
shift_probe=True)
#exit_wave = tools.interactions.ptycho_2D_round(exit_wave,
# self.obj_support * cmath.cexpi(self.obj/self.nz),
# trans)
# One final propagation to enforce the bandlimit
exit_wave = tools.propagators.near_field(exit_wave,self.as_prop)
exit_waves.append(exit_wave)
exit_waves = t.stack(exit_waves)
@@ -328,7 +359,8 @@ class Multislice2DPtycho(CDIModel):
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.probe_norm = self.probe_norm.to(*args,**kwargs)
self.probe_support = self.probe_support.to(*args,**kwargs)
#self.probe_support = self.probe_support.to(*args,**kwargs)
self.probe_fourier_support = self.probe_fourier_support.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
self.as_prop = self.as_prop.to(*args, **kwargs)
@@ -389,10 +421,14 @@ class Multislice2DPtycho(CDIModel):
('Subdominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
#('Object Amplitude',
# lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
#('Object Phase',
# lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
('Real Part of T',
lambda self, fig: p.plot_amplitude(self.obj[:,:,0].detach().cpu().numpy(), fig=fig, basis=self.probe_basis)),
('Imaginary Part of T',
lambda self, fig: p.plot_amplitude(self.obj[:,:,1].detach().cpu().numpy(), fig=fig, basis=self.probe_basis)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
('Background',
CDTools/models/s_matrix_ptycho.py
+154 -89
View File
@@ -20,7 +20,8 @@ class SMatrixPtycho(CDIModel):
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
background = None, translation_offsets=None,
probe_planes = None, mask=None,
weights = None, translation_scale = 1, saturation=None,
oversampling=1):
@@ -50,10 +51,12 @@ class SMatrixPtycho(CDIModel):
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 3:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
# Remember that for S-matrix we have several probes for different
# planes
if probe_guess.dim() > 4:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0,0].to(t.float32)))
else:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
/ self.probe_norm)
@@ -61,13 +64,14 @@ class SMatrixPtycho(CDIModel):
self.s_matrix = t.nn.Parameter(s_matrix_guess.to(t.float32))
if background is None:
ew_shape = [s_matrix_guess.shape[0] - 1 + probe_guess.shape[1],
s_matrix_guess.shape[1] - 1 + probe_guess.shape[2]]
ew_shape = [s_matrix_guess.shape[0] - 1 + probe_guess.shape[-3],
s_matrix_guess.shape[1] - 1 + probe_guess.shape[-2]]
if detector_slice is not None:
background = 1e-6 * t.ones(t.ones(ew_shape)[self.detector_slice].shape).to(t.float32)
else:
background = 1e-6 * t.ones(ew_shape).to(t.float32)
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
if weights is None:
@@ -78,7 +82,14 @@ class SMatrixPtycho(CDIModel):
if translation_offsets is None:
self.translation_offsets = None
else:
self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale)
self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale)
# This maps indices to probe planes to be used. If none, it defaults
# to always being plane 0
if probe_planes is None:
self.probe_planes = None
else:
self.probe_planes = t.LongTensor(probe_planes)
self.translation_scale = translation_scale
@@ -88,44 +99,94 @@ class SMatrixPtycho(CDIModel):
@classmethod
def from_dataset(cls, dataset, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True):
def from_dataset(cls, dataset, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1):
datasets = [dataset]
propagation_distances = [propagation_distance]
# We only return the 0th element because in the general case, the
# constructor needs to return a stacked datset in addition to
# a model, but for the case of one dataset we only need to return
# the model.
return cls.from_datasets(datasets, probe_convergence_radius,
locality_radius=locality_radius,
probe_size=probe_size,
randomize_ang=randomize_ang,
padding=padding,
n_modes=n_modes,
translation_scale=translation_scale,
saturation=saturation,
propagation_distances=propagation_distances,
scattering_mode=scattering_mode,
oversampling=oversampling)[0]
# This is for the multi-focal-plane case, where each dataset will correspond
# to a different focal plane. The guess propagation distance for each
# dataset can be set individually but otherwise the probes are
# reconstructed entirely separately. All datasets are assumed to have
# the same basic parameters (wavelength, detector geometry, etc) and share
# the same origin in the x-y plane.
@classmethod
def from_datasets(cls, datasets, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distances=None, scattering_mode=None, oversampling=1):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
det_shape = dataset[0][1].shape
distance = dataset.detector_geometry['distance']
wavelength = datasets[0].wavelength
det_basis = datasets[0].detector_geometry['basis']
det_shape = datasets[0][0][1].shape
distance = datasets[0].detector_geometry['distance']
# always do this on the cpu
get_as_args = dataset.get_as_args
dataset.get_as(device='cpu')
(indices, translations), patterns = dataset[:]
dataset.get_as(*get_as_args[0],**get_as_args[1])
# Set to none to avoid issues with things outside the detector
if auto_center:
center = tools.image_processing.centroid(t.sum(patterns,dim=0))
else:
center = None
# Then, generate the probe geometry from the dataset
ewg = tools.initializers.exit_wave_geometry
probe_basis, ew_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
center=center,
padding=padding,
opt_for_fft=False,
oversampling=oversampling)
det_shape,
wavelength,
distance,
padding=padding,
opt_for_fft=False,
oversampling=oversampling)
if propagation_distances is None:
propagation_distances = [None] * len(datasets)
# This shrinks the probe to ensure that the output wavefield
# is the correct shape
probe_shape = t.Size(np.array(ew_shape) - (2*locality_radius))
# always do this on the cpu
probe_planes = []
translations = []
patterns = []
probes = []
for i, dataset in enumerate(datasets):
get_as_args = dataset.get_as_args
dataset.get_as(device='cpu')
(indices, tx), pats = dataset[:]
dataset.get_as(*get_as_args[0],**get_as_args[1])
translations.append(tx)
probe_planes.extend([i]*tx.shape[0])
patterns.append(pats)
# Finally, initialize the probe and object using this information
if locality_radius != 0:
probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distances[i], oversampling=oversampling)[locality_radius:-locality_radius,locality_radius:-locality_radius]
else:
probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distances[i], oversampling=oversampling)
# Now we initialize all the subdominant probe modes
probe_max = t.max(cmath.cabs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
probe = t.stack([tools.propagators.inverse_far_field(probe),] + probe_stack)
probes.append(probe)
translations = t.cat(translations)
patterns = t.cat(patterns)
probes = t.stack(probes)
if hasattr(dataset, 'sample_info') and \
dataset.sample_info is not None and \
'orientation' in dataset.sample_info:
surface_normal = dataset.sample_info['orientation'][2]
if hasattr(datasets[0], 'sample_info') and \
datasets[0].sample_info is not None and \
'orientation' in datasets[0].sample_info:
surface_normal = datasets[0].sample_info['orientation'][2]
else:
surface_normal = np.array([0.,0.,1.])
@@ -151,25 +212,12 @@ class SMatrixPtycho(CDIModel):
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200+2*locality_radius)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
background = t.sqrt(datasets[0].background)
else:
background = None
# Finally, initialize the probe and object using this information
if probe_size is None:
if locality_radius != 0:
probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)[locality_radius:-locality_radius,locality_radius:-locality_radius]
else:
probe = tools.initializers.SHARP_style_probe(dataset, ew_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)
# Now we initialize all the subdominant probe modes
probe_max = t.max(cmath.cabs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
probe = t.stack([tools.propagators.inverse_far_field(probe),] + probe_stack)
s_matrix = t.zeros([2*locality_radius+1,2*locality_radius+1,obj_size[0],
obj_size[1],2])
s_matrix[locality_radius,locality_radius,:,:,:] = \
@@ -179,35 +227,43 @@ class SMatrixPtycho(CDIModel):
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
translation_offsets = 0 * (t.rand((translations.shape[0],2)) - 0.5)
weights = t.ones(len(dataset))
weights = t.ones(translations.shape[0])
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
if hasattr(datasets[0], 'mask') and datasets[0].mask is not None:
mask = datasets[0].mask.to(t.bool)
else:
mask = None
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
xs, ys = np.mgrid[:probe.shape[1],:probe.shape[2]]
probe_support = t.zeros_like(probes[0,0].to(dtype=t.float32))
xs, ys = np.mgrid[:probes.shape[-3],:probes.shape[-2]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
Rs = np.sqrt(xs**2 + ys**2)
probe_support[Rs<probe_convergence_radius] = 1
probe = probe * probe_support[None,:,:]
probe_support[Rs<probe_convergence_radius] = 1
probes = probes * probe_support[None,None,:,:]
return cls(wavelength, det_geo, probe_basis, probe, probe_support,
s_matrix,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets = translation_offsets,
weights=weights, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
oversampling=oversampling)
model = cls(wavelength, det_geo, probe_basis, probes, probe_support,
s_matrix,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets = translation_offsets,
probe_planes = probe_planes,
weights=weights, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
oversampling=oversampling)
# Now we need to produce a concatenated dataset to be used to
# train the model
dataset = Ptycho2DDataset(translations, patterns)
return model, dataset
def interaction(self, index, translations):
@@ -218,16 +274,23 @@ class SMatrixPtycho(CDIModel):
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
all_exit_waves = []
for i in range(self.probe.shape[0]):
pr = tools.propagators.inverse_far_field(self.probe[i] * self.probe_fourier_support)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc_s_matrix(
pr, self.s_matrix, pix_trans, shift_probe=True)
exit_waves = exit_waves
if self.probe_planes is not None:
probes_set = self.probe[self.probe_planes[index]]
else:
probes_set = self.probe[[0]*translations.shape[0]]
all_exit_waves = []
for i in range(probes_set.shape[1]):
exit_waves = []
for j in range(probes_set.shape[0]):
pr = tools.propagators.inverse_far_field(probes_set[j,i] * self.probe_fourier_support)
exit_wave = self.probe_norm * tools.interactions.ptycho_2D_sinc_s_matrix(pr, self.s_matrix, pix_trans[j], shift_probe=True)
exit_waves.append(exit_wave)
exit_waves = t.stack(exit_waves)
if exit_waves.dim() == 4:
exit_waves = self.weights[index][:,None,None,None] * exit_waves
else:
@@ -274,13 +337,15 @@ class SMatrixPtycho(CDIModel):
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
if self.probe_planes is not None:
self.probe_planes = self.probe_planes.to(*args, **kwargs)
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.probe_norm = self.probe_norm.to(*args,**kwargs)
self.probe_fourier_support = self.probe_fourier_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
@@ -329,16 +394,16 @@ class SMatrixPtycho(CDIModel):
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
('Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
('Subdominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Subdominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('First Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0,0], fig=fig, basis=self.probe_basis)),
('First Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0,0], fig=fig, basis=self.probe_basis)),
('Second Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1,0], fig=fig, basis=self.probe_basis),
lambda self: self.probe.shape[0] >=2),
('Second Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1,0], fig=fig, basis=self.probe_basis),
lambda self: self.probe.shape[0] >=2),
('Exit Wave Amplitude under Uniform Illumination',
lambda self, fig: p.plot_amplitude(t.sum(self.s_matrix.data,dim=(0,1)), fig=fig, basis=self.probe_basis)),
('Exit Wave Phase under Uniform Illumination',
CDTools/tools/atoms.py
+1292
View File
File diff suppressed because it is too large Load Diff
CDTools/tools/interactions.py
+1 -2
View File
@@ -14,8 +14,7 @@ import numpy as np
__all__ = ['translations_to_pixel', 'pixel_to_translations',
'project_translations_to_sample',
'ptycho_2D_round','ptycho_2D_linear','ptycho_2D_sinc',
'ptycho_2D_propagate']
'ptycho_2D_round','ptycho_2D_linear','ptycho_2D_sinc']
CDTools/tools/propagators.py
+17 -3
View File
@@ -166,7 +166,10 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
# This could potentially correct for a mistake in the implied
# propagation direction (e.g. choosing e^ikx instead of e^-ikx)
#samp_det_vec *= -1
# This appears to be correct, based on empirical evidence from
# a grazing incidence reflection experiment at 10 degrees
# on the optical table
samp_det_vec *= -1
if lens == False:
# This correctly reproduces the sample-to-each-pixel vectors
@@ -247,7 +250,7 @@ def high_NA_far_field(wavefront, k_map, intensity_map=None):
for penetrating radiation - may either not need a correction or need
a different correction due to the volumetric nature of the pixels.
If the k-map map any pixels on the detector to pixels outside of the
If the k-map maps any pixels on the detector to pixels outside of the
k-space range of the wavefront, these will be set to zero. This is in
keeping with the typical assumption that the sample is band-limited to
the Nyquist frequency for the array on which it is sampled.
@@ -309,7 +312,7 @@ def high_NA_far_field(wavefront, k_map, intensity_map=None):
def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, **kwargs):
def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, bandlimit=None, **kwargs):
"""Generates an angular-spectrum based near-field propagator from experimental quantities
This function generates an angular-spectrum based near field
@@ -322,6 +325,11 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
Formally, this propagator is the complex conjugate of the fourier
transform of the convolution kernel for light propagation in free
space
If the optional bandlimit parameter is set, the propagator will be set
to zero beyond an explicit bandlimiting frequency. This is helpful if the
propagator will be used in a repeated multiply/propagate framework such
as a multislice algorithm, where it helps to prevent aliasing.
Parameters
----------
@@ -335,6 +343,8 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
The distance to simulate propagation over
remove_z_phase : bool
Default False, whether to remove the dominant z-direction phase dependence
bandlimit : float
Optional, a fraction of the full detector radius beyond which to set the propagator to zero.
Returns
-------
@@ -359,6 +369,10 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
if remove_z_phase:
propagator *= np.exp(-1j * k0 * z)
if bandlimit is not None:
Rs = np.sqrt((Ki / np.max(ki))**2 + (Kj / np.max(kj))**2)
propagator = propagator * (Rs < bandlimit)
# Take the conjugate explicitly here instead of negating
# the previous expression to ensure that complex frequencies
# get mapped to values <1 instead of >1