From f4260837bc459d601a48e563e4f0731899b4bd2a Mon Sep 17 00:00:00 2001 From: allevitan Date: Fri, 20 Mar 2026 17:20:52 +0100 Subject: [PATCH] Add title option to plot_image and wrappers; misc fixes - Add title parameter to plot_image, plot_real, plot_imag, plot_amplitude, plot_phase, and plot_colorized - Various fixes to base.py and fancy_ptycho.py Co-Authored-By: Claude Sonnet 4.6 --- examples/fancy_ptycho.py | 5 +- src/cdtools/models/base.py | 38 +++++++-- src/cdtools/models/fancy_ptycho.py | 105 ++++++++++++++++++------- src/cdtools/tools/plotting/plotting.py | 73 ++++++++--------- 4 files changed, 145 insertions(+), 76 deletions(-) diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py index b7e6fbf..8e3a9d1 100644 --- a/examples/fancy_ptycho.py +++ b/examples/fancy_ptycho.py @@ -14,9 +14,7 @@ model = cdtools.models.FancyPtycho.from_dataset( propagation_distance=5e-3, # Propagate the initial probe guess by 5 mm units='mm', # Set the units for the live plots obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix, - exponentiate_obj=False, - panel_plot_mode=True, - plot_level=2, + panel_plot_mode=True, # Organizes the live plots into panels ) if t.cuda.is_available(): @@ -30,7 +28,6 @@ if t.cuda.is_available(): # e.g. estimates of the moments of individual parameters recon = cdtools.reconstructors.AdamReconstructor(model, dataset) - # The learning rate parameter sets the alpha for Adam. # The beta parameters are (0.9, 0.999) by default # The batch size sets the minibatch size diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 8decc4d..99a87f8 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -675,6 +675,7 @@ class CDIModel(t.nn.Module): if not condition(self, dataset): continue + figsize = plot.get('figure_size', None) if self.has_inspect_been_called and \ not replot_all and \ not plt.fignum_exists(plot['title']): @@ -682,22 +683,30 @@ class CDIModel(t.nn.Module): if not self.has_inspect_been_called: fig = plt.figure(plot['title'], + figsize=figsize, constrained_layout=True) else: with plt.rc_context({'figure.raise_window': False}): fig = plt.figure(plot['title'], + figsize = panel_def.get('figure_size', None) constrained_layout=True) try: plot['plot_func'](self, fig) - plt.title(plot['title']) + if plt.gca().get_title().strip() == '': + plt.title(plot['title']) except TypeError: if dataset is not None: try: plot['plot_func'](self, fig, dataset) - plt.title(plot['title']) + if plt.gca().get_title().strip() == '': + plt.title(plot['title']) + except KeyboardInterrupt: + raise except Exception: pass + except KeyboardInterrupt: + raise except Exception: pass @@ -723,6 +732,15 @@ class CDIModel(t.nn.Module): panel_level = panel_def.get('plot_level', 1) if panel_level > self.plot_level: continue # skip entire panel + + panel_condition = panel_def.get('condition', None) + if panel_condition is not None: + try: + if not panel_condition(self): + continue + except TypeError: + if not panel_condition(self, dataset): + continue nrows, ncols = panel_def['grid'] figsize = panel_def.get('figure_size', None) @@ -767,16 +785,22 @@ class CDIModel(t.nn.Module): try: plot['plot_func'](self, subfig) - plt.gca().set_title(plot['title']) + if plt.gca().get_title().strip() == '': + plt.title(plot['title']) except TypeError: if dataset is not None: try: plot['plot_func'](self, subfig, dataset) - plt.gca().set_title(plot['title']) - except TypeError:#Exception: + if plt.gca().get_title().strip() == '': + plt.title(plot['title']) + except KeyboardInterrupt: + raise + except Exception: pass - #except Exception: - # pass + except KeyboardInterrupt: + raise + except Exception: + pass rendered.append(fig) diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index fb9d4da..9ec668b 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -896,11 +896,35 @@ class FancyPtycho(CDIModel): def plot_illumination_intensity(self, fig, dataset): - if not hasattr(self, 'weights') or self.weights.ndim != 1: - raise NotImplementedError('Not yet implemented for OPRP') + if not hasattr(self, 'weights'): + raise NotImplementedError("I don't know how to handle having no weights") + elif self.weights.ndim == 1: + probe_intensities = self.weights.detach().cpu().numpy()**2 + else: + # The big case, with OPRP + probe_matrix = np.zeros([self.probe.shape[0]]*2, + dtype=np.complex64) + np_probes = self.probe.detach().cpu().numpy() + 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()) + + weights = self.weights.detach().cpu().numpy() + + # The outer one is a sum, because the tensordot is what broadcasts the + # probe matrix along the shot dimension - the second one doesn't have to. + weighted_probe_matrices = np.sum(np.tensordot(weights, probe_matrix, axes=1)[...,None] + * weights.conj().transpose((0,2,1))[...,None,:,:], axis=-2) + + basis_probe_intensities = np.trace(probe_matrix, axis1=-2, axis2=-1) + probe_intensities = np.trace(weighted_probe_matrices, axis1=-2, axis2=-1) + + # Imaginary part is already essentially zero up to rounding error + probe_intensities = np.real(probe_intensities / basis_probe_intensities) + p.plot_nanomap( self.corrected_translations(dataset), - self.weights**2, + probe_intensities, fig=fig, cmap='magma', cmap_label='Intensity (a.u.)', @@ -908,7 +932,7 @@ class FancyPtycho(CDIModel): convention='probe', invert_xaxis=True ) - + def plot_translations_and_originals(self, fig, dataset): """Only used to make a plot for the plot list.""" @@ -987,6 +1011,7 @@ class FancyPtycho(CDIModel): (self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe)), fig=fig, + title='Basis Probe', basis=self.probe_basis, units=self.units), }, @@ -997,6 +1022,7 @@ class FancyPtycho(CDIModel): (self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe)), fig=fig, + title='Basis Probe', basis=self.probe_basis, units=self.units), }, @@ -1014,7 +1040,9 @@ class FancyPtycho(CDIModel): 'plot_func': lambda self, fig: p.plot_colorized( (self.probe if self.fourier_probe else tools.propagators.far_field(self.probe)), - fig=fig), + fig=fig, + title='Basis Probe, Fourier', + ), }, { 'title': 'Basis Probes, Fourier Amplitude', @@ -1022,13 +1050,14 @@ class FancyPtycho(CDIModel): 'plot_func': lambda self, fig: p.plot_amplitude( (self.probe if self.fourier_probe else tools.propagators.far_field(self.probe)), - fig=fig), + fig=fig, + title='Basis Probe, Fourier', + ), }, { 'title': 'Illumination Intensity', 'subplot': (0,1), 'plot_func': lambda self, fig, dataset: self.plot_illumination_intensity(fig, dataset), - 'condition': lambda self: hasattr(self, 'weights') and self.weights.ndim == 1 }, { 'title': 'Detector Background', @@ -1047,11 +1076,48 @@ class FancyPtycho(CDIModel): }, ], }, + { + 'title': 'Unstable Probe Refinement Details', + 'plot_level': 2, + 'figure_size': (9,3.5), + 'grid': (1,2), + 'condition': lambda self: len(self.weights.shape) >= 2, + 'plots': [ + { + 'title': '% of Power in Top Mode', + 'subplot': (0,0), + 'plot_func': lambda self, fig, dataset: p.plot_nanomap( + self.corrected_translations(dataset), + 100 * t.stack([ + analysis.calc_mode_power_fractions( + self.probe.data, + weight_matrix=self.weights.data[i])[0] + for i in range(self.weights.shape[0]) + ], dim=0), + fig=fig, + units=self.units), + 'condition': lambda self: len(self.weights.shape) >= 2 + }, + { + 'title': 'Average Weight Matrix Amplitudes', + 'subplot': (0,1), + 'plot_func': lambda self, fig: p.plot_amplitude( + np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0), + fig=fig), + 'condition': lambda self: len(self.weights.shape) >= 2 + }, + ] + } ] plot_list = [ + {'title': 'Quantum Efficiency Mask', + 'plot_level': 2, + 'plot_func': lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig), + 'condition': lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None)}, {'title': 'Per-Exposure Probe Intensity', 'plot_level': 3, + 'figure_size': (8,5.3), 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( dataset, fig=fig, @@ -1061,6 +1127,7 @@ class FancyPtycho(CDIModel): 'condition': lambda self: len(self.weights.shape) >= 2}, {'title': 'Per-Exposure Probe Amplitudes', 'plot_level': 3, + 'figure_size': (8,5.3), 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( dataset, fig=fig, @@ -1070,6 +1137,7 @@ class FancyPtycho(CDIModel): 'condition': lambda self: len(self.weights.shape) >= 2}, {'title': 'Per-Exposure Probe Phases', 'plot_level': 3, + 'figure_size': (8,5.3), 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( dataset, fig=fig, @@ -1077,29 +1145,6 @@ class FancyPtycho(CDIModel): image_title='Probe Phases (scroll to view modes)', image_colorbar_title='Probe Phase'), 'condition': lambda self: len(self.weights.shape) >= 2}, - {'title': 'Average Weight Matrix Amplitudes', - 'plot_level': 1, - 'plot_func': lambda self, fig: p.plot_amplitude( - np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0), - fig=fig), - 'condition': lambda self: len(self.weights.shape) >= 2}, - {'title': '% of Power in Top Mode', - 'plot_level': 3, - 'plot_func': lambda self, fig, dataset: p.plot_nanomap( - self.corrected_translations(dataset), - 100 * t.stack([ - analysis.calc_mode_power_fractions( - self.probe.data, - weight_matrix=self.weights.data[i])[0] - for i in range(self.weights.shape[0]) - ], dim=0), - fig=fig, - units=self.units), - 'condition': lambda self: len(self.weights.shape) >= 2}, - {'title': 'Quantum Efficiency Mask', - 'plot_level': 3, - 'plot_func': lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig), - 'condition': lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None)}, ] diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py index 367c906..7069130 100644 --- a/src/cdtools/tools/plotting/plotting.py +++ b/src/cdtools/tools/plotting/plotting.py @@ -105,6 +105,7 @@ def plot_image( vmin=None, vmax=None, interpolation=None, + title=None, **kwargs ): """Plots an image with a colorbar and on an appropriate spatial grid @@ -169,12 +170,13 @@ def plot_image( # 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() - try: - title = fig.axes[0].get_title() - except IndexError: - title = '' + if title is not None: + ax_title = title + else: + try: + ax_title = fig.axes[0].get_title() + except IndexError: + ax_title = '' # 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 @@ -190,7 +192,6 @@ def plot_image( # By only updating the data, and not redrawing the fig, we # don't "reset" the home positions of the other if hasattr(fig, '_current_im'): - print('Just changing data') fig._current_im.set_data(to_plot) fig._current_im.autoscale() # We need to go to the "home" position before updating it @@ -200,10 +201,11 @@ def plot_image( if fig.canvas.toolbar is not None: fig.canvas.toolbar.home() fig.canvas.toolbar.update() - # Replace existing mode number - for artist in fig.texts: - artist.set_text(f'Mode {fig.plot_idx}') - + + if len(im.shape) >= 3: + base = title if title is not None else '('.join(ax_title.split('(')[:-1])[:-1] + fig.axes[0].set_title(base + f' ({fig.plot_idx+1} of {num_images})') + return fig fig.clear() @@ -291,10 +293,10 @@ def plot_image( ax.set_xlabel('j (pixels)') ax.set_ylabel('i (pixels)') - ax.set_title(title) - + if title is not None: + ax.set_title(ax_title) if len(im.shape) >= 3: - fig.text(0.03, 0.03, f'Mode {fig.plot_idx}', fontsize=14) + ax.set_title(ax_title + f' ({fig.plot_idx+1} of {num_images})') if fig.canvas.toolbar is not None: fig.canvas.toolbar.update() @@ -335,7 +337,7 @@ def plot_image( return result_fig -def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs): +def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', title=None, **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 @@ -369,11 +371,11 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_ 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) + title=title, **kwargs) -def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs): +def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', title=None, **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 @@ -407,10 +409,10 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_ 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) + title=title, **kwargs) -def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Amplitude (a.u.)', **kwargs): +def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Amplitude (a.u.)', title=None, **kwargs): """Plots the amplitude of a complex array with dimensions NxM If a figure is given explicitly, it will clear that existing figure and @@ -444,7 +446,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', 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) + title=title, **kwargs) def plot_phase( @@ -456,6 +458,7 @@ def plot_phase( cmap_label='Phase (rad)', vmin=None, vmax=None, + title=None, **kwargs ): """ Plots the phase of a complex array with dimensions NxM @@ -506,14 +509,14 @@ def plot_phase( return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, units=units, cmap=cmap, cmap_label=cmap_label, - vmin=vmin,vmax=vmax, + vmin=vmin, vmax=vmax, title=title, **kwargs) def plot_amplitude_surfacenorm(): pass -def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs): +def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', title=None, **kwargs): """ Plots the colorized version of a complex array with dimensions NxM The darkness corresponds to the intensity of the image, and the color @@ -545,7 +548,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs): """ plot_func = lambda x: colorize(x) return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, - units=units, show_cbar=False, **kwargs) + units=units, show_cbar=False, title=title, **kwargs) def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, clear_fig=True, label=None, color=None, marker='.', **kwargs): @@ -713,20 +716,19 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # mode, i.e. on a figure that already has this thing showing. if fig is None: - fig = plt.figure(figsize=(8,5.3)) + fig = plt.figure(figsize=(20,4.5), constrained_layout=True) else: - plt.figure(fig.number) - plt.gcf().clear() + fig = plt.figure(fig.number, figsize=(20,4.5), constrained_layout=True) + fig.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) + gs = fig.add_gridspec(2, 2, height_ratios=[0.9,0.1], width_ratios=[1,1]) - 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]) + axes = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1])] + axslider = fig.add_subplot(gs[1, :]) # full width # This gets the set of sizes for the points in the nanomap def calculate_sizes(idx): @@ -779,11 +781,12 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non 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) + axes[0].set_aspect('equal') 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) + 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 @@ -836,8 +839,8 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal', format='%.2e', - ticks=ticker.LinearLocator(numticks=5), - pad=0.17,fraction=0.1) + 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))