diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index 794d192..ce46fc4 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -337,7 +337,11 @@ class Ptycho2DDataset(CDataset): idx = slider.val - 1 elif event.key == 'down' or event.button == 'down' or event.key == 'left': idx = slider.val + 1 - + else: + # This prevents errors from being thrown on irrelevant key + # or mouse input + return + # Handle the wraparound and trigger the update idx = int(idx) % len(self) slider.set_val(idx) diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 58af8be..8604d7b 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -24,7 +24,7 @@ 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'] +__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI', 'UnifiedModePtycho'] from CDTools.models.base import CDIModel from CDTools.models.simple_ptycho import SimplePtycho @@ -34,3 +34,5 @@ from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho 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 +from CDTools.models.unified_mode_ptycho2 import UnifiedModePtycho2 diff --git a/CDTools/models/base.py b/CDTools/models/base.py index e5fb9c3..cefa169 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -36,6 +36,9 @@ from matplotlib import pyplot as plt from matplotlib.widgets import Slider from matplotlib import ticker import numpy as np +import threading +import queue +import time __all__ = ['CDIModel'] @@ -97,7 +100,7 @@ class CDIModel(t.nn.Module): raise NotImplementedError() def AD_optimize(self, iterations, data_loader, optimizer,\ - scheduler=None, regularization_factor=None): + scheduler=None, regularization_factor=None, thread=True): """Runs a round of reconstruction using the provided optimizer This is the basic automatic differentiation reconstruction tool @@ -118,44 +121,75 @@ class CDIModel(t.nn.Module): Optional, a learning rate scheduler to use regularization_factor : float or list(float) Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method + thread : bool + Default True, whether to run the computation in a separate thread to allow interaction with plots during computation """ # First, calculate the normalization normalization = 0 for inputs, patterns in data_loader: normalization += t.sum(patterns).cpu().numpy() - - for it in range(iterations): + + def run_iteration(stop_event=None): loss = 0 N = 0 for inputs, patterns in data_loader: N += 1 def closure(): + # This is just used to allow graceful exit when threading + if stop_event is not None and stop_event.is_set(): + exit() optimizer.zero_grad() sim_patterns = self.forward(*inputs) if hasattr(self, 'mask'): loss = self.loss(patterns,sim_patterns, mask=self.mask) else: loss = self.loss(patterns,sim_patterns) - + if regularization_factor is not None \ and hasattr(self, 'regularizer'): loss += self.regularizer(regularization_factor) - + #print(loss) loss.backward() return loss - + loss += optimizer.step(closure).detach().cpu().numpy() - + loss /= normalization if scheduler is not None: scheduler.step(loss) - - yield loss + + return loss + + if thread: + result_queue = queue.Queue() + stop_event = threading.Event() + def target(): + result_queue.put(run_iteration(stop_event)) + + for it in range(iterations): + if thread: + calc = threading.Thread(target=target, name='calculator', daemon=True) + try: + calc.start() + while calc.is_alive(): + if hasattr(self, 'figs'): + 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.') + calc.join() + raise e + + yield result_queue.get() + else: + yield run_iteration() def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False, amsgrad=False, subset=None, - regularization_factor=None): + regularization_factor=None, thread=True): """Runs a round of reconstruction using the Adam optimizer This is generally accepted to be the most robust algorithm for use @@ -179,6 +213,8 @@ class CDIModel(t.nn.Module): Optional, a pattern index or list of pattern indices to use regularization_factor : float or list(float) Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method + thread : bool + Default True, whether to run the computation in a separate thread to allow interaction with plots during computation """ if subset is not None: @@ -203,12 +239,13 @@ class CDIModel(t.nn.Module): return self.AD_optimize(iterations, data_loader, optimizer, scheduler=scheduler, - regularization_factor=regularization_factor) + regularization_factor=regularization_factor, + thread=thread) def LBFGS_optimize(self, iterations, dataset, batch_size=None, lr=0.1,history_size=2, subset=None, - regularization_factor=None): + regularization_factor=None, thread=True): """Runs a round of reconstruction using the L-BFGS optimizer This algorithm is often less stable that Adam, however in certain @@ -232,6 +269,8 @@ class CDIModel(t.nn.Module): Optional, a pattern index or list of pattern indices to ues regularization_factor : float or list(float) Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method + thread : bool + Default True, whether to run the computation in a separate thread to allow interaction with plots during computation """ if subset is not None: # if just one pattern, turn into a list for convenience @@ -252,12 +291,14 @@ class CDIModel(t.nn.Module): lr = lr, history_size=history_size) return self.AD_optimize(iterations, data_loader, optimizer, - regularization_factor=regularization_factor) + regularization_factor=regularization_factor, + thread=thread) def SGD_optimize(self, iterations, dataset, batch_size=None, lr=0.01, momentum=0, dampening=0, weight_decay=0, - nesterov=False, subset=None, regularization_factor=None): + nesterov=False, subset=None, regularization_factor=None, + thread=True): """Runs a round of reconstruction using the SGDoptimizer This algorithm is often less stable that Adam, but it is simpler @@ -279,6 +320,8 @@ class CDIModel(t.nn.Module): Optional, a pattern index or list of pattern indices to use regularization_factor : float or list(float) Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method + thread : bool + Default True, whether to run the computation in a separate thread to allow interaction with plots during computation """ if subset is not None: @@ -303,7 +346,8 @@ class CDIModel(t.nn.Module): nesterov=nesterov) return self.AD_optimize(iterations, data_loader, optimizer, - regularization_factor=regularization_factor) + regularization_factor=regularization_factor, + thread=thread) # By default, the plot_list is empty @@ -337,6 +381,7 @@ class CDIModel(t.nn.Module): Whether to update existing plots or plot new ones """ + first_update = False if update and hasattr(self, 'figs') and self.figs: figs = self.figs @@ -375,10 +420,10 @@ class CDIModel(t.nn.Module): try: plotter(self, fig, dataset) plt.title(name) - except (IndexError, KeyError, AttributeError) as e: + except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e: pass - except (IndexError, KeyError, AttributeError) as e: + except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e: pass idx += 1 @@ -386,7 +431,7 @@ class CDIModel(t.nn.Module): if update: plt.draw() fig.canvas.start_event_loop(0.001) - + if first_update: plt.pause(0.05 * len(self.figs)) diff --git a/CDTools/models/bragg_2d_ptycho.py b/CDTools/models/bragg_2d_ptycho.py index 2dfec4c..e2a00fe 100644 --- a/CDTools/models/bragg_2d_ptycho.py +++ b/CDTools/models/bragg_2d_ptycho.py @@ -170,12 +170,13 @@ 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 - det_shape = self.probe[0].shape[:-1] - + probe_shape = self.probe[0].shape[:-1] + self.k_map, self.intensity_map = \ tools.propagators.generate_high_NA_k_intensity_map( - self.probe_basis, self.detector_geometry['basis'], - det_shape, + self.probe_basis, + self.detector_geometry['basis']/ oversampling, + probe_shape, self.detector_geometry['distance'], self.wavelength,dtype=t.float32, lens=lens) diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 1bfb3d3..936ef1d 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -144,7 +144,7 @@ class FancyPtycho(CDIModel): outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1]) outgoing_dir /= np.linalg.norm(outgoing_dir) surface_normal = outgoing_dir + np.array([0.,0.,1.]) - surface_normal /= np.linalg.norm(surface_normal) + surface_normal /= -np.linalg.norm(surface_normal) # Next generate the object geometry from the probe geometry and diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index f529189..0032044 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -31,7 +31,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'): @@ -95,22 +95,28 @@ class Multislice2DPtycho(CDIModel): if obj_support is not None: self.obj_support = obj_support - self.obj.data = self.obj * obj_support + if self.obj.dim() == 3: + self.obj.data = self.obj * obj_support + elif self.obj.dim() == 4: + self.obj.data = self.obj * obj_support[None,...] else: - self.obj_support = t.ones_like(self.obj) - + if self.obj.dim() == 3: + self.obj_support = t.ones_like(self.obj) + elif self.obj.dim() == 4: + self.obj_support = t.ones_like(self.obj[0]) self.oversampling = oversampling spacing = np.linalg.norm(self.probe_basis,axis=0) shape = np.array(self.probe.shape[1:-1]) self.bandlimit = bandlimit - + + # Big question: Should there be a minus sign before self.dz, or not? 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, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=4/5): + def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=4/5, replicate_slice=False): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -185,8 +191,14 @@ class Multislice2DPtycho(CDIModel): 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)) - + # Consider a different start + obj = t.zeros(obj_size+(2,)) + #obj = tools.cmath.expi(t.zeros(obj_size)) + # If we will use a separate object per slice + if not replicate_slice: + obj = t.stack([obj]*nz) + + det_geo = dataset.detector_geometry translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) @@ -212,7 +224,10 @@ class Multislice2DPtycho(CDIModel): ro = restrict_obj os = np.array(obj_size) ps = np.array(probe_shape) - obj_support = t.zeros_like(obj.to(dtype=t.float32)) + if replicate_slice: + obj_support = t.zeros_like(obj.to(dtype=t.float32)) + else: + obj_support = t.zeros_like(obj[0].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: @@ -251,7 +266,13 @@ class Multislice2DPtycho(CDIModel): if self.translation_offsets is not None: pix_trans += self.translation_scale * self.translation_offsets[index] - + if len(pix_trans.shape) == 1: + pix_trans = [pix_trans] + index = [index] + strip_first_index = True + else: + strip_first_index = False + all_exit_waves = [] for i in range(self.probe.shape[0]): # For a Fourier-space probe @@ -263,45 +284,32 @@ class Multislice2DPtycho(CDIModel): #print(self.probe_norm) #for i in range(self.nz): exit_waves = [] - if len(pix_trans.shape) == 1: - pix_trans = [pix_trans] + for trans in pix_trans: exit_wave = self.probe_norm * pr - for i in range(self.nz-1): - - exit_wave = tools.interactions.ptycho_2D_sinc(exit_wave, - self.obj_support * cmath.cexpi(self.obj/self.nz),#self.obj.data, - trans, - shift_probe=True) - #exit_wave = tools.interactions.ptycho_2D_round(exit_wave, - # self.obj_support * cmath.cexpi(self.obj.data/self.nz), - # trans) - exit_wave = tools.propagators.near_field(exit_wave,self.as_prop) + for i in range(self.nz): + # If only one object slice + if self.obj.dim() == 3: + exit_wave = tools.interactions.ptycho_2D_sinc( + exit_wave, + self.obj_support*cmath.cexpi(self.obj/self.nz), + trans, shift_probe=True) + # If separate slices + elif self.obj.dim() == 4: + exit_wave = tools.interactions.ptycho_2D_sinc( + exit_wave, + self.obj_support*cmath.cexpi(self.obj[i]/self.nz), + trans, shift_probe=True) + + exit_wave = tools.propagators.near_field( + exit_wave,self.as_prop) #tools.plotting.plot_amplitude(exit_wave) #plt.show() - - # only final layer gets a derivative - exit_wave = tools.interactions.ptycho_2D_sinc(exit_wave, - self.obj_support * cmath.cexpi(self.obj/self.nz),#self.obj, - trans, - shift_probe=True) - #exit_wave = tools.interactions.ptycho_2D_round(exit_wave, - # self.obj_support * cmath.cexpi(self.obj/self.nz), - # trans) - # One final propagation to enforce the bandlimit - exit_wave = tools.propagators.near_field(exit_wave,self.as_prop) - exit_waves.append(exit_wave) + exit_waves.append(exit_wave) exit_waves = t.stack(exit_waves) - - if np.array(index).size == 1: - index = [index] - strip_first_index = True - else: - strip_first_index = False - if exit_waves.dim() == 4: exit_waves = self.weights[index][:,None,None,None] * exit_waves else: @@ -312,7 +320,6 @@ class Multislice2DPtycho(CDIModel): all_exit_waves.append(exit_waves) - return t.stack(all_exit_waves) @@ -422,14 +429,10 @@ class Multislice2DPtycho(CDIModel): ('Subdominant Probe Phase', lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), lambda self: len(self.probe) >=2), - #('Object Amplitude', - # lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), - #('Object Phase', - # lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)), - ('Real Part of T', - lambda self, fig: p.plot_amplitude(self.obj[:,:,0].detach().cpu().numpy(), fig=fig, basis=self.probe_basis)), - ('Imaginary Part of T', - lambda self, fig: p.plot_amplitude(self.obj[:,:,1].detach().cpu().numpy(), fig=fig, basis=self.probe_basis)), + ('Integrated Real Part of T', + lambda self, fig: p.plot_real(t.mean(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis)), + ('Integrated Imaginary Part of T', + lambda self, fig: p.plot_imag(t.mean(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis)), ('Corrected Translations', lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)), ('Background', diff --git a/CDTools/models/unified_mode_ptycho.py b/CDTools/models/unified_mode_ptycho.py new file mode 100644 index 0000000..afaa19d --- /dev/null +++ b/CDTools/models/unified_mode_ptycho.py @@ -0,0 +1,449 @@ +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, rhos_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)) + + self.rhos = t.nn.Parameter(t.Tensor(rhos_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, mixing_mode='unified'): + + 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) + + # + if mixing_mode.lower().strip() == 'unified': + rhos = t.zeros(len(dataset),n_modes,n_modes) + rhos[:,0,0] = 1 + for i in range(1,n_modes): + rhos[:,i,i] = 1/n_modes + + 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, rhos, + 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] + + probes = [] + all_exit_waves = [] + for i in range(self.probe.shape[0]): + # from storing the probe in Fourier space + #pr = tools.propagators.inverse_far_field(self.probe[i]) * self.probe_support + pr = self.probe[i] * self.probe_support + exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr, + self.obj_support * self.obj, + pix_trans, + shift_probe=True) + exit_waves = exit_waves * self.probe_support[...,:,:] + + all_exit_waves.append(exit_waves) + + + return t.stack(all_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, indices): + #return tools.measurements.density_matrix(wavefields,self.rhos[indices], + # detector_slice=self.detector_slice, + # saturation=self.saturation, + # oversampling=self.oversampling) + return tools.measurements.quadratic_background(wavefields, + self.background, self.rhos[indices], + detector_slice=self.detector_slice, + measurement=tools.measurements.density_matrix, + saturation=self.saturation, + oversampling=self.oversampling) + + + def forward(self, *args): + """The complete forward model + + We need to override this to enable the wavefield mixing at the + level of the measurement function + """ + indices = args[0] + return self.measurement(self.forward_propagator(self.interaction(*args)),indices) + + 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): + + rhos_out = np.zeros([self.rhos.shape[0], + self.rhos.shape[1],self.rhos.shape[2]], + dtype=np.complex64) + for (i,j) in ((i,j) for i in range(rhos_out.shape[1]) + for j in range(rhos_out.shape[2])): + if i == j: + rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy() + if i < j: # upper triangle, real part + rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy() + rhos_out[:,j,i] += self.rhos.data[:,i,j].cpu().detach().numpy() + if i > j: # upper triangle, real part + rhos_out[:,j,i] += 1j * self.rhos.data[:,i,j].cpu().detach().numpy() + rhos_out[:,i,j] -= 1j * self.rhos.data[:,i,j].cpu().detach().numpy() + return rhos_out + + 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 + + """ + 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.swapaxes(np.dot(Atrans,np.dot(rhos,Aconj)),0,1) + + new_rhos /= normalization + ortho_probes *= np.sqrt(normalization) + #print(np.dot(Aconjinv,rhos).shape) + #print(np.dot(rhos,Aconj).shape) + new_rhos = cmath.complex_to_torch(new_rhos).to( + dtype=self.rhos.dtype,device=self.rhos.device) + + # This repacks the data into the format used internally + for (i,j) in ((i,j) for i in range(new_rhos.shape[1]) + for j in range(new_rhos.shape[2])): + if i == j: + self.rhos.data[:,i,j] = new_rhos[:,i,j,0] + if i < j: # upper triangle, real part + self.rhos.data[:,i,j] = new_rhos[:,i,j,0] + self.rhos.data[:,j,i] = new_rhos[:,i,j,1] + + + 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 = [ + ('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), + ('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.probe) >=2), + ('Von Neumann Entropy (only accurate after tidy_probes)', + lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_vn_entropy(self.get_rhos()), fig=fig)), + ('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 + 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_ptycho2.py b/CDTools/models/unified_mode_ptycho2.py new file mode 100644 index 0000000..dbab1f5 --- /dev/null +++ b/CDTools/models/unified_mode_ptycho2.py @@ -0,0 +1,490 @@ +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__ = ['UnifiedModePtycho2'] + +class UnifiedModePtycho2(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(UnifiedModePtycho2,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] + + if type(index) == type(0): + index = [index] + Ws = [Ws] + pix_trans = [pix_trans] + single_frame = True + else: + single_frame = False + + probes = [] + all_exit_waves = [] + + # This is the case if a purely stable, incoherent model is defined. + if len(Ws[0].shape) == 0: + # What we do here is generate an identity matrix, and multiply + # that identity matrix by the per-frame weight. + Ws = [W * t.stack([t.eye(self.probe.shape[0]), + t.zeros([self.probe.shape[0]]*2)],dim=-1).to( + dtype=W.dtype, device=W.device) + for W in Ws] + + + # Outer iteration is the mode index iteration + for i in range(Ws[0].shape[0]): + # Now we need to separately treat each mode + exit_waves = [] + for W, pix_tran in zip(Ws, pix_trans): + # from storing the probe in Fourier space + pr = [cmath.cmult(W[i,j,:], self.probe[j] * self.probe_support) + for j in range(self.probe.shape[0])] + + pr = t.sum(t.stack(pr), axis=0) + + exit_waves.append(self.probe_norm * + tools.interactions.ptycho_2D_sinc( + pr, self.obj_support * self.obj, + pix_tran, shift_probe=True)) + + exit_waves = t.stack(exit_waves) + + if single_frame: + exit_waves = exit_waves[0] + # Multiply again by probe support to suppress the fringes from the + # sinc-interpolated shift + exit_waves = exit_waves * self.probe_support[...,:,:] + + all_exit_waves.append(exit_waves) + + + return t.stack(all_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(UnifiedModePtycho2, 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 = [ + ('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), + ('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 1811087..538d6e8 100644 --- a/CDTools/tools/analysis/analysis.py +++ b/CDTools/tools/analysis/analysis.py @@ -12,14 +12,14 @@ import numpy as np from CDTools.tools import cmath from CDTools.tools import image_processing as ip from scipy import fftpack +from scipy import linalg as sla -__all__ = ['orthogonalize_probes','standardize', 'synthesize_reconstructions', +__all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions', 'calc_consistency_prtf', 'calc_deconvolved_cross_correlation', - 'calc_frc'] + 'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction'] -from matplotlib import pyplot as plt -def orthogonalize_probes(probes): +def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False): """Orthogonalizes a set of incoherently mixing probes The strategy is to define a reduced orthogonal basis that spans @@ -27,55 +27,104 @@ def orthogonalize_probes(probes): defined by the probes in that basis. After diagonalization, the eigenvectors can be recast into the original basis and returned + By default, it assumes that the set of probes are defined just as + standard incoherently mixing probe modes, and orthogonalizes them. + However, if a density matrix is explicitly given, it will instead + consider the problem of extracting the eigenbasis of the matrix + probes * denstity_matrix * probes^dagger, where probes is the + column matrix of the given probe functions. This latter problem arises + in the generalization of the probe mixing model, and reduces to the + simpler case when the density matrix is equal to the identity matrix + + If the parameter "keep_transform" is set, the function will additionally + return the matrix A such that A * ortho_probes^dagger = probes^dagger + + If the parameter "normalize" is False (as is the default), the variation + in intensities in the probe modes will be kept in the probe modes, as is + natural for a purely incoherent model. If it is set to "True", the + returned probe modes will all be normalized instead. + Parameters ---------- probes : array An l x n x m complex array representing a stack of probes + density_matrix : np.array + An optional l x l density matrix further elaborating on the state + keep_transform : bool + Default False, whether to return the map from probes to ortho_probes + normalize : bool + Default False, whether to normalize the probe modes Returns ------- ortho_probes: array An l x n x m complex array representing a stack of probes """ - + try: probes = cmath.torch_to_complex(probes.detach().cpu()) send_to_torch = True except: send_to_torch = False - bases = [] - coefficients = np.zeros((probes.shape[0],probes.shape[0]), dtype=np.complex64) - for i, probe in enumerate(probes): - ortho_probe = np.copy(probe) - for j, basis in enumerate(bases): - coefficients[j,i] = np.sum(basis.conj()*ortho_probe) - ortho_probe -= basis * coefficients[j,i] + from matplotlib import pyplot as plt + # We can do the orthogonalization with an SVD, so first we have to + # reshape the final two dimensions (the image shape) into a single + # vectorized dimension. This matrix is probes^dagger, hence the + # conjugation + probes_mat = probes.reshape(probes.shape[0], + probes.shape[1]*probes.shape[2]) + if density_matrix is None: + density_matrix = np.eye(probes.shape[0]) + + + # next we want to extract the eigendecomposition of the density matrix + # itself + w,v = sla.eigh(density_matrix) + w,v = np.linalg.eigh(density_matrix) + w = w[::-1] + v = v[:,::-1] - coefficients[i,i] = np.sqrt(np.sum(np.abs(ortho_probe)**2)) - bases.append(ortho_probe / coefficients[i,i]) + # We do this just to avoid total failure when the density + # matrix is not positive definite. + w = np.abs(w) + # Note: this will fail if any w are less than zero + # Probably should figure out a good way to deal with that + B_dagger = np.dot(np.diag(np.sqrt(w)), v.conj().transpose()) + #u,s,vh = np.linalg.svd(np.dot(B_dagger,probes_mat), full_matrices=False) + u,s,vh = sla.svd(np.dot(B_dagger,probes_mat), full_matrices=False) + - density_mat = coefficients.dot(np.conj(coefficients).transpose()) - eigvals, eigvecs = np.linalg.eigh(density_mat) + if normalize: + ortho_probes = vh.reshape(probes.shape[0], + probes.shape[1], + probes.shape[2]) - ortho_probes = [] - for i in range(len(eigvals)): - coefficients = np.sqrt(eigvals[i]) * eigvecs[:,i] - probe = np.zeros(bases[0].shape, dtype=np.complex64) - for coefficient, basis in zip(coefficients, bases): - probe += basis * coefficient - ortho_probes.append(probe) - - - if send_to_torch: - return cmath.complex_to_torch(np.stack(ortho_probes[::-1])) + B_dagger_inv = np.linalg.pinv(B_dagger) + A = np.dot(B_dagger_inv,np.dot(u,np.diag(s))) + #A_dagger = np.dot(np.linalg.pinv(np.diag(s)), + # np.dot(np.transpose(u).conj(),B_dagger)) else: - return np.stack(ortho_probes[::-1]) - + ortho_probes = np.dot(np.diag(s),vh).reshape(probes.shape[0], + probes.shape[1], + probes.shape[2]) + B_dagger_inv = np.linalg.pinv(B_dagger) + A = np.dot(B_dagger_inv,u) + #A_dagger = np.dot(np.transpose(u).conj(),B_dagger) + + if send_to_torch: + ortho_probes = cmath.complex_to_torch(np.stack(ortho_probes)) + A = cmath.complex_to_torch(A) + #A_dagger = cmath.complex_to_torch(A_dagger) + if keep_transform: + return ortho_probes, A#_dagger + else: + return ortho_probes + def standardize(probe, obj, obj_slice=None, correct_ramp=False): """Standardizes a probe and object to prepare them for comparison @@ -512,3 +561,66 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.): threshold = t.tensor(threshold) return bins[:-1], frc, threshold + + +def calc_vn_entropy(matrix): + """Calculates the Von Neumann entropy of a density matrix + + Will either accept a single matrix, or a stack of matrices. Matrices + are assumed to be Hermetian and positive definite, to be well-formed + density matrices + + Parameters + ---------- + matrix : np.array + The nxn matrix or lxnxn stack of matrices to calculate the entropy of + + + Returns + ------- + entropy: float or np.array + The entropy or entropies of the arrays + """ + + if len(matrix.shape) == 3: + # Get the eigenvalues + eigs = [np.linalg.eigh(mat)[0] for mat in matrix] + # Normalize them to match standard density matrix form + eigs = [eig / np.sum(eig) for eig in eigs] + # And calculate the VN entropy! + entropies = [-np.sum(eig*np.log(eig)) for eig in eigs] + return np.array(entropies) + else: + eig = np.linalg.eigh(matrix)[0] + entropy = -np.sum(eig*np.log(eig))/np.sum(eig) + return -np.trace(np.dot(matrix,sla.logm(matrix))) + +def calc_top_mode_fraction(matrix): + """Calculates the fraction of total power in the top mode of a density matrix + + Will either accept a single matrix, or a stack of matrices. Matrices + are assumed to be Hermetian and positive definite, to be well-formed + density matrices + + Parameters + ---------- + matrix : np.array + The nxn matrix or lxnxn stack of matrices to work from + + + Returns + ------- + entropy: float or np.array + The fraction of power in the top mode of each matrix + """ + + if len(matrix.shape) == 3: + # Get the eigenvalues + eigs = [np.linalg.eigh(mat)[0] for mat in matrix] + # Normalize them to match standard density matrix form + fractions = [np.max(eig) / np.sum(eig) for eig in eigs] + return np.array(fractions) + else: + eig = np.linalg.eigh(matrix)[0] + fraction = np.max(eig) / np.sum(eig) + return fraction diff --git a/CDTools/tools/cmath/cmath.py b/CDTools/tools/cmath/cmath.py index d969f53..282b840 100644 --- a/CDTools/tools/cmath/cmath.py +++ b/CDTools/tools/cmath/cmath.py @@ -13,7 +13,7 @@ import torch as t __all__ = ['complex_to_torch', 'torch_to_complex', 'cabssq', 'cabs', 'cconj', - 'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi'] + 'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi', 'cexpi'] # diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py index 7b7e5ad..42f272d 100644 --- a/CDTools/tools/interactions/interactions.py +++ b/CDTools/tools/interactions/interactions.py @@ -140,7 +140,15 @@ def project_translations_to_sample(sample_basis, translations): describing the location of the probe's intersection with the sample plane, assuming the basis found in sample_basis is used. Second, an amount that the probe must be propagated along the z-axis to reach - that location on the sample plane. + the sample plane at the given location. This includes both the effect of + the tilted sample plane and any explicitly defined motion along the z axis + of the probe-forming optic as included in the input translations. + + + Note that because of the sign convention (that this function returns th + relative amount the probe needs to be propagated to reach any given + location), a positive motion along the z-axis of the probe forming optics + will lead to a negative propagation distance. The assumed geometry is incoming radiation with a wavevector parallel to the +z axis, [0,0,1]. @@ -177,6 +185,8 @@ def project_translations_to_sample(sample_basis, translations): dtype=surface_normal.dtype) # Here we're setting up a matrix-vector equation mat*answer=input + # At some point ger will need to be replaced by outer, but for now + # outer many places still don't have new enough versions of torch. mat = t.cat((I - t.ger(propagation_dir,propagation_dir), surface_normal.unsqueeze(0))) @@ -205,8 +215,9 @@ def project_translations_to_sample(sample_basis, translations): single_translation = True pixel_translations = t.mm(translations, sample_projection) - propagations = t.mm(translations, prop_projection) - + propagations = t.mm(translations, prop_projection) \ + - t.mm(translations,propagation_dir[:,None]) + if single_translation: return pixel_translations[0], propagations[0] else: diff --git a/CDTools/tools/measurements/measurements.py b/CDTools/tools/measurements/measurements.py index 31f7c2f..9d3fba4 100644 --- a/CDTools/tools/measurements/measurements.py +++ b/CDTools/tools/measurements/measurements.py @@ -17,7 +17,8 @@ from torch.nn.functional import avg_pool2d # intensity pattern on a detector # -__all__ = ['intensity', 'incoherent_sum', 'quadratic_background'] +__all__ = ['intensity', 'incoherent_sum', 'density_matrix', + 'quadratic_background'] def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1): @@ -66,6 +67,101 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove return t.clamp(output + epsilon,0,saturation) +def density_matrix(wavefields, density_matrix, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1): + """Returns the intensities associated with a given density matrix state + + The essential idea is that the most general description of a light field + at the detector plane will consist of a density matrix state. Here, that + low rank density matrix state is encoded as a set of basis wavefields + and a density matrix in that basis. + + For computational efficiency, the density matrix is coded in an unusual + format. The density matrix formally is a complex Hermetian matrix, + which also happens to be positive definite. Here, we store it as a + real-valued matrix, where the upper triangle corresponds to the real + part of the elements in the upper triangle, and the lower triangle + corresponds to the imaginary parts. The elements on the diagonal are + purely real, and are stored as they are. + + As with other multi-mode measurement functions, the modes are stored in + the first index, and the index of the diffraction pattern in the stack + of diffraction patterns is the second index. The stack-direction index + can be omitted if only a single pattern needs to be simulated + + It is important to note that this method does not inforce the positive + definiteness of the density matrix, this it is possible for negative + values of intensity to appear if the underlying density matrix passed + to this method is not positive definite + + Parameters + ---------- + wavefields : torch.Tensor + An Lx(Jx)MxNx2 stack of complex wavefields + density_matrix : torch.Tensor + A (Jx)LxL stack of real-valued representations of density matrices, as per above + saturation : float + Optional, a maximum saturation value to clamp the resulting intensities to + oversampling : int + Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel + + Returns + ------- + sim_patterns : torch.Tensor + A real Lx(Jx)MxN array storing the incoherently summed intensities + + """ + + #if wavefields.dim() == 4: + # wavefields.unsqueeze(1) + # single_frame = True + #elif wavefields.dim() == 5: + # single_frame=False + + output = t.zeros(wavefields.shape[1:-1], + dtype=wavefields.dtype, + device=wavefields.device) + + # flat is better than nested, but simple is better than complex... + for (i,j) in ((i,j) for i in range(density_matrix.shape[-2]) + for j in range(density_matrix.shape[-1])): + if i == j: # diagonal + output += density_matrix[...,i,j,None,None] \ + * cmath.cabssq(wavefields[i]) + if i < j: # upper triangle, real part + output += 2 * density_matrix[...,i,j,None,None] \ + * (wavefields[i,...,0] * wavefields[j,...,0] + + wavefields[i,...,1] * wavefields[j,...,1]) + if i > j: # lower triangle, imaginary part + # We pull the i,jth element from the density matrix, + # but this correponds to wavefield j and wavefield i, + # unlike above where it was wavefield i and j (swapped order). + # We also get the one negative sign because this is the imaginary + # part + output += 2 * density_matrix[...,i,j,None,None] \ + * (wavefields[j,...,0] * wavefields[i,...,1] + - wavefields[j,...,1] * wavefields[i,...,0]) + + # Now we apply oversampling + if oversampling != 1: + if wavefields.dim() == 4: + output = avg_pool2d(output.unsqueeze(0), oversampling)[0] + else: + output = avg_pool2d(output, oversampling) + + # Then we grab the detector slice + if detector_slice is not None: + if wavefields.dim() == 4: + output = output[detector_slice] + else: + output = output[(np.s_[:],) + detector_slice] + + if saturation is None: + return t.clamp(output,min=0) + epsilon + else: + return t.clamp(output + epsilon,0,saturation) + + + def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1): """Returns the incoherent sum of the intensities of the wavefields @@ -94,7 +190,6 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non sim_patterns : torch.Tensor A real JXMxN array storing the incoherently summed intensities """ - # This syntax just adds an axis to the slice to preserve the J direction output = t.sum(cmath.cabssq(wavefields),dim=0) @@ -118,7 +213,7 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non return t.clamp(output + epsilon,0,saturation) -def quadratic_background(wavefield, background, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1): +def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1): """Returns the intensity of a wavefield plus a background The intensity is calculated via the given measurment function @@ -148,10 +243,10 @@ def quadratic_background(wavefield, background, detector_slice=None, measurement """ if detector_slice is None: - output = measurement(wavefield, epsilon=epsilon, + output = measurement(wavefield, *args, epsilon=epsilon, oversampling=oversampling) + background**2 else: - output = measurement(wavefield, detector_slice, + output = measurement(wavefield, *args, detector_slice=detector_slice, epsilon=epsilon, oversampling=oversampling) \ + background**2 diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py index 619fe2c..14c9a4d 100644 --- a/CDTools/tools/plotting/plotting.py +++ b/CDTools/tools/plotting/plotting.py @@ -17,7 +17,7 @@ from matplotlib.colors import hsv_to_rgb __all__ = ['colorize', 'plot_amplitude', 'plot_phase', 'plot_colorized', 'plot_translations', 'get_units_factor', - 'plot_nanomap'] + 'plot_nanomap', 'plot_real', 'plot_imag'] def colorize(z): @@ -82,6 +82,140 @@ def get_units_factor(units): return factor +def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs): + """Plots the real part of a complex array with dimensions NxM + + If a figure is given explicitly, it will clear that existing figure and + plot over it. Otherwise, it will generate a new figure. + + If a basis is explicitly passed, the image will be plotted in real-space + coordinates + + Parameters + ---------- + im : array + An complex array with dimensions NxM + fig : matplotlib.figure.Figure + Default is a new figure, a matplotlib figure to use to plot + basis : np.array + Optional, the 3x2 probe basis + units : str + The length units to mark on the plot, default is um + cmap : str + Default is 'viridis', the colormap to plot with + \\**kwargs + All other args are passed to fig.add_subplot(111, \\**kwargs) + + Returns + ------- + used_fig : matplotlib.figure.Figure + The figure object that was actually plotted to. + """ + if fig is None: + fig = plt.figure() + ax = fig.add_subplot(111, **kwargs) + else: + plt.figure(fig.number) + plt.gcf().clear() + + if isinstance(im, t.Tensor): + real = im[...,0].detach().cpu().numpy() + else: + real = np.real(im) + + #Plot in a basis if it exists, otherwise dont + if basis is not None: + if isinstance(basis,t.Tensor): + basis = basis.detach().cpu().numpy() + # This fails if the + basis_norm = np.linalg.norm(basis, axis = 0) + basis_norm = basis_norm * get_units_factor(units) + + extent = [0, real.shape[-1]*basis_norm[1], 0, real.shape[-2]*basis_norm[0]] + else: + extent=None + + plt.imshow(real, cmap = cmap, extent = extent) + cbar = plt.colorbar() + cbar.set_label('Real Part (a.u.)') + + if basis is not None: + plt.xlabel('X (' + units + ')') + plt.ylabel('Y (' + units + ')') + else: + plt.xlabel('j (pixels)') + plt.ylabel('i (pixels)') + + return fig + + +def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs): + """Plots the imaginary part of a complex array with dimensions NxM + + If a figure is given explicitly, it will clear that existing figure and + plot over it. Otherwise, it will generate a new figure. + + If a basis is explicitly passed, the image will be plotted in real-space + coordinates + + Parameters + ---------- + im : array + An complex array with dimensions NxM + fig : matplotlib.figure.Figure + Default is a new figure, a matplotlib figure to use to plot + basis : np.array + Optional, the 3x2 probe basis + units : str + The length units to mark on the plot, default is um + cmap : str + Default is 'viridis', the colormap to plot with + \\**kwargs + All other args are passed to fig.add_subplot(111, \\**kwargs) + + Returns + ------- + used_fig : matplotlib.figure.Figure + The figure object that was actually plotted to. + """ + if fig is None: + fig = plt.figure() + ax = fig.add_subplot(111, **kwargs) + else: + plt.figure(fig.number) + plt.gcf().clear() + + if isinstance(im, t.Tensor): + imag = im[...,1].detach().cpu().numpy() + else: + imag = np.imag(im) + + #Plot in a basis if it exists, otherwise dont + if basis is not None: + if isinstance(basis,t.Tensor): + basis = basis.detach().cpu().numpy() + # This fails if the + basis_norm = np.linalg.norm(basis, axis = 0) + basis_norm = basis_norm * get_units_factor(units) + + extent = [0, imag.shape[-1]*basis_norm[1], 0, imag.shape[-2]*basis_norm[0]] + else: + extent=None + + plt.imshow(imag, cmap = cmap, extent = extent) + cbar = plt.colorbar() + cbar.set_label('Imaginary Part (a.u.)') + + if basis is not None: + plt.xlabel('X (' + units + ')') + plt.ylabel('Y (' + units + ')') + else: + plt.xlabel('j (pixels)') + plt.ylabel('i (pixels)') + + return fig + + def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs): """Plots the amplitude of a complex array with dimensions NxM diff --git a/CDTools/tools/propagators/propagators.py b/CDTools/tools/propagators/propagators.py index 3e8ab1d..2eca125 100644 --- a/CDTools/tools/propagators/propagators.py +++ b/CDTools/tools/propagators/propagators.py @@ -205,7 +205,7 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance, intensity_map = t.Tensor(intensity_map).to(*args, **kwargs) - + # This accounts for the implied phase ramp along the exit wave direction # In other words, it prevents the diffraction pattern from sliding off the # detector when the sample is tilted but represented by an object with diff --git a/examples/example_reconstructions/gold_balls.pickle b/examples/example_reconstructions/gold_balls.pickle index ff2f4be..9115f28 100644 Binary files a/examples/example_reconstructions/gold_balls.pickle and b/examples/example_reconstructions/gold_balls.pickle differ diff --git a/examples/example_reconstructions/lab_ptycho.pickle b/examples/example_reconstructions/lab_ptycho.pickle index 25e6ab5..745ca6d 100644 Binary files a/examples/example_reconstructions/lab_ptycho.pickle and b/examples/example_reconstructions/lab_ptycho.pickle differ diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 7909157..6917c0e 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -3,6 +3,7 @@ from __future__ import division, print_function, absolute_import import CDTools from matplotlib import pyplot as plt import pickle +import time # First, we load an example dataset from a .cxi file filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi' @@ -18,8 +19,8 @@ dataset.get_as(device='cuda') for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)): # And we liveplot the updates to the model as they happen - model.inspect(dataset) print(i,loss) + model.inspect(dataset) # And we save the reconstruction out to a file with open('example_reconstructions/gold_balls.pickle', 'wb') as f: diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index d9ff18e..866a7a0 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -6,7 +6,7 @@ from matplotlib import pyplot as plt # First, we load an example dataset from a .cxi file filename = 'example_data/lab_ptycho_data.cxi' dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) - +plt.ion() # Next, we create a ptychography model from the dataset model = CDTools.models.SimplePtycho.from_dataset(dataset) diff --git a/examples/unified_modes.py b/examples/unified_modes.py new file mode 100644 index 0000000..05dbc27 --- /dev/null +++ b/examples/unified_modes.py @@ -0,0 +1,39 @@ +from __future__ import division, print_function, absolute_import + +import CDTools +from matplotlib import pyplot as plt +import pickle + +filename = 'example_data/lab_ptycho_data.cxi' +dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) + +# dataset.inspect() +# plt.show() + +#model = CDTools.models.UnifiedModePtycho.from_dataset(dataset, oversampling=2,n_modes=3)#, probe_support_radius=90) +model = CDTools.models.UnifiedModePtycho2.from_dataset(dataset, oversampling=1,n_modes=3)#, probe_support_radius=90) +#model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=1) + +model.to(device='cuda') +dataset.get_as(device='cuda') + + +model.translation_offsets.requires_grad = False +for i, loss in enumerate(model.Adam_optimize(100, dataset)): + model.inspect(dataset) + print(i,loss) + +model.tidy_probes() + +for i, loss in enumerate(model.Adam_optimize(20, dataset, lr=0.0001)): + model.inspect(dataset) + print(i,loss) + +model.tidy_probes() +model.inspect(dataset) + +with open('example_reconstructions/unified_modes.pickle', 'wb') as f: + pickle.dump(model.save_results(dataset),f) + +model.compare(dataset) +plt.show() diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index c0feade..12e897f 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -22,8 +22,8 @@ def test_orthogonalize_probes(): probe_Ys, probe_Xs = np.meshgrid(probe_ys, probe_xs) probe_Rs = np.sqrt(probe_Xs**2 + probe_Ys**2) - probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2)), - 3*np.exp(-probe_Rs**2 / (2 * 12**2)), + probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2 + 1j)), + 3*np.exp(-probe_Rs**2 / (2 * 12**2 - 3j)), 1*np.exp(-probe_Rs**2 / (2 * 15**2))]).astype(np.complex64) # test that it works on numpy arrays @@ -31,13 +31,15 @@ def test_orthogonalize_probes(): # test that it also works on torch tensors ortho_probes_t = cmath.torch_to_complex(analysis.orthogonalize_probes(cmath.complex_to_torch(probes))) - + + # This tests for orthogonality for p1,p2 in combinations(ortho_probes,2): assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6 for p1,p2 in combinations(ortho_probes_t,2): assert np.sum(np.conj(p1)*p2) / np.sum(np.abs(p1)**2) < 1e-6 + probe_intensity = np.sum(np.abs(probes)**2,axis=0) ortho_probe_intensity = np.sum(np.abs(ortho_probes)**2,axis=0) ortho_probe_t_intensity = np.sum(np.abs(ortho_probes_t)**2,axis=0) @@ -45,8 +47,35 @@ def test_orthogonalize_probes(): assert np.allclose(probe_intensity,ortho_probe_intensity) assert np.allclose(probe_intensity,ortho_probe_t_intensity) + # Check that it returns normalized probes if we ask + ortho_probes = analysis.orthogonalize_probes(probes, normalize=True) + assert np.allclose([1,1,1],np.sum(np.abs(ortho_probes)**2,axis=(1,2))) + # And now we check that the A matrices actually work + ortho_probes, A = analysis.orthogonalize_probes(probes, keep_transform=True, + normalize=False) + assert np.allclose(np.tensordot(A, ortho_probes,axes=1),probes) + # And now we check that the A matrices actually work + ortho_probes, A = analysis.orthogonalize_probes(probes, keep_transform=True, + normalize=False) + assert np.allclose(np.tensordot(A, ortho_probes,axes=1),probes) + + # The big problem here is that we haven't tested if it returns the + # probes or their complex conjugate.. we can test that by sending in + # one probe and checking that we get the right thing out + + probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 20**2 + 100j))]) + ortho_probes = analysis.orthogonalize_probes(probes) + # We need to correct for the arbitrary phase offset to be able to + # compare them + correction = ortho_probes[0]/probes[0] + correction = correction / np.abs(correction) + ortho_probes *= correction + + assert np.allclose(ortho_probes,probes) + + def test_standardize(): # Start by making a probe and object that should meet the standardization diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index bb3a85e..4f6871c 100644 --- a/tests/tools/test_initializers.py +++ b/tests/tools/test_initializers.py @@ -192,6 +192,32 @@ def test_SHARP_style_probe(ptycho_cxi_1): def test_RPI_spectral_init(): - # Figure out a good way to test this - assert 0 - RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, background=None) + # I think we can only really meaningfully test that it doesn't throw errors, + # since the original implementation is in numpy and there aren't any clear + # cases that can be calculated analytically. + + pattern = np.random.rand(230,253).astype(np.float32) + probe = np.random.rand(230,253).astype(np.complex64) + obj_shape = [37,53] + mask = t.Tensor(np.random.rand(*pattern.shape) > 0.04) + background = t.Tensor(np.random.rand(*pattern.shape) .astype(np.float32)* 0.05) + + probe = cmath.complex_to_torch(probe) + pattern = t.Tensor(pattern) + + obj = initializers.RPI_spectral_init(pattern, probe, obj_shape) + assert list(obj.shape) == [1]+obj_shape+[2] + + obj = initializers.RPI_spectral_init(pattern, probe, obj_shape, + n_modes=2, mask=mask) + assert list(obj.shape) == [2]+obj_shape+[2] + + obj = initializers.RPI_spectral_init(pattern, probe, obj_shape, + n_modes=2, background=background) + assert list(obj.shape) == [2]+obj_shape+[2] + + obj = initializers.RPI_spectral_init(pattern, probe, obj_shape, + n_modes=2, mask=mask, + background=background) + assert list(obj.shape) == [2]+obj_shape+[2] + diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py index e4f1177..79e2bfe 100644 --- a/tests/tools/test_interactions.py +++ b/tests/tools/test_interactions.py @@ -4,6 +4,8 @@ from CDTools.tools import cmath from CDTools.tools import interactions import numpy as np import torch as t +from numpy import fft +from scipy.fftpack import fftshift, ifftshift import pytest @@ -98,8 +100,35 @@ def test_pixel_to_translations(): def test_project_translations_to_sample(): - # Needs to be tested - assert 0 + # First, try the case where everything is ones and simple + basis = t.Tensor([[0,-1,0],[-1,0,0]]).t() + translations = t.rand((10,3)) + pixels, props = interactions.project_translations_to_sample(basis, translations) + + assert np.allclose(pixels[:,0].numpy(),-translations[:,1]) + assert np.allclose(pixels[:,1].numpy(),-translations[:,0]) + assert np.allclose(props.numpy(),-translations[:,2:].numpy()) + + # Next, a simple tilt along one axis. This is a 45 degree rotation + # around the positive y-axis + # Thus, y-axis translations are unaffected, but x-axis translations + # induce a motion of 1/sqrt(2) in the j- pixel space, as well as + # creating a propagation (negative propagation for positive x) + basis = t.Tensor([[0,-1e-3,0],[-np.sqrt(2)*1e-3,0,np.sqrt(2)*1e-3]]).t() + translations = t.rand((10,3)) + pixels, props = interactions.project_translations_to_sample(basis, translations) + + print(props.numpy()) + print(-translations[:,2:].numpy() - translations[:,:1].numpy()) + assert np.allclose(pixels[:,0].numpy(),-translations[:,1]*1e3) + assert np.allclose(pixels[:,1].numpy(),-translations[:,0]*1e3/np.sqrt(2)) + assert np.allclose(props.numpy(),-translations[:,2:].numpy() - translations[:,:1].numpy()) + + # Finally, we check a non-orthogonal case + + + + def test_ptycho_2D_round(random_probe, random_obj): # Test a stack of images @@ -227,22 +256,39 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj): def test_RPI_interaction(random_probe, random_obj): - random_obj1 = cmath.complex_to_torch(random_obj[:79,:68]) - random_probe1 = cmath.complex_to_torch(random_probe) - output = interactions.RPI_interaction(random_probe1, random_obj1) + random_obj1 = random_obj[:79,:68] + random_probe1 = random_probe + t_random_obj1 = cmath.complex_to_torch(random_obj1) + t_random_probe1 = cmath.complex_to_torch(random_probe1) + t_output1 = interactions.RPI_interaction(t_random_probe1, t_random_obj1) - random_obj1 = cmath.complex_to_torch(random_obj[:256,:256]) - random_probe1 = cmath.complex_to_torch(random_probe) - output = interactions.RPI_interaction(random_probe1, random_obj1) + obj1_fourier = fftshift(fft.fft2(ifftshift(random_obj1), norm='ortho')) + obj1_ups = np.zeros(random_probe1.shape[:2]).astype(np.complex128) + obj1_ups[(random_probe1.shape[0]-79)//2: + (random_probe1.shape[0]-79)//2 + 79, + (random_probe1.shape[1]-68)//2: + (random_probe1.shape[1]-68)//2 + 68] = obj1_fourier + output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups), + norm='ortho')) + + assert np.allclose(cmath.torch_to_complex(t_output1), output1) - random_obj1 = cmath.complex_to_torch(random_obj[:42,:103]) - random_probe1 = cmath.complex_to_torch(random_probe) - output = interactions.RPI_interaction(random_probe1, random_obj1) + random_obj2 = np.stack([random_obj[:64,:89]]*3) + random_probe2 = random_probe[3:,5:] + t_random_obj2 = cmath.complex_to_torch(random_obj2) + t_random_probe2 = cmath.complex_to_torch(random_probe2) + t_output2 = interactions.RPI_interaction(t_random_probe2, t_random_obj2) + + obj2_fourier = fftshift(fft.fft2(ifftshift(random_obj2), norm='ortho')) + obj2_ups = np.zeros((3,)+random_probe2.shape[:2]).astype(np.complex128) + obj2_ups[:,(random_probe2.shape[0]-64)//2: + (random_probe2.shape[0]-64)//2 + 64, + (random_probe2.shape[1]-89)//2: + (random_probe2.shape[1]-89)//2 + 89] = obj2_fourier + output2 = random_probe2 * fftshift(fft.ifft2(ifftshift(obj2_ups), + norm='ortho')) + - # Need to actually test against a numpy implementation - print(random_probe.shape) - print(random_obj.shape) - assert 0