diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 5691059..b49da15 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -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 diff --git a/CDTools/models/bragg_2d_ptycho.py b/CDTools/models/bragg_2d_ptycho.py index e2a00fe..36372cb 100644 --- a/CDTools/models/bragg_2d_ptycho.py +++ b/CDTools/models/bragg_2d_ptycho.py @@ -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() diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index a64e50d..42ed20c 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -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: diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index db82260..8768cc9 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -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= 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 diff --git a/CDTools/models/pinhole_plane_ptycho.py b/CDTools/models/pinhole_plane_ptycho.py deleted file mode 100644 index 9cb9d3d..0000000 --- a/CDTools/models/pinhole_plane_ptycho.py +++ /dev/null @@ -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} diff --git a/CDTools/models/unified_mode_ptycho.py b/CDTools/models/unified_mode_ptycho.py deleted file mode 100644 index a2c9392..0000000 --- a/CDTools/models/unified_mode_ptycho.py +++ /dev/null @@ -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} diff --git a/CDTools/tools/analysis/analysis.py b/CDTools/tools/analysis/analysis.py index 4a9cb5c..a99723e 100644 --- a/CDTools/tools/analysis/analysis.py +++ b/CDTools/tools/analysis/analysis.py @@ -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 diff --git a/CDTools/tools/atoms/atoms.py b/CDTools/tools/atoms/atoms.py index db4212b..9493214 100644 --- a/CDTools/tools/atoms/atoms.py +++ b/CDTools/tools/atoms/atoms.py @@ -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 diff --git a/CDTools/tools/data/data.py b/CDTools/tools/data/data.py index 8aeafb7..b3f05c4 100644 --- a/CDTools/tools/data/data.py +++ b/CDTools/tools/data/data.py @@ -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 diff --git a/CDTools/tools/image_processing/image_processing.py b/CDTools/tools/image_processing/image_processing.py index c437cf5..b8a5826 100644 --- a/CDTools/tools/image_processing/image_processing.py +++ b/CDTools/tools/image_processing/image_processing.py @@ -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) diff --git a/CDTools/tools/initializers/initializers.py b/CDTools/tools/initializers/initializers.py index 9671a25..2bf04c7 100644 --- a/CDTools/tools/initializers/initializers.py +++ b/CDTools/tools/initializers/initializers.py @@ -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