diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 166749e..fdfa7a4 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -39,6 +39,7 @@ import numpy as np import threading import queue import time +import pytorch_warmup __all__ = ['CDIModel'] @@ -105,7 +106,7 @@ class CDIModel(t.nn.Module): def AD_optimize(self, iterations, data_loader, optimizer,\ scheduler=None, regularization_factor=None, thread=True, - calculation_width=10): + calculation_width=10, warmup_scheduler=None): """Runs a round of reconstruction using the provided optimizer This is the basic automatic differentiation reconstruction tool @@ -135,7 +136,8 @@ class CDIModel(t.nn.Module): normalization = 0 for inputs, patterns in data_loader: normalization += t.sum(patterns).cpu().numpy() - + + def run_iteration(stop_event=None): loss = 0 N = 0 @@ -173,8 +175,20 @@ class CDIModel(t.nn.Module): loss = self.regularizer(regularization_factor) loss.backward() return total_loss - + + if warmup_scheduler is not None: + old_lrs = [group['lr'] for group in + optimizer.param_groups] + warmup_scheduler.dampen() + #print([group['lr'] for group in + # optimizer.param_groups], end='\r') + loss += optimizer.step(closure).detach().cpu().numpy() + + if warmup_scheduler is not None: + for old_lr, group in zip(old_lrs, optimizer.param_groups): + group['lr'] = old_lr + loss /= normalization if scheduler is not None: @@ -201,6 +215,7 @@ class CDIModel(t.nn.Module): self.figs[0].canvas.start_event_loop(0.01) else: calc.join() + except KeyboardInterrupt as e: stop_event.set() print('\nAsking execution thread to stop cleanly - please be patient.') @@ -215,7 +230,7 @@ class CDIModel(t.nn.Module): def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False, amsgrad=False, subset=None, regularization_factor=None, thread=True, - calculation_width=10): + calculation_width=10, warmup=False): """Runs a round of reconstruction using the Adam optimizer This is generally accepted to be the most robust algorithm for use @@ -266,8 +281,14 @@ class CDIModel(t.nn.Module): else: scheduler = None + if warmup: + warmup_scheduler = pytorch_warmup.UntunedLinearWarmup(optimizer) + else: + warmup_scheduler = None + return self.AD_optimize(iterations, data_loader, optimizer, scheduler=scheduler, + warmup_scheduler=warmup_scheduler, regularization_factor=regularization_factor, thread=thread, calculation_width=calculation_width) diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index c875939..9736c6a 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -30,7 +30,10 @@ class Multislice2DPtycho(CDIModel): bandlimit=None, subpixel=True, exponentiate_obj=True, - fourier_probe=False, units='um'): + fourier_probe=False, + apodization=None, + phase_only=False, + units='um'): super(Multislice2DPtycho,self).__init__() self.wavelength = t.Tensor([wavelength]) @@ -56,6 +59,7 @@ class Multislice2DPtycho(CDIModel): self.exponentiate_obj = exponentiate_obj self.fourier_probe = fourier_probe self.units = units + self.phase_only=phase_only if mask is None: self.mask = mask @@ -73,7 +77,7 @@ class Multislice2DPtycho(CDIModel): / 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]) @@ -97,6 +101,11 @@ class Multislice2DPtycho(CDIModel): self.probe_fourier_support = t.Tensor(probe_fourier_support).to(t.float32) + if apodization is not None: + self.apodization = t.Tensor(apodization).to(t.float32) + else: + self.apodization = None + self.oversampling = oversampling spacing = np.linalg.norm(self.probe_basis,axis=0) @@ -108,7 +117,7 @@ class Multislice2DPtycho(CDIModel): @classmethod - def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, 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): + def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, 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, apodize_prop=False, phase_only=False): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -157,7 +166,6 @@ class Multislice2DPtycho(CDIModel): 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) @@ -175,14 +183,17 @@ class Multislice2DPtycho(CDIModel): 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)) + 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)] + probe = t.stack([probe,] + probe_stack) + # For a Fourier space probe if fourier_probe: probe = tools.propagators.far_field(probe) - - # Now we initialize all the subdominant probe modes - probe_max = t.max(cmath.cabs(probe)) - probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)] - probe = t.stack([probe,] + probe_stack) + # Consider a different start if exponentiate_obj: @@ -193,7 +204,17 @@ class Multislice2DPtycho(CDIModel): if not replicate_slice: obj = t.stack([obj]*nz) - + if apodize_prop: + shape = probe.shape[-3:-1] + #window_x = (1+np.cos(np.pi+2*np.pi*np.arange(shape[0])/shape[0]))/2 + #window_y = (1+np.cos(np.pi+2*np.pi*np.arange(shape[1])/shape[1]))/2 + window_x = np.cos(-np.pi/2+np.pi*np.arange(shape[0])/shape[0]) + window_y = np.cos(-np.pi/2+np.pi*np.arange(shape[1])/shape[1]) + Wx, Wy = np.meshgrid(window_x,window_y, indexing='ij') + apodization=t.Tensor(Wx*Wy).to(dtype=probe.dtype) + else: + apodization=None + det_geo = dataset.detector_geometry translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) @@ -228,7 +249,9 @@ class Multislice2DPtycho(CDIModel): bandlimit=bandlimit, subpixel=subpixel, exponentiate_obj=exponentiate_obj, - units=units, fourier_probe=fourier_probe) + units=units, fourier_probe=fourier_probe, + apodization=apodization, + phase_only=phase_only) def interaction(self, index, translations): @@ -247,8 +270,12 @@ class Multislice2DPtycho(CDIModel): prs = self.probe*self.probe_fourier_support[None,:,:] # Here is where the mixing would happen, if it happened + if self.exponentiate_obj: - obj = cmath.cexpi(self.obj/self.nz) + if self.phase_only: + obj = cmath.expi(self.obj[...,0]) + else: + obj = cmath.cexpi(self.obj) else: obj = self.obj @@ -266,6 +293,10 @@ class Multislice2DPtycho(CDIModel): exit_waves = tools.interactions.ptycho_2D_round( exit_waves, obj, pix_trans, multiple_modes=True) + if self.apodization is not None: + exit_waves = exit_waves * self.apodization[None,:,:,None] + + elif self.obj.dim() == 4: # If separate slices @@ -277,6 +308,8 @@ class Multislice2DPtycho(CDIModel): exit_waves = tools.interactions.ptycho_2D_round( exit_waves, obj[i], pix_trans, multiple_modes=True) + if self.apodization is not None: + exit_waves = exit_waves * self.apodization[None,:,:,None] if i < self.nz-1: #on all but the last iteration exit_waves = tools.propagators.near_field( @@ -330,6 +363,8 @@ class Multislice2DPtycho(CDIModel): if self.mask is not None: self.mask = self.mask.to(*args, **kwargs) + if self.apodization is not None: + self.apodization = self.apodization.to(*args,**kwargs) self.min_translation = self.min_translation.to(*args,**kwargs) @@ -338,6 +373,7 @@ class Multislice2DPtycho(CDIModel): #self.probe_support = self.probe_support.to(*args,**kwargs) self.probe_fourier_support = self.probe_fourier_support.to(*args,**kwargs) + self.surface_normal = self.surface_normal.to(*args, **kwargs) self.as_prop = self.as_prop.to(*args, **kwargs) @@ -396,13 +432,13 @@ class Multislice2DPtycho(CDIModel): ('Probe Real Space Phase', lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)), ('Slice by Slice Real Part of T', - lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), + lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'), lambda self: self.exponentiate_obj), ('Slice by Slice Imaginary Part of T', lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), 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), + 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), ('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), @@ -411,13 +447,13 @@ class Multislice2DPtycho(CDIModel): 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), + 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), ('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), + 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), ('Corrected Translations', lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)), diff --git a/CDTools/tools/initializers/initializers.py b/CDTools/tools/initializers/initializers.py index 68487ef..aa7df9d 100644 --- a/CDTools/tools/initializers/initializers.py +++ b/CDTools/tools/initializers/initializers.py @@ -9,7 +9,8 @@ import numpy as np import torch as t __all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian', - 'gaussian_probe', 'SHARP_style_probe', 'RPI_spectral_init'] + 'gaussian_probe', 'SHARP_style_probe', 'RPI_spectral_init', + 'generate_subdominant_modes'] from CDTools.tools import cmath from CDTools.tools.propagators import * @@ -18,6 +19,7 @@ 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 * def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0, oversampling=1): @@ -438,3 +440,79 @@ def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, backgroun return cmath.complex_to_torch(z0).to(dtype=t.float32) + +def generate_subdominant_modes(dominant_mode, n_modes, circular=True): + """Generates guesses of subdominant modes based on spatial derivatives + + The idea here is that vibration is an extremely common cause of + spatial incoherence. This typically presents as subdominant modes that + look like the derivative of the dominant mode along various axes. + + Therefore, it is a good starting guess if the guess of the dominant + mode is reasonable. The output probes will all be normalized to have + the same intensity as the input dominant probe + + Parameters + ---------- + dominant_mode : array + The complex-valued dominant mode to work from + n_modes : int + The number of additional modes to create + circular : bool + Default True, whether to use circular modes (x+iy) or linear (x,y) + Returns + ------- + torch.Tensor + The complex-style tensor storing the probe guesses + """ + + # This generates a list of tuples, where each tuple (a,b) corresponds + # a term (kx^a)*(ky^b), or (kx+iky)^a*(kx-iky)^b for circular modes + def make_orders(n_orders): + if n_orders >1024: + raise KeyError('Are you sure you want this many orders?') + i = 1 + n = 0 + total = 0 + while True: + if total >= n_orders: + break + + yield (n,i-n) + total += 1 + if n