Make a program to visualize the shot to shot fluctuations when using the unified mode model

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)

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:

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),

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)

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)
+
+

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)

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)