diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 8604d7b..f3b4823 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -35,4 +35,3 @@ 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 306abd9..138485b 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -311,7 +311,7 @@ class CDIModel(t.nn.Module): # Define the optimizer optimizer = t.optim.LBFGS(self.parameters(), lr = lr, history_size=history_size) - + return self.AD_optimize(iterations, data_loader, optimizer, regularization_factor=regularization_factor, thread=thread, diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index d68d7bf..c89e50c 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -153,6 +153,7 @@ class FancyPtycho(CDIModel): obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200) + if hasattr(dataset, 'background') and dataset.background is not None: background = t.sqrt(dataset.background) else: @@ -336,16 +337,10 @@ class FancyPtycho(CDIModel): # Needs to be updated to allow for plotting to an existing figure plot_list = [ - ('Dominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Dominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Subdominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), - ('Subdominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), + ('Probe Amplitude (scroll to view modes)', + lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)), + ('Probe Phase (scroll to view modes)', + lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)), ('Object Amplitude', lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), ('Object Phase', diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index 4ae8ed0..c875939 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -387,26 +387,32 @@ class Multislice2DPtycho(CDIModel): # Needs to be updated to allow for plotting to an existing figure plot_list = [ - ('Dominant Probe Fourier Space Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[0] if self.fourier_probe else tools.propagators.inverse_far_field(self.probe[0]), fig=fig)), - ('Dominant Probe Fourier Space Phase', - lambda self, fig: p.plot_phase(self.probe[0] if self.fourier_probe else tools.propagators.inverse_far_field(self.probe[0]), fig=fig)), - ('Dominant Probe Real Space Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[0] if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe[0]), fig=fig, basis=self.probe_basis, units=self.units)), - ('Dominant Probe Real Space Phase', - lambda self, fig: p.plot_phase(self.probe[0] if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe[0]), fig=fig, basis=self.probe_basis, units=self.units)), - ('Subdominant Probe Real Space Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[1] if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe[1]), fig=fig, basis=self.probe_basis, units=self.units), - lambda self: len(self.probe) >=2), - ('Subdominant Probe Real Space Phase', - lambda self, fig: p.plot_phase(self.probe[1] if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe[1]), fig=fig, basis=self.probe_basis, units=self.units), - lambda self: len(self.probe) >=2), + ('Probe Fourier Space Amplitude', + lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), + ('Probe Fourier Space Phase', + lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), + ('Probe Real Space Amplitude', + lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)), + ('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: 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: 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), lambda self: self.exponentiate_obj), + ('Slice by Slice Amplitude of Object Function', + lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), + lambda self: not self.exponentiate_obj), + ('Slice by Slice Phase of Object Function', + lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units), + 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), diff --git a/CDTools/models/unified_mode_ptycho.py b/CDTools/models/unified_mode_ptycho.py index afaa19d..60921e9 100644 --- a/CDTools/models/unified_mode_ptycho.py +++ b/CDTools/models/unified_mode_ptycho.py @@ -18,7 +18,7 @@ class UnifiedModePtycho(CDIModel): def __init__(self, wavelength, detector_geometry, probe_basis, - probe_guess, obj_guess, rhos_guess, + probe_guess, obj_guess, Ws_guess, detector_slice=None, surface_normal=np.array([0.,0.,1.]), min_translation = t.Tensor([0,0]), @@ -71,7 +71,10 @@ class UnifiedModePtycho(CDIModel): self.background = t.nn.Parameter(t.Tensor(background).to(t.float32)) - self.rhos = t.nn.Parameter(t.Tensor(rhos_guess).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 @@ -95,7 +98,7 @@ class UnifiedModePtycho(CDIModel): @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'): + 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'] @@ -175,13 +178,22 @@ class UnifiedModePtycho(CDIModel): 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 - + # 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: @@ -207,7 +219,7 @@ class UnifiedModePtycho(CDIModel): else: obj_support = None - return cls(wavelength, det_geo, probe_basis, probe, obj, rhos, + return cls(wavelength, det_geo, probe_basis, probe, obj, Ws, detector_slice=det_slice, surface_normal=surface_normal, min_translation=min_translation, @@ -229,22 +241,61 @@ class UnifiedModePtycho(CDIModel): if self.translation_offsets is not None: pix_trans += self.translation_scale * self.translation_offsets[index] + Ws = self.Ws[index] + + # This probably needs to be fixed, I doubt it will really still work + # to deal with single-pattern sims. + #if type(index) == type(0): + # print('hi') + # index = [index] + # Ws = [Ws] + # pix_trans = [pix_trans] + # single_frame = True + #else: + # single_frame = False + 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) + + # 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] + #print(Ws) + + # This restricts the basis probes with the probe support + basis_prs = self.probe * self.probe_support[...,:,:] - return t.stack(all_exit_waves) + if len(Ws[0].shape) == 0: + # If a purely stable coherent illumination is defined + prs = cmath.cmult(Ws[...,None,None,None,:],basis_prs) + else: + # If a frame-by-frame weight matrix is defined + # This takes the dot product of all the weight matrices with + # the probes. The output has dimensions of translation, then + # coherent mode index, then x,y, and then complex index + prs = t.sum(cmath.cmult(Ws[...,None,None,:], basis_prs), + axis=-4) + + exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs, self.obj_support * self.obj,pix_trans, + shift_probe=True, multiple_modes=True) + + exit_waves = exit_waves * self.probe_support[...,:,:] + + if hasattr(self,'weights') and self.weights is not None: + if exit_waves.dim() == 5: + exit_waves = self.weights[index][:,None,None,None,None] \ + * exit_waves + else: + exit_waves = self.weights[index] * exit_waves + + return exit_waves def forward_propagator(self, wavefields): @@ -255,28 +306,14 @@ class UnifiedModePtycho(CDIModel): 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) + def measurement(self, wavefields): return tools.measurements.quadratic_background(wavefields, - self.background, self.rhos[indices], + self.background, detector_slice=self.detector_slice, - measurement=tools.measurements.density_matrix, + measurement=tools.measurements.incoherent_sum, 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) @@ -344,22 +381,15 @@ class UnifiedModePtycho(CDIModel): mask=mask) def get_rhos(self): - - rhos_out = np.zeros([self.rhos.shape[0], - self.rhos.shape[1],self.rhos.shape[2]], + # 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) - 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 @@ -368,6 +398,18 @@ class UnifiedModePtycho(CDIModel): 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()) @@ -377,24 +419,26 @@ class UnifiedModePtycho(CDIModel): normalize=True) Aconj = A.conj() Atrans = np.transpose(A) - new_rhos = np.swapaxes(np.dot(Atrans,np.dot(rhos,Aconj)),0,1) + new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj)) 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] + 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) @@ -408,21 +452,16 @@ class UnifiedModePtycho(CDIModel): # Needs to be updated to allow for plotting to an existing figure plot_list = [ - ('Dominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Dominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)), - ('Subdominant Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), - ('Subdominant Probe Phase', - lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), - lambda self: len(self.probe) >=2), + ('Basis Probe Amplitudes', + lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)), + ('Basis Probe Phases', + lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)), ('Average Density Matrix Amplitudes', lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig), - lambda self: len(self.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)), + 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', @@ -441,9 +480,9 @@ class UnifiedModePtycho(CDIModel): 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() + Ws = cmath.torch_to_complex(self.Ws.detach().cpu()) return {'basis':basis, 'translation':translations, 'probe':probe,'obj':obj, 'background':background, - 'weights':weights} + 'Ws':Ws} diff --git a/CDTools/models/unified_mode_ptycho2.py b/CDTools/models/unified_mode_ptycho2.py deleted file mode 100644 index dbab1f5..0000000 --- a/CDTools/models/unified_mode_ptycho2.py +++ /dev/null @@ -1,490 +0,0 @@ -from __future__ import division, print_function, absolute_import - -import torch as t -from CDTools.models import CDIModel -from CDTools.datasets import Ptycho2DDataset -from CDTools import tools -from CDTools.tools import cmath -from CDTools.tools import analysis -from CDTools.tools import plotting as p -from matplotlib import pyplot as plt -from datetime import datetime -import numpy as np -from copy import copy - -__all__ = ['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/plotting/plotting.py b/CDTools/tools/plotting/plotting.py index 14c9a4d..5d66f72 100644 --- a/CDTools/tools/plotting/plotting.py +++ b/CDTools/tools/plotting/plotting.py @@ -82,7 +82,146 @@ def get_units_factor(units): return factor -def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs): +def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, **kwargs): + """Plots an image with a colorbar and on an appropriate spatial grid + + 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 + + Finally, if a function is passed to the plot_func argument, this function + will be called on each slice of data before it is plotted. This is used + internally to enable the plot_real, plot_image, plot_phase, etc. functions. + + + Parameters + ---------- + im : array + An complex array with dimensions NxM + plot_func : callable + A function which maps numpy arrays to the image to be plotted + 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 + cmap_label : str + What to label the colorbar when plotting + \\**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. + """ + + # convert to numpy + if isinstance(im, t.Tensor): + # If final dimension is 2, assume it is a complex array. If not, + # assume it represents a real array + if im.shape[-1] == 2: + im = cmath.torch_to_complex(im.detach().cpu()) + else: + im = im.detach().cpu().numpy() + + if fig is None: + fig = plt.figure() + ax = fig.add_subplot(111, **kwargs) + + # This nukes everything and updates either the appropriate image from the + # stack of images, or the only image if only a single image has been + # given + def make_plot(idx): + plt.figure(fig.number) + title = plt.gca().get_title() + fig.clear() + + + # If im only has two dimensions, this reshape will add a leading + # dimension, and update will be called on index 0. If it has 3 or more + # dimensions, then all the leading dimensions will be compressed into + # one long dimension which can be scrolled through. + s = im.shape + reshaped_im = im.reshape(-1,s[-2],s[-1]) + num_images = reshaped_im.shape[0] + fig.plot_idx = idx % num_images + + to_plot = plot_func(reshaped_im[fig.plot_idx]) + + #Plot in a basis if it exists, otherwise dont + if basis is not None: + if isinstance(basis,t.Tensor): + np_basis = basis.detach().cpu().numpy() + else: + np_basis = basis + # This fails if the basis is not rectangular + basis_norm = np.linalg.norm(np_basis, axis = 0) + basis_norm = basis_norm * get_units_factor(units) + + extent = [0, to_plot.shape[-1]*basis_norm[1], 0, + to_plot.shape[-2]*basis_norm[0]] + else: + extent=None + + plt.imshow(to_plot, cmap = cmap, extent = extent) + cbar = plt.colorbar() + if cmap_label is not None: + cbar.set_label(cmap_label) + + if basis is not None: + plt.xlabel('X (' + units + ')') + plt.ylabel('Y (' + units + ')') + else: + plt.xlabel('j (pixels)') + plt.ylabel('i (pixels)') + + + plt.title(title) + + if len(im.shape) >= 3: + plt.text(0.03, 0.03, str(fig.plot_idx), fontsize=14, transform=plt.gcf().transFigure) + return fig + + if hasattr(fig, 'plot_idx'): + result = make_plot(fig.plot_idx) + else: + result = make_plot(0) + + update = make_plot + + + def on_action(event): + if not hasattr(event, 'button'): + event.button = None + if not hasattr(event, 'key'): + event.key = None + + if event.key == 'up' or event.button == 'up': + update(fig.plot_idx - 1) + elif event.key == 'down' or event.button == 'down': + update(fig.plot_idx + 1) + plt.draw() + + if len(im.shape) >=3: + if not hasattr(fig,'my_callbacks'): + fig.my_callbacks = [] + + for cid in fig.my_callbacks: + fig.canvas.mpl_disconnect(cid) + fig.my_callbacks = [] + fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action)) + fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action)) + + return result + + +def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **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 @@ -103,6 +242,8 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa The length units to mark on the plot, default is um cmap : str Default is 'viridis', the colormap to plot with + cmap_label : str + What to label the colorbar when plotting \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -111,45 +252,14 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa 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 + plot_func = lambda x: np.real(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, cmap=cmap, cmap_label=cmap_label, + **kwargs) + -def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs): +def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **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 @@ -170,6 +280,8 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa The length units to mark on the plot, default is um cmap : str Default is 'viridis', the colormap to plot with + cmap_label : str + What to label the colorbar when plotting \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -178,53 +290,21 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa 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 + plot_func = lambda x: np.imag(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, cmap=cmap, cmap_label=cmap_label, + **kwargs) -def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs): +def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Amplitude (a.u.)', **kwargs): """Plots the amplitude 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 - + coordinates. + Parameters ---------- im : array @@ -237,6 +317,8 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', The length units to mark on the plot, default is um cmap : str Default is 'viridis', the colormap to plot with + cmap_label : str + What to label the colorbar when plotting \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -245,45 +327,13 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', 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): - absolute = cmath.cabs(im).detach().cpu().numpy() - else: - absolute = np.absolute(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, absolute.shape[-1]*basis_norm[1], 0, absolute.shape[-2]*basis_norm[0]] - else: - extent=None - - plt.imshow(absolute, cmap = cmap, extent = extent) - cbar = plt.colorbar() - cbar.set_label('Amplitude (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 + plot_func = lambda x: np.absolute(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, cmap=cmap, cmap_label=cmap_label, + **kwargs) -def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs): +def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', cmap_label='Phase (rad)', **kwargs): """ Plots the phase of a complex array with dimensions NxMx2 If a figure is given explicitly, it will clear that existing figure and @@ -304,6 +354,8 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs) The length units to mark on the plot, default is um cmap : str Default is 'viridis', the colormap to plot with + cmap_label : str + What to label the colorbar when plotting \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -312,49 +364,18 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs) 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): - phase = cmath.cphase(im).detach().cpu().numpy() - else: - phase = np.angle(im) - - if basis is not None: - if isinstance(basis,t.Tensor): - basis = basis.detach().cpu().numpy() - basis_norm = np.linalg.norm(basis, axis = 0) - basis_norm = basis_norm * get_units_factor(units) - - extent = [0, phase.shape[-1]*basis_norm[1], 0, phase.shape[-2]*basis_norm[0]] - else: - extent=None - - - # If the user has matplotlib >=3.0, use the preferred colormap if cmap == 'auto': - try: - plt.imshow(phase, cmap = 'twilight', extent=extent) - except: - plt.imshow(phase, cmap = 'hsv', extent=extent) - else: - plt.imshow(phase, cmap = cmap, extent=extent) + if 'twilight' in plt.colormaps(): + cmap = 'twilight' + elif 'hsv' in plt.colormaps(): + cmap = 'hsv' + else: + raise AttributeError('Neither twilight or hsv colormap exists in this screwed up matplotlib install') - cbar = plt.colorbar() - cbar.set_label('Phase (rad)') - - if basis is not None: - plt.xlabel('X (' + units + ')') - plt.ylabel('Y (' + units + ')') - else: - plt.xlabel('j (pixels)') - plt.ylabel('i (pixels)') - - return fig + plot_func = lambda x: np.angle(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, cmap=cmap, cmap_label=cmap_label, + **kwargs) def plot_amplitude_surfacenorm(): @@ -390,38 +411,9 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs): 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): - im = cmath.torch_to_complex(im.detach().cpu()) - - if basis is not None: - if isinstance(basis,t.Tensor): - basis = basis.detach().cpu().numpy() - basis_norm = np.linalg.norm(basis, axis = 0) - basis_norm = basis_norm * get_units_factor(units) - - extent = [0, im.shape[-1]*basis_norm[1], 0, im.shape[-2]*basis_norm[0]] - else: - extent=None - - colorized = colorize(im) - plt.imshow(colorized, extent=extent) - - if basis is not None: - plt.xlabel('X (' + units + ')') - plt.ylabel('Y (' + units + ')') - else: - plt.xlabel('j (pixels)') - plt.ylabel('i (pixels)') - - return fig - + plot_func = lambda x: colorize(x) + return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, + units=units, cmap=cmap, **kwargs) def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs): diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 88eba2f..f822b94 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -17,8 +17,8 @@ model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2) model.to(device='cuda') dataset.get_as(device='cuda') -#for i, loss in enumerate(model.Adam_optimize(20, dataset, batch_size=50)): -for i, loss in enumerate(model.LBFGS_optimize(20, dataset, lr=1, history_size=5)): +for i, loss in enumerate(model.Adam_optimize(20, dataset, batch_size=50)): +#for i, loss in enumerate(model.LBFGS_optimize(20, dataset, lr=1, history_size=5)): # And we liveplot the updates to the model as they happen print(i,loss) model.inspect(dataset) diff --git a/examples/unified_modes.py b/examples/unified_modes.py index 05dbc27..f8c7c85 100644 --- a/examples/unified_modes.py +++ b/examples/unified_modes.py @@ -10,8 +10,7 @@ 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.UnifiedModePtycho.from_dataset(dataset, oversampling=1,n_modes=3, dm_rank=-1)#, probe_support_radius=90) #model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=1) model.to(device='cuda')