diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index 2331dcc..e16b664 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -62,7 +62,7 @@ class PolarizedFancyPtycho(FancyPtycho): @classmethod def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True): - model = FancyPtycho.from_dataset(dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um') + model = FancyPtycho.from_dataset(dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True) # Mutate the class to its subclass diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py index b5c65a8..f1451b2 100644 --- a/CDTools/tools/plotting/plotting.py +++ b/CDTools/tools/plotting/plotting.py @@ -17,7 +17,11 @@ 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_with_images'] + 'plot_nanomap_with_images', + 'polarized_plot_component_amplitudes', + 'polarized_plot_phase_ret', + 'polarized_plot_global_phases', + 'polarized_plot_ellipses'] def colorize(z): @@ -84,7 +88,7 @@ def get_units_factor(units): 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. @@ -94,7 +98,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', 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 ---------- @@ -120,7 +124,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', 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, @@ -142,7 +146,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', 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 @@ -151,9 +155,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', 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): @@ -181,7 +185,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', plt.xlabel('j (pixels)') plt.ylabel('i (pixels)') - + plt.title(title) if len(im.shape) >= 3: @@ -194,14 +198,14 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', 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': @@ -217,9 +221,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', 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 @@ -256,7 +260,7 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_ 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', cmap_label='Imaginary Part (a.u.)', **kwargs): @@ -304,7 +308,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', If a basis is explicitly passed, the image will be plotted in real-space coordinates. - + Parameters ---------- im : array @@ -520,12 +524,12 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr 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 @@ -558,11 +562,11 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non 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 @@ -571,29 +575,29 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non 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], @@ -614,7 +618,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # 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) @@ -634,7 +638,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # 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: @@ -655,9 +659,9 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non 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), @@ -665,7 +669,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non 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 @@ -675,7 +679,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # 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: @@ -686,22 +690,22 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non 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") @@ -715,7 +719,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # 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] @@ -733,7 +737,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non 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': @@ -742,7 +746,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # 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) @@ -753,16 +757,16 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # 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] + 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