diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index ce46fc4..0b84d54 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -198,169 +198,27 @@ class Ptycho2DDataset(CDataset): position. """ - # We start by making the figure and axes - fig, axes = plt.subplots(1,2,figsize=(8,5.3)) - fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96]) - axslider = plt.axes([0.15,0.06,0.75,0.03]) - - # - # Then we define some helper functions for getting the right data - # that are used both in the initial setup and the updates - # - - def get_data(idx): + def get_images(idx): inputs, output = self[idx] meas_data = output.detach().cpu().numpy() if hasattr(self, 'mask') and self.mask is not None: mask = self.mask.detach().cpu().numpy() else: mask = 1 - - return mask, meas_data - - - def calculate_sizes(idx): - bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted()) - s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch - s0 /= 4 # A rough value to make the size work out - s = np.ones(len(self)) * s0 - - s[idx] *= 4 - return s - - def update_colorbar(im): - # - # This solves the problem of the colorbar being changed - # when the forward and back buttons are used!!! - # - if hasattr(im, 'norecurse') and im.norecurse: - im.norecurse=False - return - - im.norecurse=True - # This is needed to update the colorbar - # only change limits if array contains multiple values - if np.min(im.get_array()) != np.max(im.get_array()): - im.set_clim(vmin=np.min(im.get_array()), - vmax=np.max(im.get_array())) - - # - # The meatiest part of this program, here we just go through and - # set up the plot how we want it - # - - # First we set up the left-hand plot, which shows an overview map - axes[0].set_title('Relative Displacement Map') + + if logarithmic: + return np.log(meas_data) / np.log(10) * mask + else: + return meas_data * mask translations = self.translations.detach().cpu().numpy() nanomap_values = (self.mask.to(t.float32) * self.patterns).sum(dim=(1,2)).detach().cpu().numpy() - s = calculate_sizes(0) - units_factor = plotting.get_units_factor(units) - nanomap = axes[0].scatter(units_factor * translations[:,0],units_factor * translations[:,1],s=s,c=nanomap_values, picker=True) - - axes[0].invert_xaxis() - axes[0].set_facecolor('k') - axes[0].set_xlabel('Translation x ('+units+')', labelpad=1) - axes[0].set_ylabel('Translation y ('+units+')', labelpad=1) - cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal', - format='%.2e', - ticks=ticker.LinearLocator(numticks=5), - pad=0.17,fraction=0.1) - cb1.ax.set_title('Integrated Intensity', size="medium", pad=5) - cb1.ax.tick_params(labelrotation=20) - - - # Now we set up the second plot, which shows the individual - # diffraction patterns - axes[1].set_title('Diffraction Pattern') - mask, meas_data = get_data(0) if logarithmic: - meas = axes[1].imshow(np.log(meas_data) / np.log(10) * mask) + cbar_title='Log Base 10 of Diffraction Intensity' else: - meas = axes[1].imshow(meas_data * mask) + cbar_title='Diffraction Intensity' + + plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title) - cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal', - format='%.2e', - ticks=ticker.LinearLocator(numticks=5), - pad=0.17,fraction=0.1) - cb2.ax.tick_params(labelrotation=20) - cb2.ax.set_title('Pixel Intensity', size="medium", pad=5) - cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas)) - - # This function handles all the updating, except for moving the - # slider value. This is done because the slider widget is - # ultimately responsible for triggering an update, so all other - # updates are done by changing the slider widget value which - # then triggers this - def update(idx): - # We have to explicitly make it an integer because the slider will - # output floats (even if they are still integer-valued) - idx = int(idx) - - # Get the new data for this index - mask, meas_data = get_data(idx) - - # Now we resize the nanomap to show the new selection - axes[0].collections[0].set_sizes(calculate_sizes(idx)) - - # And we update the data in the image as well - meas = axes[1].images[-1] - if logarithmic: - meas.set_data(np.log(meas_data) / np.log(10) * mask) - else: - meas.set_data(meas_data * mask) - - update_colorbar(meas) - - - # - # Now we define the functions to handle various kinds of events - # that can be thrown our way - # - - # We start by creating the slider here, so it can be used - # by the update hooks. - slider = Slider(axslider, 'Pattern #', 0, len(self)-1, valstep=1, valfmt="%d") - - # This handles scroll wheel and keypress events - def on_action(event): - # Otherwise the if statements can throw errors when the - # event type isn't right, this way they just don't trigger - if not hasattr(event, 'button'): - event.button = None - if not hasattr(event, 'key'): - event.key = None - - if event.key == 'up' or event.button == 'up' or event.key == 'right': - 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) - - # This handles "pick" events in the nanomap - def on_pick(event): - # If we don't filter on type of event, this will also capture, - # for example, scroll events that happen over the nanomap - if event.mouseevent.button == 1: - slider.set_val(event.ind[0]) - - - # Here we connect the various update functions - fig.canvas.mpl_connect('pick_event',on_pick) - fig.canvas.mpl_connect('key_press_event',on_action) - fig.canvas.mpl_connect('scroll_event',on_action) - slider.on_changed(update) - - # Throw an extra update into the mix just to get rid of any things - # (like the nanomap dot sizes) that otherwise would change on the - # first update - update(0) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index fdfa7a4..0166734 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -472,7 +472,7 @@ class CDIModel(t.nn.Module): plotter = plots[1] if figs is None: - fig = plt.figure() + fig = plt.figure() self.figs.append(fig) else: fig = figs[idx] @@ -480,6 +480,7 @@ class CDIModel(t.nn.Module): try: plotter(self,fig) plt.title(name) + except TypeError as e: if dataset is not None: try: diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 683a320..310c3e9 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -271,7 +271,6 @@ class FancyPtycho(CDIModel): translations, surface_normal=self.surface_normal) pix_trans -= self.min_translation - # We then add on any recovered translation offset, if they exist if self.translation_offsets is not None: pix_trans += self.translation_scale * self.translation_offsets[index] @@ -392,6 +391,7 @@ class FancyPtycho(CDIModel): 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 + def get_rhos(self): # If this is the general unified mode model if self.weights.dim() >= 2: @@ -470,13 +470,61 @@ class FancyPtycho(CDIModel): self.probe.data = cmath.complex_to_torch(ortho_probes).to( device=self.probe.device,dtype=self.probe.dtype) + + + def plot_wavefront_variation(self, dataset,fig=None,mode='amplitude',**kwargs): + def get_probes(idx): + basis_prs = self.probe * self.probe_support[...,:,:] + prs = t.sum(cmath.cmult(self.weights[idx,:,:,None,None,:], + basis_prs), axis=-4) + ortho_probes = analysis.orthogonalize_probes(prs) + + #return np.abs(cmath.torch_to_complex(prs.detach().cpu())) + if mode.lower() == 'amplitude': + return np.abs(cmath.torch_to_complex(ortho_probes.detach().cpu())) + if mode.lower() == 'root_sum_intensity': + return np.sum(np.abs(cmath.torch_to_complex(ortho_probes.detach().cpu()))**2,axis=0) + if mode.lower() == 'phase': + return np.angle(cmath.torch_to_complex(ortho_probes.detach().cpu())) + + probe_matrix = np.zeros([self.probe.shape[0]]*2, + dtype=np.complex64) + np_probes = cmath.torch_to_complex(self.probe.detach().cpu()) + for i in range(probe_matrix.shape[0]): + for j in range(probe_matrix.shape[0]): + probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj()) - # Needs to be updated to allow for plotting to an existing figure + weights = cmath.torch_to_complex(self.weights.detach().cpu()) + + probe_intensities = np.sum(np.tensordot(weights,probe_matrix,axes=1)* + weights.conj(),axis=2) + + # Imaginary part is already essentially zero up to rounding error + probe_intensities = np.real(probe_intensities) + + values = np.sum(probe_intensities,axis=1) + if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity': + cmap = 'viridis' + else: + cmap = 'twilight' + + p.plot_nanomap_with_images(self.corrected_translations(dataset), get_probes, values=values, fig=fig, units=self.units, basis=self.probe_basis, nanomap_colorbar_title='Total Probe Intensity',cmap=cmap,**kwargs), + + plot_list = [ - ('Probe Amplitude (scroll to view modes)', + ('', + lambda self, fig, dataset: self.plot_wavefront_variation(dataset,fig=fig,mode='root_sum_intensity',image_title='Root Summed Probe Intensities',image_colorbar_title='Square Root of Intensity'), + lambda self: len(self.weights.shape) >= 2), + ('', + lambda self, fig, dataset: self.plot_wavefront_variation(dataset,fig=fig,mode='amplitude',image_title='Probe Amplitudes (scroll to view modes)',image_colorbar_title='Probe Amplitude'), + lambda self: len(self.weights.shape) >= 2), + ('', + lambda self, fig, dataset: self.plot_wavefront_variation(dataset,fig=fig,mode='phase',image_title='Probe Phases (scroll to view modes)',image_colorbar_title='Probe Phase'), + lambda self: len(self.weights.shape) >= 2), + ('Basis Probe Amplitudes (scroll to view modes)', lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis,units=self.units)), - ('Probe Phase (scroll to view modes)', + ('Basis Probe Phases (scroll to view modes)', lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis,units=self.units)), ('Average Density Matrix Amplitudes', lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()),axis=0), fig=fig), diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index 9ec50e4..8b449c2 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -4,9 +4,8 @@ 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 cmath, analysis, image_processing from CDTools.tools import plotting as p -from CDTools.tools import analysis from matplotlib import pyplot as plt from datetime import datetime import numpy as np @@ -31,7 +30,7 @@ class Multislice2DPtycho(CDIModel): subpixel=True, exponentiate_obj=True, fourier_probe=False, - apodization=None, + prevent_aliasing=True, phase_only=False, units='um'): @@ -60,6 +59,7 @@ class Multislice2DPtycho(CDIModel): self.fourier_probe = fourier_probe self.units = units self.phase_only=phase_only + self.prevent_aliasing=prevent_aliasing if mask is None: self.mask = mask @@ -111,23 +111,21 @@ class Multislice2DPtycho(CDIModel): self.probe_support = t.Tensor(probe_support).to(t.float32) - if apodization is not None: - self.apodization = t.Tensor(apodization).to(t.float32) - else: - self.apodization = None - self.oversampling = oversampling spacing = np.linalg.norm(self.probe_basis,axis=0) shape = np.array(self.probe.shape[1:-1]) + if prevent_aliasing: + shape *= 2 + spacing /= 2 self.bandlimit = bandlimit - + self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=self.bandlimit) @classmethod - def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, apodize_prop=False, phase_only=False): + def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -216,17 +214,6 @@ class Multislice2DPtycho(CDIModel): if not replicate_slice: obj = t.stack([obj]*nz) - if apodize_prop: - shape = probe.shape[-3:-1] - #window_x = (1+np.cos(np.pi+2*np.pi*np.arange(shape[0])/shape[0]))/2 - #window_y = (1+np.cos(np.pi+2*np.pi*np.arange(shape[1])/shape[1]))/2 - window_x = np.cos(-np.pi/2+np.pi*np.arange(shape[0])/shape[0]) - window_y = np.cos(-np.pi/2+np.pi*np.arange(shape[1])/shape[1]) - Wx, Wy = np.meshgrid(window_x,window_y, indexing='ij') - apodization=t.Tensor(Wx*Wy).to(dtype=probe.dtype) - else: - apodization=None - det_geo = dataset.detector_geometry translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) @@ -277,8 +264,8 @@ class Multislice2DPtycho(CDIModel): subpixel=subpixel, exponentiate_obj=exponentiate_obj, units=units, fourier_probe=fourier_probe, - apodization=apodization, - phase_only=phase_only) + phase_only=phase_only, + prevent_aliasing=prevent_aliasing) def interaction(self, index, translations): @@ -289,13 +276,19 @@ class Multislice2DPtycho(CDIModel): if self.translation_offsets is not None: pix_trans += self.translation_scale * self.translation_offsets[index] - + # This restricts the basis probes to stay within the probe support basis_prs = self.probe * self.probe_support[...,:,:] - + + # For a Fourier-space probe, we take an IFT + if self.fourier_probe: + basis_prs = tools.propagators.inverse_far_field(basis_prs) + + if self.prevent_aliasing: + pix_trans = pix_trans * 2 + basis_prs = image_processing.fourier_upsample(basis_prs) + # Now we construct the probes for each shot from the basis probes - # This is fine to happen in real or Fourier space, so we do it - # regardless of whether fourier_probe is True Ws = self.weights[index] if len(self.weights[0].shape) == 0: @@ -310,9 +303,6 @@ class Multislice2DPtycho(CDIModel): prs = t.sum(cmath.cmult(Ws[...,None,None,:], basis_prs), axis=-4) - # For a Fourier-space probe, we take an IFT - if self.fourier_probe: - prs = tools.propagators.inverse_far_field(prs) if self.exponentiate_obj: if self.phase_only: @@ -321,7 +311,10 @@ class Multislice2DPtycho(CDIModel): obj = cmath.cexpi(self.obj) else: obj = self.obj - + + if self.prevent_aliasing: + obj = image_processing.fourier_upsample(obj) + exit_waves = self.probe_norm * prs for i in range(self.nz): # If only one object slice @@ -336,8 +329,6 @@ class Multislice2DPtycho(CDIModel): exit_waves = tools.interactions.ptycho_2D_round( exit_waves, obj, pix_trans, multiple_modes=True) - if self.apodization is not None: - exit_waves = exit_waves * self.apodization[None,:,:,None] @@ -351,8 +342,6 @@ class Multislice2DPtycho(CDIModel): exit_waves = tools.interactions.ptycho_2D_round( exit_waves, obj[i], pix_trans, multiple_modes=True) - if self.apodization is not None: - exit_waves = exit_waves * self.apodization[None,:,:,None] if i < self.nz-1: #on all but the last iteration exit_waves = tools.propagators.near_field( @@ -362,11 +351,12 @@ class Multislice2DPtycho(CDIModel): def forward_propagator(self, wavefields): - return tools.propagators.far_field(wavefields) + left = [self.probe.shape[-3]//2,self.probe.shape[-2]//2] + right = [self.probe.shape[-3]//2+self.probe.shape[-3], + self.probe.shape[-2]//2+self.probe.shape[-2]] - - def backward_propagator(self, wavefields): - return tools.propagators.inverse_far_field(wavefields) + return tools.propagators.far_field(wavefields)[...,left[0]:right[0], + left[1]:right[1],:] def measurement(self, wavefields): @@ -397,9 +387,6 @@ class Multislice2DPtycho(CDIModel): if self.mask is not None: self.mask = self.mask.to(*args, **kwargs) - if self.apodization is not None: - self.apodization = self.apodization.to(*args,**kwargs) - self.min_translation = self.min_translation.to(*args,**kwargs) self.probe_basis = self.probe_basis.to(*args,**kwargs) diff --git a/CDTools/tools/image_processing/image_processing.py b/CDTools/tools/image_processing/image_processing.py index 4ea5d3c..8681262 100644 --- a/CDTools/tools/image_processing/image_processing.py +++ b/CDTools/tools/image_processing/image_processing.py @@ -10,11 +10,11 @@ a way that it is safe to include them in automatic differentiation models. from __future__ import division, print_function, absolute_import import numpy as np import torch as t -from CDTools.tools import cmath +from CDTools.tools import cmath, propagators __all__ = ['centroid', 'centroid_sq', 'sinc_subpixel_shift', 'find_subpixel_shift', 'find_pixel_shift', 'find_shift', - 'convolve_1d'] + 'convolve_1d', 'fourier_upsample'] def centroid(im, dims=2): @@ -325,3 +325,17 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True): return conv_im[...,0] else: return conv_im + + +def fourier_upsample(ims): + upsampled = t.zeros(ims.shape[:-3]+(2*ims.shape[-3],2*ims.shape[-2])+(2,), + dtype=ims.dtype, + device=ims.device) + left = [ims.shape[-3]//2,ims.shape[-2]//2] + right = [ims.shape[-3]//2+ims.shape[-3], + ims.shape[-2]//2+ims.shape[-2]] + + upsampled[...,left[0]:right[0],left[1]:right[1],:] = propagators.far_field(ims) + return propagators.inverse_far_field(upsampled) + + diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py index 313a88c..c454f2e 100644 --- a/CDTools/tools/plotting/plotting.py +++ b/CDTools/tools/plotting/plotting.py @@ -13,11 +13,14 @@ import torch as t import numpy as np import matplotlib.pyplot as plt from matplotlib.colors import hsv_to_rgb +from matplotlib.widgets import Slider +from matplotlib import ticker, patheffects __all__ = ['colorize', 'plot_amplitude', 'plot_phase', 'plot_colorized', 'plot_translations', 'get_units_factor', - 'plot_nanomap', 'plot_real', 'plot_imag'] + 'plot_nanomap', 'plot_real', 'plot_imag', + 'plot_nanomap_with_images'] def colorize(z): @@ -516,3 +519,254 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr plt.colorbar() return fig + + +def plot_nanomap_with_images(translations, get_image_func, values=None, mask=None, basis=None, fig=None, nanomap_units='$\\mu$m', image_units='$\\mu$m', convention='probe', image_title='Image', image_colorbar_title='Image Amplitude', nanomap_colorbar_title='Integrated Intensity', cmap='viridis', **kwargs): + """Plots a nanomap, with an image or stack of images for each point + + In many situations, ptychography data or the output of ptychography + reconstructions is formatted as a set of images associated with various + points in real space. This function is designed to allow for browsing + through this kind of data, by making it possible to visualize a + + """ + + # This should pull heavily from the dataset.inspect function + # In fact, I should be able to replace most of that function with a + # call to this function once it's built + # We start by making the figure and axes + + # The key will be writing this so it works okay when called in "update" + # mode, i.e. on a figure that already has this thing showing. + + if fig is None: + fig = plt.figure(figsize=(8,5.3)) + else: + plt.figure(fig.number) + plt.gcf().clear() + if hasattr(fig, 'nanomap_cids'): + for cid in fig.nanomap_cids: + fig.canvas.mpl_disconnect(cid) + + # Does figsize work with the fig.subplots, or just for plt.subplots? + axes = fig.subplots(1,2) + + fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96]) + plt.subplots_adjust(wspace=0.25) #avoids overlap of labels with plots + axslider = plt.axes([0.15,0.06,0.75,0.03]) + + # This gets the set of sizes for the points in the nanomap + def calculate_sizes(idx): + bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted()) + s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch + s0 /= 4 # A rough value to make the size work out + s = np.ones(translations.shape[0]) * s0 + + s[idx] *= 4 + return s + + + def update_colorbar(im): + # + # This solves the problem of the colorbar being changed + # when the forward and back buttons are used!!! + # + if hasattr(im, 'norecurse') and im.norecurse: + im.norecurse=False + return + + im.norecurse=True + # This is needed to update the colorbar + # only change limits if array contains multiple values + if np.min(im.get_array()) != np.max(im.get_array()): + im.set_clim(vmin=np.min(im.get_array()), + vmax=np.max(im.get_array())) + + # + # The meatiest part of this program, here we just go through and + # set up the plot how we want it + # + + # First we set up the left-hand plot, which shows an overview map + axes[0].set_title('Relative Displacement Map') + + translations = translations.detach().cpu().numpy() + + if convention.lower() != 'probe': + translations = translations * -1 + + s = calculate_sizes(0) + + nanomap_units_factor = get_units_factor(nanomap_units) + nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0], + nanomap_units_factor * translations[:,1], + s=s,c=values, picker=True) + + axes[0].invert_xaxis() + axes[0].set_facecolor('k') + axes[0].set_xlabel('Translation x ('+nanomap_units+')', labelpad=1) + axes[0].set_ylabel('Translation y ('+nanomap_units+')', labelpad=1) + cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal', + format='%.2e', + ticks=ticker.LinearLocator(numticks=5), + pad=0.17,fraction=0.1) + cb1.ax.set_title(nanomap_colorbar_title, size="medium", pad=5) + cb1.ax.tick_params(labelrotation=20) + if values is None: + # This seems to do a good job of leaving the appropriate space + # where the colorbar should have been to avoid stretching the + # nanomap plot, while still not showing the (now useless) colorbar. + cb1.remove() + + # Now we set up the second plot, which shows the individual + # diffraction patterns + axes[1].set_title(image_title) + #Plot in a basis if it exists, otherwise dont + if basis is not None: + axes[1].set_xlabel('X (' + image_units + ')') + axes[1].set_ylabel('Y (' + image_units + ')') + + example_im = get_image_func(0) + if isinstance(example_im, t.Tensor): + example_im = example_im.cpu().numpy() + + 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(image_units) + + extent = [0, example_im.shape[-1]*basis_norm[1], 0, + example_im.shape[-2]*basis_norm[0]] + else: + axes[1].set_xlabel('j (pixels)') + axes[1].set_ylabel('i (pixels)') + extent=None + + im=get_image_func(0) + if len(im.shape) >= 3: + im_idx=0 + axes[1].image_idx = im_idx + im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx] + axes[1].text_box = axes[1].text(0.98, 0.98, str(im_idx), color='w', + fontsize=14, + horizontalalignment='right', + verticalalignment='top', + transform=axes[1].transAxes) + axes[1].text_box.set_path_effects( + [patheffects.Stroke(linewidth=2, foreground='black'), + patheffects.Normal()]) + + meas = axes[1].imshow(im, extent=extent, cmap=cmap) + + cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal', + format='%.2e', + ticks=ticker.LinearLocator(numticks=5), + pad=0.17,fraction=0.1) + cb2.ax.tick_params(labelrotation=20) + cb2.ax.set_title(image_colorbar_title, size="medium", pad=5) + cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas)) + + # This function handles all the updating, except for moving the + # slider value. This is done because the slider widget is + # ultimately responsible for triggering an update, so all other + # updates are done by changing the slider widget value which + # then triggers this + def update(idx, im_idx=None): + # We have to explicitly make it an integer because the slider will + # output floats (even if they are still integer-valued) + idx = int(idx) + + # Get the new data for this index + im = get_image_func(idx) + if len(im.shape) >= 3: + if im_idx == None and hasattr(axes[1],'image_idx'): + im_idx = axes[1].image_idx + elif im_idx == None: + im_idx=0 + axes[1].image_idx = im_idx + axes[1].text_box.set_text(str(im_idx)) + im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx] + + # Now we resize the nanomap to show the new selection + axes[0].collections[0].set_sizes(calculate_sizes(idx)) + + # And we update the data in the image as well + + ax_im = axes[1].images[-1] + ax_im.set_data(im) + update_colorbar(ax_im) + + + # + # Now we define the functions to handle various kinds of events + # that can be thrown our way + # + + # We start by creating the slider here, so it can be used + # by the update hooks. + slider = Slider(axslider, 'Image #', 0, translations.shape[0]-1, valstep=1, valfmt="%d") + + # This handles scroll wheel and keypress events + def on_action(event): + im = get_image_func(0) + if event.inaxes is axes[1] and len(im.shape) >=3: + # This is triggered if the data to display has more than 2 + # dimensions (i.e. is an image stack) and the event originates + # while the mouse is within the image display + im = im.reshape(-1,im.shape[-2],im.shape[-1]) + im_idx = axes[1].image_idx + + if event.key == 'up' or event.button == 'up' \ + or event.key == 'left': + im_idx = (im_idx - 1) % im.shape[0] + if event.key == 'down' or event.button == 'down' \ + or event.key == 'right': + im_idx = (im_idx + 1) % im.shape[0] + + axes[1].image_idx=im_idx + slider.set_val(slider.val)#update(slider.val,im_idx=im_idx) + return # This prevents the rest from also happening + + # Otherwise the if statements can throw errors when the + # event type isn't right, this way they just don't trigger + if not hasattr(event, 'button'): + event.button = None + if not hasattr(event, 'key'): + event.key = None + + if event.key == 'up' or event.button == 'up' or event.key == 'left': + idx = slider.val - 1 + elif event.key == 'down' or event.button == 'down' or event.key == 'right': + 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) % translations.shape[0] + slider.set_val(idx) + + # This handles "pick" events in the nanomap + def on_pick(event): + # If we don't filter on type of event, this will also capture, + # for example, scroll events that happen over the nanomap + if event.mouseevent.button == 1: + slider.set_val(event.ind[0]) + + + # Here we connect the various update functions + cid1 = fig.canvas.mpl_connect('pick_event',on_pick) + cid2 = fig.canvas.mpl_connect('key_press_event',on_action) + cid3 = fig.canvas.mpl_connect('scroll_event',on_action) + # It's so dumb that matplotlib doesn't automatically track this for you + fig.nanomap_cids = [cid1,cid2,cid3] + slider.on_changed(update) + + # Throw an extra update into the mix just to get rid of any things + # (like the nanomap dot sizes) that otherwise would change on the + # first update + update(0) diff --git a/examples/unified_modes.py b/examples/unified_modes.py index ac64690..28d36ab 100644 --- a/examples/unified_modes.py +++ b/examples/unified_modes.py @@ -1,14 +1,17 @@ from __future__ import division, print_function, absolute_import import CDTools +from CDTools.tools import plotting as p from matplotlib import pyplot as plt import pickle +import torch as t +from CDTools.tools import cmath filename = 'example_data/lab_ptycho_data.cxi' dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename) -# dataset.inspect() -# plt.show() +#dataset.inspect(units='mm') +#plt.show() # dm_rank=-1 tells it to use a full-rank unified mode approximation model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=1,n_modes=3, dm_rank=-1) @@ -18,13 +21,13 @@ model.to(device='cuda') dataset.get_as(device='cuda') model.translation_offsets.requires_grad = False -for i, loss in enumerate(model.Adam_optimize(20, dataset)): +for i, loss in enumerate(model.Adam_optimize(200, dataset)): model.inspect(dataset) print(i,loss) model.tidy_probes() -for i, loss in enumerate(model.Adam_optimize(20, dataset, lr=0.0001)): +for i, loss in enumerate(model.Adam_optimize(200, dataset, lr=0.0001)): model.inspect(dataset) print(i,loss)