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

get tests working

Browse Source
This commit is contained in:
Abe Levitan
2021-06-21 16:55:34 -04:00
parent 8f58242efe
commit 76e6d5efae
18 changed files with 291 additions and 1072 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/__init__.py
+3 -7
View File
@@ -20,18 +20,14 @@ defining a new ptychography model before attempting to do so.
"""
from __future__ import division, print_function, absolute_import
# I don't believe that __all__ really needed, but it's nice to define it
# to be explicit that import * is safe
__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI', 'UnifiedModePtycho']
#__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI']
from CDTools.models.base import CDIModel
from CDTools.models.simple_ptycho import SimplePtycho
from CDTools.models.fancy_ptycho import FancyPtycho
#from CDTools.models.pinhole_plane_ptycho import PinholePlanePtycho
#from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
#from CDTools.models.s_matrix_ptycho import SMatrixPtycho
#from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
#from CDTools.models.rpi import RPI
#from CDTools.models.unified_mode_ptycho import UnifiedModePtycho
CDTools/models/bragg_2d_ptycho.py
+40 -85
View File
@@ -4,7 +4,6 @@ import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import plotting as p
from CDTools.tools.propagators import generate_generalized_angular_spectrum_propagator as ggasp
from matplotlib import pyplot as plt
@@ -110,9 +109,9 @@ class Bragg2DPtycho(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)))
self.probe_norm = 1 * t.max(t.abs(probe_guess[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(t.abs(probe_guess).to(t.float32))
# Not strictly necessary but otherwise it will return
# a probe with the stuff outside of the support unchanged after
@@ -124,42 +123,44 @@ class Bragg2DPtycho(CDIModel):
# you look at the phase map
probe_guess[probe_guess == 0] = 0
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape)
else:
background = 1e-6 * t.ones(self.probe[0].shape[:-1])
background = 1e-6 * t.ones(self.probe[0].shape)
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
self.background = t.nn.Parameter(t.as_tensor(background,dtype=t.float32))
if weights is None:
self.weights = None
else:
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
self.weights = t.nn.Parameter(t.as_tensor(weights,dtype=t.float32))
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.as_tensor(translation_offsets,dtype=t.float32) /
translation_scale)
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_support = t.ones(self.probe[0].shape,dtype=t.bool)
if obj_support is not None:
self.obj_support = obj_support
self.obj.data = self.obj * obj_support
else:
self.obj_support = t.ones_like(self.obj)
self.obj_support = t.ones_like(self.obj, dtype=t.bool)
self.oversampling = oversampling
@@ -170,7 +171,7 @@ class Bragg2DPtycho(CDIModel):
# recall that here we always want the shape of the detector
# before it's cut down by the detector slice to match the
# physical detector region
probe_shape = self.probe[0].shape[:-1]
probe_shape = self.probe[0]
self.k_map, self.intensity_map = \
tools.propagators.generate_high_NA_k_intensity_map(
@@ -188,11 +189,11 @@ class Bragg2DPtycho(CDIModel):
# This propagator should be able to be multiplied by the propagation
# distance each time to get a propagator
self.universal_propagator = cmath.cphase(ggasp(self.probe.shape[1:],
self.universal_propagator = t.angle(ggasp(self.probe.shape[1:],
self.probe_basis, self.wavelength,
t.Tensor([0,0,self.wavelength/(2*np.pi)]),
propagation_vector=self.prop_dir,
dtype=t.float32,
dtype=t.complex64,
propagate_along_offset=True))
@@ -299,11 +300,11 @@ class Bragg2DPtycho(CDIModel):
# Now we initialize all the subdominant probe modes
probe_max = t.max(cmath.cabs(probe))
probe_max = t.max(t.abs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
probe = t.stack([probe,] + probe_stack)
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
obj = t.exp(1j*(randomize_ang * (t.rand(obj_size)-0.5)))
det_geo = dataset.detector_geometry
@@ -317,8 +318,8 @@ class Bragg2DPtycho(CDIModel):
mask = None
if probe_support_radius is not None:
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
p_cent = np.array(probe.shape[1:3]).astype(int) // 2
probe_support = t.zeros(probe[0].shape, dtype=t.bool)
p_cent = np.array(probe[0].shape).astype(int) // 2
psr = int(probe_support_radius)
probe_support[p_cent[0]-psr:p_cent[0]+psr,
p_cent[1]-psr:p_cent[1]+psr] = 1
@@ -329,7 +330,7 @@ class Bragg2DPtycho(CDIModel):
ro = restrict_obj
os = np.array(obj_size)
ps = np.array(probe_shape)
obj_support = t.zeros_like(obj.to(dtype=t.float32))
obj_support = t.zeros(obj.shape,dtype=t.bool)
obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
else:
@@ -372,49 +373,23 @@ class Bragg2DPtycho(CDIModel):
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
single_translation = False
if translations.dim() == 1:
translations = translations[None,:]
pix_trans = pix_trans[None,:]
single_translation = True
all_exit_waves = []
for i in range(self.probe.shape[0]):
pr = self.probe[i] * self.probe_support
exit_waves = []
for j in range(translations.size()[0]):
if self.propagate_probe:
# Minus sign is empirical
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')
#plt.imshow(np.abs(cmath.torch_to_complex(prop_pr.detach().cpu())))
#plt.show()
else:
prop_pr = pr
exit_waves.append(self.probe_norm *
tools.interactions.ptycho_2D_sinc(prop_pr,
self.obj_support * self.obj,
pix_trans[j],
shift_probe=True))
exit_waves = t.stack(exit_waves)
if not single_translation:
exit_waves = self.weights[index][:,None,None,None] * exit_waves
else:
exit_waves = self.weights[index] * exit_waves[0,:,:,:]
all_exit_waves.append(exit_waves)
Ws = self.weights[index]
prs = Ws[...,None,None,None] * self.probe
# Now we need to propagate each of the probes
return t.stack(all_exit_waves)
for j in range(prs.shape[0]):
propagator = t.exp(
1j*(props[j]*(2*np.pi)/self.wavelength)
* self.universal_propagator)
prs[j] = tools.propagators.near_field(prs[j], propagator)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj_support * self.obj,pix_trans,
shift_probe=True, multiple_modes=True)
return exit_waves
def forward_propagator(self, wavefields):
if self.correct_tilt:
@@ -514,32 +489,12 @@ class Bragg2DPtycho(CDIModel):
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
translations = dataset.translations.to(dtype=self.probe.real.dtype,
device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
# 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),
# ('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)),
# ('Corrected Translations',
# lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
# ('Background',
# lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
# ]
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig)),
@@ -565,9 +520,9 @@ class Bragg2DPtycho(CDIModel):
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
CDTools/models/fancy_ptycho.py
+1 -1
View File
@@ -179,7 +179,7 @@ class FancyPtycho(CDIModel):
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)
else:
CDTools/models/multislice_2d_ptycho.py
+75 -78
View File
@@ -1,10 +1,8 @@
from __future__ import division, print_function, absolute_import
import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath, analysis, image_processing
from CDTools.tools import analysis, image_processing
from CDTools.tools import plotting as p
from matplotlib import pyplot as plt
from datetime import datetime
@@ -37,7 +35,7 @@ class Multislice2DPtycho(CDIModel):
super(Multislice2DPtycho,self).__init__()
self.wavelength = t.Tensor([wavelength])
self.detector_geometry = copy(detector_geometry)
self.dz = -dz
self.dz = dz
self.nz = nz
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
@@ -69,23 +67,23 @@ class Multislice2DPtycho(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)))
self.probe_norm = 1 * t.max(t.abs(probe_guess[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(t.abs(probe_guess).to(t.float32))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape)
else:
background = 1e-6 * t.ones(self.probe[0].shape[:-1])
background = 1e-6 * t.ones(self.probe[0].shape)
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
self.background = t.nn.Parameter(t.as_tensor(background,dtype=t.float32))
if weights is None:
self.weights = None
@@ -94,38 +92,41 @@ class Multislice2DPtycho(CDIModel):
# weights and complex-valued per-mode weight matrices
if len(weights.shape) == 1:
# This is if it's just a list of numbers
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
self.weights = t.nn.Parameter(t.as_tensor(weights,
dtype=t.float32))
else:
# Now this is a matrix of weights, so we
if type(weights) == type(t.zeros(1)):
self.weights = t.nn.Parameter(weights.to(t.float32))
else:
self.weights = t.nn.Parameter(cmath.complex_to_torch(weights).to(t.float32))
self.weights = t.nn.Parameter(t.as_tensor(weights,
dtype=t.complex64))
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.as_tensor(translation_offsets,dtype=t.float32)/ translation_scale)
self.translation_scale = translation_scale
self.probe_support = t.Tensor(probe_support).to(t.float32)
if probe_support is not None:
self.probe_support = t.as_tensor(probe_support,dtype=t.bool)
else:
self.probe_support = None#t.ones_like(self.probe,dtype=t.bool)#None
self.oversampling = oversampling
spacing = np.linalg.norm(self.probe_basis,axis=0)
shape = np.array(self.probe.shape[1:-1])
shape = np.array(self.probe.shape[1:])
if prevent_aliasing:
shape *= 2
spacing /= 2
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_convergence_radius, probe_size=None, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True):
def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, probe_size=None, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -186,13 +187,10 @@ class Multislice2DPtycho(CDIModel):
background = None
# 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, oversampling=oversampling)
else:
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
probe = tools.initializers.STEM_style_probe(dataset, probe_shape, det_slice, probe_convergence_semiangle, propagation_distance=propagation_distance, oversampling=oversampling)
# Now we initialize all the subdominant probe modes
probe_max = t.max(cmath.cabs(probe))
probe_max = t.max(t.abs(probe))
if n_modes >=2:
probe_stack = list(0.01*tools.initializers.generate_subdominant_modes(probe,n_modes-1,circular=False))
#probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
@@ -207,9 +205,9 @@ class Multislice2DPtycho(CDIModel):
# Consider a different start
if exponentiate_obj:
obj = t.zeros(obj_size+(2,))
obj = t.zeros(obj_size, dtype=t.complex64)
else:
obj = tools.cmath.expi(t.zeros(obj_size))
obj = t.exp(1j*t.zeros(obj_size))
# If we will use a separate object per slice
if not replicate_slice:
obj = t.stack([obj]*nz)
@@ -241,14 +239,13 @@ class Multislice2DPtycho(CDIModel):
else:
mask = None
probe_support = t.zeros_like(probe[0])
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,:,:]
# probe_support = t.zeros(probe[0].shape,dtype=t.bool)
# xs, ys = np.mgrid[:probe.shape[-2],:probe.shape[-1]]
# xs = xs - np.mean(xs)
# ys = ys - np.mean(ys)
# Rs = np.sqrt(xs**2 + ys**2)
probe = probe
return cls(wavelength, det_geo, probe_basis, probe, obj, dz, nz,
detector_slice=det_slice,
@@ -258,7 +255,6 @@ class Multislice2DPtycho(CDIModel):
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
oversampling=oversampling,
bandlimit=bandlimit,
subpixel=subpixel,
@@ -277,15 +273,18 @@ class Multislice2DPtycho(CDIModel):
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
# This restricts the basis probes to stay within the probe support
basis_prs = self.probe * self.probe_support[...,:,:]
# This restricts the basis probes to stay within the probe support
if self.probe_support is not None:
basis_prs = self.probe * self.probe_support[...,:,:]
else:
basis_prs = self.probe
# For a Fourier-space probe, we take an IFT
if self.fourier_probe:
basis_prs = tools.propagators.inverse_far_field(basis_prs)
if self.prevent_aliasing:
pix_trans = pix_trans * 2
#pix_trans = pix_trans * 2
basis_prs = image_processing.fourier_upsample(basis_prs)
# Now we construct the probes for each shot from the basis probes
@@ -294,31 +293,30 @@ class Multislice2DPtycho(CDIModel):
if len(self.weights[0].shape) == 0:
# If a purely stable coherent illumination is defined
# No cmult because Ws is real in this case
prs = Ws[...,None,None,None,None] * basis_prs
prs = Ws[...,None,None,None] * basis_prs
else:
# If a frame-by-frame weight matrix is defined
# This takes the dot product of all the weight matrices with
# the probes. The output has dimensions of translation, then
# coherent mode index, then x,y, and then complex index
prs = t.sum(cmath.cmult(Ws[...,None,None,:], basis_prs),
axis=-4)
prs = t.sum(Ws[...,None,None] * basis_prs, axis=-3)
if self.exponentiate_obj:
if self.phase_only:
obj = cmath.expi(self.obj[...,0])
obj = t.exp(1j*self.obj.real)
else:
obj = cmath.cexpi(self.obj)
obj = t.exp(1j*self.obj)
else:
obj = self.obj
if self.prevent_aliasing:
obj = image_processing.fourier_upsample(obj)
#if self.prevent_aliasing:
# obj = image_processing.fourier_upsample(obj)
exit_waves = self.probe_norm * prs
for i in range(self.nz):
# If only one object slice
if self.obj.dim() == 3:
if self.obj.dim() == 2:
if i == 0 and self.subpixel:
# We only need to apply the subpixel shift to the first
# slice, because it shifts the probe
@@ -328,14 +326,9 @@ class Multislice2DPtycho(CDIModel):
else:
exit_waves = tools.interactions.ptycho_2D_round(
exit_waves, obj, pix_trans,
multiple_modes=True)
#if self.iteration_count >= 10:
# plt.figure()
# plt.imshow(np.abs(cmath.torch_to_complex(exit_waves[0,0].detach().cpu())))
# plt.show()
multiple_modes=True,upsample_obj=self.prevent_aliasing)
elif self.obj.dim() == 4:
elif self.obj.dim() == 3:
# If separate slices
if i == 0 and self.subpixel:
exit_waves = tools.interactions.ptycho_2D_sinc(
@@ -349,17 +342,20 @@ class Multislice2DPtycho(CDIModel):
if i < self.nz-1: #on all but the last iteration
exit_waves = tools.propagators.near_field(
exit_waves,self.as_prop)
#plt.imshow(t.abs(exit_waves[0,0].detach().cpu()))
#plt.show()
return exit_waves
def forward_propagator(self, wavefields):
if self.prevent_aliasing:
left = [self.probe.shape[-3]//2,self.probe.shape[-2]//2]
right = [self.probe.shape[-3]//2+self.probe.shape[-3],
self.probe.shape[-2]//2+self.probe.shape[-2]]
left = [self.probe.shape[-2]//2,self.probe.shape[-1]//2]
right = [self.probe.shape[-2]//2+self.probe.shape[-2],
self.probe.shape[-1]//2+self.probe.shape[-1]]
return tools.propagators.far_field(wavefields)[...,left[0]:right[0],
left[1]:right[1],:]
left[1]:right[1]]
else:
return tools.propagators.far_field(wavefields)
@@ -396,7 +392,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)
if self.probe_support is not None:
self.probe_support = self.probe_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
@@ -441,7 +438,7 @@ class Multislice2DPtycho(CDIModel):
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
translations = dataset.translations.to(dtype=self.probe.real.dtype,device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
@@ -449,7 +446,7 @@ class Multislice2DPtycho(CDIModel):
def get_rhos(self):
# If this is the general unified mode model
if self.weights.dim() >= 2:
Ws = cmath.torch_to_complex(self.weights.detach().cpu())
Ws = self.weights.detach().cpu().numpy()
rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
return rhos_out
# This is the purely incoherent case
@@ -473,10 +470,10 @@ class Multislice2DPtycho(CDIModel):
# to catch this case, but because it's so much simpler than the
# unified mode case I think it's appropriate
if self.weights.dim() == 1:
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = self.probe.detach().cpu().numpy()
ortho_probes = analysis.orthogonalize_probes(probe)
self.probe.data = cmath.complex_to_torch(ortho_probes).to(
device=self.probe.device,dtype=self.probe.dtype)
self.probe.data = t.as_tensor(ortho_probes,
device=self.probe.device,dtype=self.probe.dtype)
return
# This is for the unified mode case
@@ -487,7 +484,7 @@ class Multislice2DPtycho(CDIModel):
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = self.probe.detach().cpu().numpy()
ortho_probes, A = analysis.orthogonalize_probes(probe,
density_matrix=overall_rho,
keep_transform=True,
@@ -520,10 +517,10 @@ class Multislice2DPtycho(CDIModel):
new_Ws = np.array(new_Ws)
self.weights.data = cmath.complex_to_torch(new_Ws).to(
self.weights.data = t.as_tensor(new_Ws,
dtype=self.weights.dtype,device=self.weights.device)
self.probe.data = cmath.complex_to_torch(ortho_probes).to(
self.probe.data = t.as_tensor(ortho_probes,
device=self.probe.device,dtype=self.probe.dtype)
@@ -551,22 +548,22 @@ class Multislice2DPtycho(CDIModel):
lambda self: self.exponentiate_obj),
('Integrated Real Part of T',
lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
lambda self: self.exponentiate_obj),
lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3),
('Integrated Imaginary Part of T',
lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: self.exponentiate_obj),
lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3),
('Slice by Slice Amplitude of Object Function',
lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: not self.exponentiate_obj),
('Slice by Slice Phase of Object Function',
lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units,cmap='cividis'),
lambda self: not self.exponentiate_obj),
('Amplitude of Stacked Object Function',
lambda self, fig: p.plot_amplitude(reduce(cmath.cmult, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: not self.exponentiate_obj),
('Amplitude of Stacked Object Function',
lambda self, fig: p.plot_amplitude(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: (not self.exponentiate_obj) and self.obj.dim() >=3),
('Phase of Stacked Object Function',
lambda self, fig: p.plot_phase(reduce(cmath.cmult, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units,cmap='cividis'),
lambda self: not self.exponentiate_obj),
lambda self, fig: p.plot_phase(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
@@ -582,9 +579,9 @@ class Multislice2DPtycho(CDIModel):
else:
probe = self.probe
probe = cmath.torch_to_complex(probe.detach().cpu())
probe = probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
dz = self.dz
CDTools/models/pinhole_plane_ptycho.py
-384
View File
@@ -1,384 +0,0 @@
from __future__ import division, print_function, absolute_import
import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import plotting as p
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from copy import copy
__all__ = ['PinholePlanePtycho']
class PinholePlanePtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
pinhole_guess, obj_guess, illumination_guess,
pinhole_propagation,
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,
pinhole_support = None, obj_support=None, oversampling=1):
super(PinholePlanePtycho,self).__init__()
self.wavelength = t.Tensor([wavelength])
self.detector_geometry = copy(detector_geometry)
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = t.Tensor(det_geo['distance'])
if hasattr(det_geo, 'basis'):
det_geo['basis'] = t.Tensor(det_geo['basis'])
if hasattr(det_geo, 'corner'):
det_geo['corner'] = t.Tensor(det_geo['corner'])
self.min_translation = t.Tensor(min_translation)
self.probe_basis = t.Tensor(probe_basis)
self.detector_slice = detector_slice
self.surface_normal = t.Tensor(surface_normal)
self.saturation = saturation
if mask is None:
self.mask = mask
else:
self.mask = t.ByteTensor(mask)
# We rescale the pinhole here so it learns at the same rate as the
# object
self.pinhole_norm = 1 * t.max(tools.cmath.cabs(pinhole_guess.to(t.float32)))
self.pinhole = t.nn.Parameter(pinhole_guess.to(t.float32)
/ self.pinhole_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
self.illumination = t.nn.Parameter(illumination_guess.to(t.float32))
spacing = np.abs([self.probe_basis.numpy()[0,1], self.probe_basis.numpy()[1,0]])
self.pinhole_propagator = tools.propagators.generate_angular_spectrum_propagator(self.pinhole.shape[:-1], spacing, self.wavelength, pinhole_propagation).to(dtype=self.pinhole.dtype)
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(self.pinhole[self.detector_slice].shape[:-1])
else:
background = 1e-6 * t.ones(self.pinhole.shape[:-1])
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
if weights is None:
self.weights = None
else:
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
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_scale = translation_scale
if pinhole_support is not None:
self.pinhole_support = pinhole_support
else:
self.pinhole_support = t.ones_like(self.pinhole)
if obj_support is not None:
self.obj_support = obj_support
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, pinhole_propagation, pinhole_size=None, randomize_ang=0, padding=0, translation_scale = 1, saturation=None, pinhole_support_radius=None, propagation_distance=None, restrict_obj=-1, 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']
# 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
center = tools.image_processing.centroid(t.sum(patterns,dim=0))
# Then, generate the pinhole geometry from the dataset
ewg = tools.initializers.exit_wave_geometry
probe_basis, pinhole_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
center=center,
padding=padding,
opt_for_fft=False,
oversampling=oversampling)
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]
else:
surface_normal = np.array([0.,0.,1.])
# If this information is supplied when the function is called,
# then we override the information in the .cxi file
if scattering_mode in {'t', 'transmission'}:
surface_normal = np.array([0.,0.,1.])
elif scattering_mode in {'r', 'reflection'}:
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)
# 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(pinhole_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# Finally, initialize the pinhole and object using this information
if pinhole_size is None:
pinhole = tools.initializers.SHARP_style_probe(dataset, pinhole_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
else:
pinhole = tools.initializers.gaussian_probe(dataset, probe_basis, pinhole_shape, pinhole_size, propagation_distance=propagation_distance)
illumination = tools.cmath.expi(t.zeros(obj_size))
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
weights = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.uint8)
else:
mask = None
if pinhole_support_radius is not None:
pinhole_support = t.zeros_like(pinhole.to(dtype=t.float32))
p_cent = np.array(pinhole.shape[:2]).astype(int) // 2
psr = int(pinhole_support_radius)
pinhole_support[p_cent[0]-psr:p_cent[0]+psr,
p_cent[1]-psr:p_cent[1]+psr] = 1
pinhole = pinhole_support * pinhole
else:
pinhole_support = None;
if restrict_obj != -1:
ro = restrict_obj
os = np.array(obj_size)
ps = np.array(pinhole_shape)
obj_support = t.zeros_like(obj.to(dtype=t.float32))
obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
else:
obj_support = None
return cls(wavelength, det_geo, probe_basis, pinhole, obj,
illumination, pinhole_propagation,
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,
pinhole_support=pinhole_support,
obj_support=obj_support,
oversampling=oversampling)
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
pr = self.pinhole * self.pinhole_support
pinhole_exit_waves = self.pinhole_norm * \
tools.interactions.ptycho_2D_sinc(pr,
self.illumination,
pix_trans,
shift_probe=True)
sample_illumination = tools.propagators.near_field(pinhole_exit_waves,
self.pinhole_propagator)
exit_waves = []
if pinhole_exit_waves.dim() == 4:
for pix_tran, illum in zip(pix_trans, sample_illumination):
exit_waves.append(tools.interactions.ptycho_2D_sinc(illum,
self.obj_support * self.obj,
pix_tran,
shift_probe=True))
exit_waves = t.stack(exit_waves)
else:
exit_waves = tools.interactions.ptycho_2D_sinc(sample_illumination,
self.obj_support * self.obj,
pix_trans,
shift_probe=True)
if exit_waves.dim() == 4:
exit_waves = self.weights[index][:,None,None,None] * exit_waves
else:
exit_waves = self.weights[index] * exit_waves
return exit_waves
def forward_propagator(self, wavefields):
return tools.propagators.far_field(wavefields)
def backward_propagator(self, wavefields):
return tools.propagators.inverse_far_field(wavefields)
def measurement(self, wavefields):
return tools.measurements.quadratic_background(wavefields,
self.background,
detector_slice=self.detector_slice,
saturation=self.saturation,
oversampling=self.oversampling)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
def to(self, *args, **kwargs):
super(PinholePlanePtycho, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args,**kwargs)
# move the detector geometry too
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
if hasattr(det_geo, 'basis'):
det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
if hasattr(det_geo, 'corner'):
det_geo['corner'] = det_geo['corner'].to(*args,**kwargs)
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.pinhole_norm = self.pinhole_norm.to(*args,**kwargs)
self.pinhole_support = self.pinhole_support.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
self.pinhole_propagator = self.pinhole_propagator.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'CDTools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
surface_normal = self.surface_normal.detach().cpu().numpy()
xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
xsurfacevec /= np.linalg.norm(xsurfacevec)
ysurfacevec = np.cross(surface_normal, xsurfacevec)
ysurfacevec /= np.linalg.norm(ysurfacevec)
orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
sample_info = {'description': 'A simulated sample',
'orientation': orientation}
detector_geometry = self.detector_geometry
mask = self.mask
wavelength = self.wavelength
indices, translations = args_list
# Then we simulate the results
data = self.forward(indices, translations)
# And finally, we make the dataset
return Ptycho2DDataset(translations, data,
entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.pinhole.dtype,device=self.pinhole.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Pinhole Amplitude',
lambda self, fig: p.plot_amplitude(self.pinhole, fig=fig, basis=self.probe_basis)),
('Pinhole Phase',
lambda self, fig: p.plot_phase(self.pinhole, fig=fig, basis=self.probe_basis)),
('Illumination Amplitude',
lambda self, fig: p.plot_amplitude(self.illumination, fig=fig, basis=self.probe_basis),
lambda self: len(self.pinhole) >=2),
('Illumination Phase',
lambda self, fig: p.plot_phase(self.illumination, fig=fig, basis=self.probe_basis),
lambda self: len(self.pinhole) >=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)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
]
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.pinhole.detach().cpu())
probe = probe * self.pinhole_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
return {'basis':basis, 'translation':translations,
'probe':probe,'obj':obj,
'background':background,
'weights':weights}
CDTools/models/unified_mode_ptycho.py
-457
View File
@@ -1,457 +0,0 @@
from __future__ import division, print_function, absolute_import
import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import analysis
from CDTools.tools import plotting as p
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from copy import copy
__all__ = ['UnifiedModePtycho']
class UnifiedModePtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess, Ws_guess,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
translation_scale = 1, saturation=None,
probe_support = None, obj_support=None, oversampling=1):
super(UnifiedModePtycho,self).__init__()
self.wavelength = t.Tensor([wavelength])
self.detector_geometry = copy(detector_geometry)
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = t.Tensor(det_geo['distance'])
if hasattr(det_geo, 'basis'):
det_geo['basis'] = t.Tensor(det_geo['basis'])
if hasattr(det_geo, 'corner'):
det_geo['corner'] = t.Tensor(det_geo['corner'])
self.min_translation = t.Tensor(min_translation)
self.probe_basis = t.Tensor(probe_basis)
self.detector_slice = detector_slice
self.surface_normal = t.Tensor(surface_normal)
self.saturation = saturation
if mask is None:
self.mask = mask
else:
self.mask = t.BoolTensor(mask)
# 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)))
else:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
if background is None:
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))
if type(Ws_guess) == type(t.zeros(1)):
self.Ws = t.nn.Parameter(Ws_guess.to(t.float32))
else:
self.Ws = t.nn.Parameter(cmath.complex_to_torch(Ws_guess).to(t.float32))
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_scale = translation_scale
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
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, oversampling=1, auto_center=True, opt_for_fft=False, dm_rank=0):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
det_shape = dataset[0][1].shape
distance = dataset.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, probe_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
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]
else:
surface_normal = np.array([0.,0.,1.])
# If this information is supplied when the function is called,
# then we override the information in the .cxi file
if scattering_mode in {'t', 'transmission'}:
surface_normal = np.array([0.,0.,1.])
elif scattering_mode in {'r', 'reflection'}:
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(surface_normal)
# 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=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# 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, 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([probe,] + probe_stack)
#probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
# dm_rank defines the rank of the shot-by-shot density matrices
if dm_rank > n_modes:
raise KeyError('Density matrix rank cannot be greater than the number of modes')
elif dm_rank != 0:
if dm_rank == -1:
dm_rank = n_modes
Ws = t.zeros(len(dataset),dm_rank,n_modes,2)
Ws[:,0,0,0] = 1
for i in range(1,dm_rank):
Ws[:,i,i,0] = 1/np.sqrt(n_modes)
else:
# dm_rank=0 is a special case defining a purely stable, incoherent
# mode mixing model. This is passed on by defining a set of weights
# which only has one index
Ws = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
else:
mask = None
if probe_support_radius is not None:
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
p_cent = np.array(probe.shape[1:3]).astype(int) // 2
psr = int(probe_support_radius)
probe_support[p_cent[0]-psr:p_cent[0]+psr,
p_cent[1]-psr:p_cent[1]+psr] = 1
probe = probe * probe_support[None,:,:]
else:
probe_support = None;
if restrict_obj != -1:
ro = restrict_obj
os = np.array(obj_size)
ps = np.array(probe_shape)
obj_support = t.zeros_like(obj.to(dtype=t.float32))
obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
else:
obj_support = None
return cls(wavelength, det_geo, probe_basis, probe, obj, Ws,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets = translation_offsets,
mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
obj_support=obj_support,
oversampling=oversampling)
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
Ws = self.Ws[index]
# This restricts the basis probes with the probe support
basis_prs = self.probe * self.probe_support[...,:,:]
if len(self.Ws[0].shape) == 0:
# If a purely stable coherent illumination is defined
prs = cmath.cmult(Ws[...,None,None,None,:],basis_prs)
else:
# If a frame-by-frame weight matrix is defined
# This takes the dot product of all the weight matrices with
# the probes. The output has dimensions of translation, then
# coherent mode index, then x,y, and then complex index
prs = t.sum(cmath.cmult(Ws[...,None,None,:], basis_prs),
axis=-4)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj_support * self.obj,pix_trans,
shift_probe=True, multiple_modes=True)
exit_waves = exit_waves * self.probe_support[...,:,:]
return exit_waves
def forward_propagator(self, wavefields):
return tools.propagators.far_field(wavefields)
def backward_propagator(self, wavefields):
return tools.propagators.inverse_far_field(wavefields)
def measurement(self, wavefields):
return tools.measurements.quadratic_background(wavefields,
self.background,
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
def to(self, *args, **kwargs):
super(UnifiedModePtycho, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args,**kwargs)
# move the detector geometry too
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
if hasattr(det_geo, 'basis'):
det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
if hasattr(det_geo, 'corner'):
det_geo['corner'] = det_geo['corner'].to(*args,**kwargs)
if self.mask is not None:
self.mask = self.mask.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_support = self.probe_support.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'CDTools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
surface_normal = self.surface_normal.detach().cpu().numpy()
xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
xsurfacevec /= np.linalg.norm(xsurfacevec)
ysurfacevec = np.cross(surface_normal, xsurfacevec)
ysurfacevec /= np.linalg.norm(ysurfacevec)
orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
sample_info = {'description': 'A simulated sample',
'orientation': orientation}
detector_geometry = self.detector_geometry
mask = self.mask
wavelength = self.wavelength
indices, translations = args_list
# Then we simulate the results
data = self.forward(indices, translations)
# And finally, we make the dataset
return Ptycho2DDataset(translations, data,
entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def get_rhos(self):
# If this is not a purely stable model
if len(self.Ws.shape) >= 2:
Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
return rhos_out
else:
return np.array([np.eye(self.probe.shape[0])]*self.Ws.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
# Must also implement a version that works appropriately with
# a purely incoherent model
#
# Note to future: We could probably do this more cleanly with an
# SVD directly on the Ws matrix, instead of an eigendecomposition
# of the rho matrix. This could avoid potential stability issues
# due to the existence of zero eigenvalues in the full rho matrix
# when dm_rank < n_modes
#
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = cmath.torch_to_complex(self.probe.detach().cpu())
ortho_probes, A = analysis.orthogonalize_probes(probe,
density_matrix=overall_rho,
keep_transform=True,
normalize=True)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
dm_rank = self.Ws.shape[1]
new_Ws = []
for rho in new_rhos:
# These are returned from smalles to largest - we want to keep
# the largest ones
w,v = np.linalg.eigh(rho)
w = w[::-1][:dm_rank]
v = v[:,::-1][:,:dm_rank]
new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
new_Ws = np.array(new_Ws)
self.Ws.data = cmath.complex_to_torch(new_Ws).to(
dtype=self.Ws.dtype,device=self.Ws.device)
self.probe.data = cmath.complex_to_torch(ortho_probes).to(
device=self.probe.device,dtype=self.probe.dtype)
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Basis Probe Amplitudes',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
('Basis Probe Phases',
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig),
lambda self: len(self.Ws.shape) >=2),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig),
lambda self: len(self.Ws.shape) >=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)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
]
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
background = self.background.detach().cpu().numpy()**2
Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
return {'basis':basis, 'translation':translations,
'probe':probe,'obj':obj,
'background':background,
'Ws':Ws}
CDTools/tools/analysis/analysis.py
-1
View File
@@ -5,7 +5,6 @@ or numpy arrays, so they can be used either directly after reconstructions
on the attributes of the models themselves, or after-the-fact once the
data has been stored in numpy arrays.
"""
from __future__ import division, print_function
import torch as t
import numpy as np
CDTools/tools/atoms/atoms.py
+1 -1
View File
@@ -7,7 +7,7 @@ Appendix C. doi: 10.1007/978-1-4419-6533-2_11
They are actually taken from https://github.com/LeBeauGroup/pyMultislicer
written by Jim LeBeau, but above is the original reference.
"""
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from scipy import fftpack
CDTools/tools/data/data.py
-1
View File
@@ -5,7 +5,6 @@ specific desired information from a .cxi file. These functions should
handle all the needed conversions between standard formats (for example,
transposes of the basis arrays, shifting from object to probe motion, etc).
"""
from __future__ import division, print_function, absolute_import
import h5py
import numpy as np
CDTools/tools/image_processing/image_processing.py
+11 -3
View File
@@ -7,7 +7,6 @@ kind of tools perform common image manipulations on torch tensors, in such
a way that it is safe to include them in automatic differentiation models.
"""
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from CDTools.tools import propagators
@@ -139,6 +138,7 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
shift : torch.Tensor
The relative shift (i,j) needed to best map im1 onto im2
"""
#
# Here's my approach, perhaps it's a little unconventional. I will first
# calculate the phase correlation function as found in ____ (cite a paper
@@ -159,7 +159,7 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
try:
search_around = search_around.cpu()
except:
search_around = t.tensor(search_around)
search_around = t.as_tensor(search_around)
window_size = 15
shift_zero = tuple(-search_around + t.tensor([window_size,window_size]))
@@ -178,6 +178,7 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
# And we extract the shift from the window
print(upsampled.shape)
sh = t.as_tensor(upsampled.shape, device=upsampled.device)
cormax = t.as_tensor([t.div(t.argmax(upsampled), sh[1],
rounding_mode='floor'),
@@ -185,6 +186,9 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
device=upsampled.device)
sh_over_2 = t.div(sh,2,rounding_mode='floor')
print(cormax.shape)
print(sh.shape)
print(sh_over_2.shape)
subpixel_shift = ((cormax + sh_over_2) % sh - sh_over_2).to(dtype=upsampled.dtype)
return search_around.to(device=upsampled.device, dtype=upsampled.dtype) + \
@@ -301,7 +305,9 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True):
return conv_im
def fourier_upsample(ims):
def fourier_upsample(ims, preserve_mean=False):
# If preserve_mean is true, it preserves the mean pixel intensity
# otherwise, it preserves the total summed intensity
upsampled = t.zeros(ims.shape[:-2]+(2*ims.shape[-2],2*ims.shape[-1]),
dtype=ims.dtype,
device=ims.device)
@@ -310,6 +316,8 @@ def fourier_upsample(ims):
ims.shape[-1]//2+ims.shape[-1]]
upsampled[...,left[0]:right[0],left[1]:right[1]] = propagators.far_field(ims)
if preserve_mean:
upsampled *= 2
return propagators.inverse_far_field(upsampled)
CDTools/tools/initializers/initializers.py
+114 -16
View File
@@ -4,22 +4,22 @@ The functions in this module both do the geometric calculations needed to
initialize the reconstrucions, and the heuristic calculations for
geierating sensible initializations for the probe guess.
"""
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
'gaussian_probe', 'SHARP_style_probe', 'RPI_spectral_init',
'generate_subdominant_modes']
from CDTools.tools.propagators import *
from CDTools.tools.analysis import orthogonalize_probes
from CDTools.tools import image_processing
from scipy.fftpack import next_fast_len
from scipy.sparse import linalg as spla
from torch.nn.functional import pad
import numpy as np
from functools import *
__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
'gaussian_probe', 'SHARP_style_probe', 'STEM_style_probe',
'RPI_spectral_init',
'generate_subdominant_modes']
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
@@ -378,6 +378,104 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
return final_probe
def STEM_style_probe(dataset, shape, det_slice, convergence_semiangle, propagation_distance=None, oversampling=1):
"""Generates a STEM style probe guess from a dataset
What we call the "STEM" style probe guess is a probe generated by
a uniform aperture in Fourier space, with no optical aberrations.
This is the kind of probe than an ideal, aberration-corrected STEM
with perfect coherence would produce, so it is a good starting
guess for STEM datasets.
We set the initial intensity of the probe with relation to the
diffraction patterns so that a typical object will have an intensity
of around 1. We also set the initial phase ramp of the probe such that
the undiffracted probe has a centroid on the detector which matches the
centroid of the diffraction patterns
Parameters
----------
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
convergence_angle : float
The convergence angle of the probe, in mrad.
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
Returns
-------
torch.Tensor
The complex-style tensor storing the probe guess
"""
# The basis in the dataset should describe the basis of the probe's
# Fourier transform, even if the probe is simulated on larger stage than
# the detector. The only issue is if the probe is oversampled in
# Fourier space (simulated on a larger stage in real space). That factor
# is defined by oversampling.
probe_basis = dataset.detector_geometry['basis'] / oversampling
mean_im = t.mean(dataset.patterns,dim=0)
center = image_processing.centroid(mean_im)
Is = t.arange(shape[0], dtype=t.float32) - center[0]
Js = t.arange(shape[1], dtype=t.float32) - center[1]
Is,Js = t.meshgrid(Is,Js)
Rs = t.tensordot(probe_basis,t.stack([Is,Js]),dims=1)
forward = t.Tensor([0,0,1])
Rs = (forward * dataset.detector_geometry['distance'])[:,None,None] + Rs
dirs = Rs / t.linalg.norm(Rs,dim=(0,))
angles = t.acos(t.tensordot(forward,dirs,dims=1)) * 1000 #in mrad
probe_fft = t.zeros(shape,dtype=t.complex64)
probe_fft[angles<convergence_semiangle] = 1
probe_mean = t.mean(t.abs(probe_fft)**2)
diff_mean = t.mean(mean_im)
probe_fft = probe_fft * t.sqrt(diff_mean / probe_mean)
probe_guess = inverse_far_field(probe_fft)
if propagation_distance is not None:
# It's probably worth checking if this is correct when oversampling
# is not 1
probe_shape = t.as_tensor(tuple(probe_guess.shape))
# Start by recalculating the probe basis from the given information
det_basis = t.as_tensor(dataset.detector_geometry['basis'])
basis_dirs = det_basis / t.norm(det_basis, dim=0)
distance = dataset.detector_geometry['distance']
probe_basis = basis_dirs * dataset.wavelength * distance / \
(probe_shape * t.norm(det_basis,dim=0))
# Then package everything as it's needed
probe_spacing = t.norm(probe_basis,dim=0).numpy()
probe_shape = probe_shape.numpy().astype(np.int32)
# And generate the propagator
AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance)
probe_guess = near_field(probe_guess,AS_prop)
# Finally, place this probe in a full-sized array if there is oversampling
final_probe = t.zeros(shape,dtype=t.complex64)
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
def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, background=None):
# First, check if the probe is a single mode or many modes.
@@ -495,27 +593,27 @@ def generate_subdominant_modes(dominant_mode, n_modes, circular=True):
dominant_fft = far_field(dominant_mode)
shape = dominant_mode.shape
center = ((shape[-3]-1)//2, (shape[-2]-1)//2)
center = ((shape[-2]-1)//2, (shape[-1]-1)//2)
i, j = np.mgrid[:shape[-3], :shape[-2]]
i, j = np.mgrid[:shape[-2], :shape[-1]]
i = t.tensor(i - center[0]).to(dtype=dominant_fft.dtype,
device=dominant_fft.device)
j = t.tensor(j - center[1]).to(dtype=dominant_fft.dtype,
device=dominant_fft.device)
if circular:
a = t.stack((i,j),dim=-1)
b = t.stack((i,-j),dim=-1)
a = i + 1j * j
b = i - 1j * j
else:
a = t.stack((i,t.zeros_like(i)),dim=-1)
b = t.stack((j,t.zeros_like(j)),dim=-1)
a = i
b = j
probe_norm = t.sum(cmath.cabssq(dominant_mode))
probe_norm = t.sum(t.abs(dominant_mode)**2)
# Then we need to multiply that FFT by various phase masks and IFFT
probes = []
for a_order,b_order in make_orders(n_modes):
mask = reduce(cmath.cmult,[a]*a_order+[b]*b_order)
new_probe = inverse_far_field(cmath.cmult(mask,dominant_fft))
probes.append(new_probe * (probe_norm / t.sum(cmath.cabs(new_probe))))
mask = reduce(t.mul,[a]*a_order+[b]*b_order)
new_probe = inverse_far_field(mask * dominant_fft)
probes.append(new_probe * (probe_norm / t.sum(t.abs(new_probe))))
return t.stack(probes)
CDTools/tools/interactions/interactions.py
+18 -15
View File
@@ -5,11 +5,9 @@ function simulates some aspect of an interaction model that can be used
for ptychographic reconstruction.
"""
from __future__ import division, print_function, absolute_import
import torch as t
import numpy as np
from CDTools.tools import propagators
from CDTools.tools import propagators, image_processing
__all__ = ['translations_to_pixel', 'pixel_to_translations',
'project_translations_to_sample',
@@ -225,7 +223,7 @@ def project_translations_to_sample(sample_basis, translations):
def ptycho_2D_round(probe, obj, translations, multiple_modes=False):
def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj=False):
"""Returns a stack of exit waves without accounting for subpixel shifts
This function returns a collection of exit waves, with the first
@@ -265,10 +263,20 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False):
integer_translations = t.round(translations).to(dtype=t.int32)
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
if upsample_obj:
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2]//2,
tr[1]:tr[1]+probe.shape[-1]//2]
for tr in integer_translations])
selections = image_processing.fourier_upsample(selections,
preserve_mean=True)
else:
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
if multiple_modes:
# if the probe dimension is 4, then this hasn't yet been broadcast
# over the translation dimensions
@@ -445,26 +453,22 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
-subpixel_translations[:,1,None,None]*J))
fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2))
if probe.dim == 3: # Multi-mode probe
shifted_fft_probe = fft_probe * phase_masks[:,None,:,:]
if multiple_modes: # Multi-mode probe
shifted_fft_probe = fft_probe * phase_masks[...,None,:,:]
else:
shifted_fft_probe = fft_probe * phase_masks
shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
dim=(-1,-2)))
print('p',probe.shape)
print('fftp',fft_probe.shape)
print('sp',shifted_probe.shape)
print('sel',selections.shape)
if probe.dim == 3: # Multi-mode probe
output = shifted_probe * selections[:,None,:,:]
if multiple_modes: # Multi-mode probe
output = shifted_probe * selections[...,None,:,:]
else:
output = shifted_probe * selections
else:
raise NotImplementedError('Object shift not yet implemented')
print(output.shape)
if single_translation:
return output[0]
else:
@@ -532,7 +536,6 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
J = 2 * np.pi * J.to(t.float32) / probe.shape[-1]
I = I.to(dtype=probe.dtype,device=probe.device)
J = J.to(dtype=probe.dtype,device=probe.device)
print('hi')
for tr, sp in zip(integer_translations,
subpixel_translations):
CDTools/tools/losses/losses.py
+1 -2
View File
@@ -5,9 +5,8 @@ that returns the mean squared intensity error, and one that returns the
maximum likelihood metric for a system with Poisson statistics.
"""
from __future__ import division, print_function, absolute_import
import torch as t
import torch as t
__all__ = ['amplitude_mse', 'intensity_mse', 'poisson_nll']
CDTools/tools/measurements/measurements.py
-2
View File
@@ -5,8 +5,6 @@ model. There exist tools to simulate detectors with finite saturation
thresholds, backgrounds, and more.
"""
from __future__ import division, print_function, absolute_import
import torch as t
import numpy as np
from torch.nn.functional import avg_pool2d
CDTools/tools/plotting/plotting.py
-2
View File
@@ -6,8 +6,6 @@ and for after-the-fact analysis. Utilities for plotting complex valued
images exist, as well as plotting scan patterns and nanomaps
"""
from __future__ import division, print_function, absolute_import
import torch as t
import numpy as np
import matplotlib.pyplot as plt
CDTools/tools/propagators/propagators.py
+23 -13
View File
@@ -3,7 +3,6 @@
All the functions here are designed for use in an automatic differentiation
ptychography model. Each function implements a different propagator.
"""
from __future__ import division, print_function, absolute_import
import torch as t
from torch.nn.functional import grid_sample
@@ -316,11 +315,9 @@ 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, bandlimit=None, **kwargs):
"""Generates an angular-spectrum based near-field propagator from experimental quantities
This function generates an angular-spectrum based near field
propagator that will work on torch Tensors. The function is structured
this way - to generate the propagator first - because the
@@ -335,7 +332,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
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
----------
shape : array
@@ -358,13 +355,20 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
"""
# Internally, the generalized propagation function is used, so we start
# by creating an appropriate basis
basis = t.zeros([3,2], **kwargs).real
# This creates a real-valued tensor which matches the kind of complex
# number dtype requested in **kwargs
if 'dtype' in kwargs:
basis = t.real(t.zeros([3,2], **kwargs))
else:
basis = t.zeros([3,2], dtype=t.float32)
spacing = t.as_tensor(spacing, dtype=basis.dtype)
basis[1,0] = -spacing[0]
basis[0,1] = -spacing[1]
# And similarly, the offset is just z along the z direction
offset = t.zeros([3],**kwargs).real
offset = t.zeros([3], dtype=basis.dtype)
offset[2] = z
# And we call the generalized function!
propagator = generate_generalized_angular_spectrum_propagator(shape, basis,
wavelength, offset,
@@ -478,6 +482,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
# First we calculate a dual basis for the real space grid
inv_basis = t.linalg.pinv(basis).transpose(0,1)
# Then we calculate the frequencies in (i,j) space
ki = 2 * np.pi * t.fft.fftfreq(shape[0], dtype=inv_basis.dtype)
kj = 2 * np.pi * t.fft.fftfreq(shape[1], dtype=inv_basis.dtype)
@@ -487,10 +492,16 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
# These frequencies span the 2D plane of the input wavefield,
# hence K_ip for "in-plane"
K_ip = t.tensordot(inv_basis, K_ij, dims=1)
# Now, we need to generate the out-of-plane direction, so we can
# expand these Ks to the full Ks in 3D reciprocal space.
perpendicular_dir = t.cross(basis[:,1],basis[:,0])
# THis is broken down into steps to avoid floating point underflow
# which was a real problem that showed up for electron ptycho
b1_dir = basis[:,0] / t.linalg.norm(basis[:,0])
b2_dir = basis[:,1] / t.linalg.norm(basis[:,1])
perpendicular_dir = t.cross(b1_dir,b2_dir)
perpendicular_dir /= t.linalg.norm(perpendicular_dir)
# We set the sign of the propagation direction appropriately
@@ -514,7 +525,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
# Now, we have accurate in-plane values for K, so we can calculate the
# out-of-plane part. We start by calculating it's squared magnitude
K_oop_squared = (2*np.pi/wavelength)**2 - t.linalg.norm(K_ip,dim=0)**2
# Then, we take the square root and assign it the appropriate direction,
# adding to get the full 3D wavevectors. Note that we convert to complex
# before the square root to appropriately map negative numbers to
@@ -522,7 +533,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
K = K_ip + perpendicular_dir[:,None,None] \
* t.sqrt(t.complex(K_oop_squared,t.zeros_like(K_oop_squared)))
#
#
# In this section, we take the inner product of the calcualted
# wavevectors with the offset vector, to get the phase shift
# experienced by each plane wave.
@@ -535,7 +546,6 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
# experienced by K_0 to 0. K_0 will already have been set to 0 if
# there was no propagation vector set.
K_m_K_0 = K - K_0[:,None,None]
# We actually calculate the phase mask
phase_mask = t.tensordot(offset_vector,K_m_K_0, dims=1)
examples/gold_ball_ptycho.py
+1 -1
View File
@@ -17,7 +17,7 @@ model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2,dm_rank=0, pr
model.to(device='cuda')
dataset.get_as(device='cuda')
for loss in model.Adam_optimize(400, dataset, batch_size=50, schedule=True):
for loss in model.Adam_optimize(100, dataset, batch_size=50, schedule=True):
# And we liveplot the updates to the model as they happen
print(model.report())
model.inspect(dataset)
tests/tools/test_image_processing.py
+3 -3
View File
@@ -88,8 +88,7 @@ def test_find_subpixel_shift():
shift = t.tensor((0.8,0.75))
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift)
print(im.shape)
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift, multiple_modes=False)
retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0,0), resolution=50)
# tolerance of 0.03 on this measurement
@@ -104,7 +103,8 @@ def test_find_shift():
shift = t.tensor((0.8,0.75))
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift)[:-40,:-6]
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift,
multiple_modes=False)[:-40,:-6]
retrieved_shift = image_processing.find_shift(im, test_probe[40:,6:], resolution=50)
# tolerance of 0.03 on this measurement