diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index df347b6..f2823b7 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -58,12 +58,14 @@ class CDIModel(t.nn.Module): functions. """ - def __init__(self): + def __init__(self, panel_plot_mode=False, plot_level=np.inf): super(CDIModel, self).__init__() self.loss_history = [] self.training_history = '' self.epoch = 0 + self.panel_plot_mode = panel_plot_mode + self.plot_level = plot_level def from_dataset(self, dataset): raise NotImplementedError() @@ -547,96 +549,238 @@ class CDIModel(t.nn.Module): return msg - # By default, the plot_list is empty + # By default, the plot lists are empty + plot_panel_list = [] plot_list = [] - def inspect(self, dataset=None, update=True): - """Plots all the plots defined in the model's plot_list attribute + def inspect(self, dataset=None, replot_all=False): + """Plots all the plots defined in the model's plot_panel_list and plot_list attributes - If update is set to True, it will update any previously plotted set - of plots, if one exists, and then redraw them. Otherwise, it will - plot a new set, and any subsequent updates will update the new set + Updates any previously plotted figures that are still open. Figures + that have been closed are left closed unless replot_all=True. Optionally, a dataset can be passed, which will allow plotting of any registered plots which need to incorporate some information from the dataset (such as geometry or a comparison with measured data). - Plots can be registered in any subclass by defining the plot_list - attribute. This should be a list of tuples in the following format: - ( 'Plot Title', function_to_generate_plot(self), - function_to_determine_whether_to_plot(self)) + Plots can be registered in any subclass by defining plot_panel_list + and/or plot_list class attributes. See the CDIModel documentation for + the expected dict-based format of each. - Where the third element in the tuple (a function that returns - True if the plot is relevant) is not required. + When panel_plot_mode=True (set in __init__), plot_panel_list entries + are rendered as multi-subplot figures. When False (the default), + each subplot in plot_panel_list is rendered as its own figure, + prepended to any standalone plot_list entries. + + The plot_level attribute (set in __init__, default np.inf) controls + which plots are shown: a panel or standalone plot is only shown when + its plot_level <= self.plot_level. Parameters ---------- dataset : CDataset Optional, a dataset matched to the model type - update : bool, default: True - Whether to update existing plots or plot new ones + replot_all : bool, default: False + If True, recreate figures that were previously closed by the user. """ - # We find or create all the figures - first_update = False - if update and hasattr(self, 'figs') and self.figs: - figs = self.figs - elif update: - figs = None - self.figs = [] - first_update = True + plot_panel_list = getattr(self, 'plot_panel_list', None) or [] + plot_list = getattr(self, 'plot_list', None) or [] + + if self.panel_plot_mode and plot_panel_list: + self._inspect_panel(dataset=dataset, replot_all=replot_all) else: - figs = None - self.figs = [] + # Flatten plot_panel_list, assigning each subplot the panel's plot_level, + # then prepend to plot_list + flat = [] + for panel in plot_panel_list: + panel_level = panel.get('plot_level', 0) + for plot in panel['plots']: + flat.append({**plot, 'plot_level': panel_level}) + all_plots = flat + list(plot_list) + + if not hasattr(self, '_flat_fig_map'): + self._flat_fig_map = {} + + self.figs = self._do_inspect(all_plots, self._flat_fig_map, + dataset=dataset, + replot_all=replot_all) + + plt.pause(0.05) + + + def _do_inspect(self, plot_list, fig_map, dataset=None, replot_all=False): + """Core one-figure-per-plot rendering logic. + + fig_map is a dict {title: figure} owned by the caller and updated + in-place. It tracks which figures are open across calls. + + Behaviour: + replot_all=False — closed figures are skipped (left closed). + replot_all=True — closed figures are recreated. + + Returns the list of figures that were rendered this call. + """ + if not plot_list: + return [] + + rendered = [] + + for plot in plot_list: + # Level filter + if plot.get('plot_level', 0) > self.plot_level: + continue + + # Condition check + condition = plot.get('condition', None) + if condition is not None: + try: + if not condition(self): + continue + except TypeError: + if not condition(self, dataset): + continue + + title = plot['title'] + fig = fig_map.get(title) + + if fig is not None and not plt.fignum_exists(fig.number): + # Figure was closed by the user + if replot_all: + fig = None + del fig_map[title] + else: + continue # leave it closed + + if fig is None: + fig = plt.figure(num=title) + fig._panel_label = title + fig_map[title] = fig - idx = 0 - for plots in self.plot_list: - # If a conditional is included in the plot, we check whether - # it is True try: - if len(plots) >=3 and not plots[2](self): - continue - except TypeError as e: - if len(plots) >= 3 and not plots[2](self, dataset): - continue - - name = plots[0] - plotter = plots[1] - - if figs is None: - fig = plt.figure() - self.figs.append(fig) - else: - fig = figs[idx] - - - try: # We try just plotting using the simplest allowed signature - plotter(self,fig) - plt.title(name) - except TypeError as e: - # TypeError implies it wanted another argument, i.e. a dataset + plot['plot_func'](self, fig) + plt.title(title) + except TypeError: if dataset is not None: try: - plotter(self, fig, dataset) - plt.title(name) - except Exception as e: # Don't raise errors: it's just plots + plot['plot_func'](self, fig, dataset) + plt.title(title) + except Exception: pass - - except Exception as e: # Don't raise errors, it's just a plot + except Exception: pass - idx += 1 + rendered.append(fig) + try: + fig.canvas.draw_idle() + except Exception: + pass - if update: - # This seems to update the figure without blocking. - plt.draw() - fig.canvas.start_event_loop(0.001) + return rendered + + + def _inspect_panel(self, dataset=None, replot_all=False): + """Multi-subplot panel rendering. + + Creates one figure per plot_panel_list entry, placing each subplot's + plot_func output into the appropriate axes. Closed panels stay closed + on subsequent calls unless replot_all=True. Standalone plot_list + entries are then rendered via _do_inspect and appended to self.figs. + """ + plot_panel_list = getattr(self, 'plot_panel_list', None) or [] + plot_list = getattr(self, 'plot_list', None) or [] + n_panels = len(plot_panel_list) + + # _panel_figs: list of figures (or None if never created / closed). + # _panel_axes: dict keyed by (panel_idx, row, col) → Axes. + # _standalone_fig_map: dict {title: figure} for standalone plot_list. + first_call = not hasattr(self, '_panel_figs') + if first_call: + self._panel_figs = [None] * n_panels + self._panel_axes = {} + self._standalone_fig_map = {} + + if not hasattr(self, '_standalone_fig_map'): + self._standalone_fig_map = {} + + for panel_idx, panel_def in enumerate(plot_panel_list): + panel_level = panel_def.get('plot_level', 0) + if panel_level > self.plot_level: + continue # skip entire panel + + nrows, ncols = panel_def['grid'] + figsize = panel_def.get('figure_size', None) + title = panel_def.get('title', '') + + fig = self._panel_figs[panel_idx] + + # Detect if a previously open figure was closed by the user. + if fig is not None and not plt.fignum_exists(fig.number): + self._panel_figs[panel_idx] = None + for k in [k for k in self._panel_axes if k[0] == panel_idx]: + del self._panel_axes[k] + fig = None + + if fig is None: + if not first_call and not replot_all: + continue # was closed; leave it closed + fig = plt.figure(num=title, figsize=figsize) + fig._panel_label = title + self._panel_figs[panel_idx] = fig + else: + # Remove all axes and recreate them fresh each update. + # plt.colorbar() shrinks the parent axes to make room for + # itself, so clearing and recreating is simpler than trying + # to undo that resizing. + for ax in list(fig.axes): + ax.remove() + for k in [k for k in self._panel_axes if k[0] == panel_idx]: + del self._panel_axes[k] + + for plot in panel_def['plots']: + condition = plot.get('condition', None) + if condition is not None: + try: + if not condition(self): + continue + except TypeError: + if not condition(self, dataset): + continue + + row, col = plot['subplot'] + position = row * ncols + col + 1 # 1-indexed for matplotlib + + ax_key = (panel_idx, row, col) + ax = fig.add_subplot(nrows, ncols, position) + self._panel_axes[ax_key] = ax + + try: + plot['plot_func'](self, ax) + ax.set_title(plot['title']) + except TypeError: + if dataset is not None: + try: + plot['plot_func'](self, ax, dataset) + ax.set_title(plot['title']) + except Exception: + pass + except Exception: + pass + + try: + fig.canvas.draw_idle() + except Exception: + pass + + # Rebuild self.figs from open panel figures + rendered standalone figures. + panel_figs = [f for f in self._panel_figs if f is not None] + standalone_rendered = self._do_inspect( + list(plot_list), self._standalone_fig_map, + dataset=dataset, replot_all=replot_all, + ) + self.figs = panel_figs + standalone_rendered - if first_update: - # But this is needed the first time the figures update, or - # they won't get drawn at all - plt.pause(0.05 * len(self.figs)) def save_figures(self, prefix='', extension='.pdf'): @@ -661,14 +805,15 @@ class CDIModel(t.nn.Module): Default is .eps, the file extension to save with. """ - if hasattr(self, 'figs') and self.figs: - figs = self.figs - else: - return # No figures to save + if not (hasattr(self, 'figs') and self.figs): + return # No figures to save for fig in self.figs: - fig.savefig(prefix + fig.axes[0].get_title() + extension, - bbox_inches = 'tight') + if hasattr(fig, '_panel_label') and fig._panel_label: + label = fig._panel_label + else: + label = fig.axes[0].get_title() if fig.axes else 'figure' + fig.savefig(prefix + label + extension, bbox_inches='tight') def compare(self, dataset, logarithmic=False): diff --git a/src/cdtools/models/simple_ptycho.py b/src/cdtools/models/simple_ptycho.py index 5435c66..bcc02d2 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/src/cdtools/models/simple_ptycho.py @@ -15,10 +15,15 @@ class SimplePtycho(CDIModel): probe_guess, obj_guess, min_translation = [0,0], + panel_plot_mode=False, + plot_level=0, ): # We initialize the superclass - super().__init__() + super().__init__( + panel_plot_mode=panel_plot_mode, + plot_level=plot_level, + ) # We register all the constants, like wavelength, as buffers. This # lets the model hook into some nice pytorch features, like using @@ -43,7 +48,8 @@ class SimplePtycho(CDIModel): @classmethod - def from_dataset(cls, dataset): + def from_dataset(cls, dataset,panel_plot_mode=False, + plot_level=0, ): # We get the key geometry information from the dataset wavelength = dataset.wavelength @@ -76,7 +82,9 @@ class SimplePtycho(CDIModel): probe_basis, probe, obj, - min_translation=min_translation + min_translation=min_translation, + panel_plot_mode=panel_plot_mode, + plot_level=plot_level, ) @@ -107,15 +115,59 @@ class SimplePtycho(CDIModel): # This lists all the plots to display on a call to model.inspect() - plot_list = [ - ('Probe Amplitude', - lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)), - ('Probe Phase', - lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)), - ('Object Amplitude', - lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), - ('Object Phase', - lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)) + plot_panel_list = [ + { + # Title for window + 'title' : 'Probe Results', + # (width, height) in inches + 'figure_size': (7, 7), + # (nrows, ncols) for subplot grid + 'grid': (2, 2), + # A setting to control how many plots are produced + 'plot_level': 1, + # The list of plots to include + 'plots' : [ + { + 'title': 'Probe Amplitude', + 'subplot' : (0, 0), + 'plot_func': lambda self, fig: + p.plot_amplitude(self.probe, fig, + basis=self.probe_basis) + },{ + 'title': 'Probe Phase', + 'subplot' : (0, 1), + 'plot_func': lambda self, fig: + p.plot_phase(self.probe, fig, + basis=self.probe_basis) + } + ] + }, + { + # Title for window + 'title' : 'Object Results', + # (width, height) in inches + 'figure_size': (7, 7), + # (nrows, ncols) for subplot grid + 'grid': (2, 2), + # A setting to control how many plots are produced + 'plot_level': 1, + # The list of plots to include + 'plots' : [ + { + 'title': 'Object Amplitude', + 'subplot' : (1, 0), + 'plot_func': lambda self, fig: + p.plot_amplitude(self.obj, fig, + basis=self.probe_basis) + }, { + 'title': 'Object Phase', + 'subplot' : (1, 1), + 'plot_func': lambda self, fig: + p.plot_phase(self.obj, fig, + basis=self.probe_basis) + }, + ] + }, ] def save_results(self, dataset): diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py index 1666814..4beab9f 100644 --- a/src/cdtools/reconstructors/base.py +++ b/src/cdtools/reconstructors/base.py @@ -16,6 +16,7 @@ from torch.utils import data as td import threading import queue import time +from matplotlib import pyplot as plt from typing import List, Union if TYPE_CHECKING: @@ -357,10 +358,18 @@ class Reconstructor: try: calc.start() while calc.is_alive(): - if hasattr(self.model, 'figs'): - self.model.figs[0].canvas.start_event_loop(0.01) + figs = getattr(self.model, 'figs', []) + open_fig = next( + (f for f in figs if plt.fignum_exists(f.number)), + None, + ) + if open_fig is not None: + try: + open_fig.canvas.start_event_loop(0.01) + except Exception: + time.sleep(0.01) else: - calc.join() + time.sleep(0.01) except KeyboardInterrupt as e: stop_event.set() diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py index 0f570ea..e820764 100644 --- a/src/cdtools/tools/plotting/plotting.py +++ b/src/cdtools/tools/plotting/plotting.py @@ -164,7 +164,12 @@ def plot_image( else: im = im.detach().cpu().numpy() - if fig is None: + # Support passing an Axes object instead of a Figure + ax_mode = isinstance(fig, plt.Axes) + if ax_mode: + ax = fig + fig = ax.get_figure() + elif fig is None: fig = plt.figure() ax = fig.add_subplot(111, **kwargs) @@ -173,8 +178,13 @@ def plot_image( # given def make_plot(idx): plt.figure(fig.number) - title = plt.gca().get_title() - fig.clear() + if ax_mode: + title = ax.get_title() + ax.cla() + plt.sca(ax) + else: + title = plt.gca().get_title() + fig.clear() # If im only has two dimensions, this reshape will add a leading @@ -184,9 +194,10 @@ def plot_image( s = im.shape reshaped_im = im.reshape(-1,s[-2],s[-1]) num_images = reshaped_im.shape[0] - fig.plot_idx = idx % num_images + plot_holder = ax if ax_mode else fig + plot_holder.plot_idx = idx % num_images - to_plot = plot_func(reshaped_im[fig.plot_idx]) + to_plot = plot_func(reshaped_im[plot_holder.plot_idx]) mpl_im = plt.imshow( to_plot, @@ -273,11 +284,13 @@ def plot_image( plt.title(title) if len(im.shape) >= 3: - plt.text(0.03, 0.03, str(fig.plot_idx), fontsize=14, transform=plt.gcf().transFigure) + text_transform = ax.transAxes if ax_mode else plt.gcf().transFigure + plt.text(0.03, 0.03, str(plot_holder.plot_idx), fontsize=14, transform=text_transform) return fig - if hasattr(fig, 'plot_idx'): - result = make_plot(fig.plot_idx) + plot_holder = ax if ax_mode else fig + if hasattr(plot_holder, 'plot_idx'): + result = make_plot(plot_holder.plot_idx) else: result = make_plot(0) @@ -285,15 +298,16 @@ def plot_image( def on_action(event): + plot_holder = ax if ax_mode else fig 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) + update(plot_holder.plot_idx - 1) elif event.key == 'down' or event.button == 'down': - update(fig.plot_idx + 1) + update(plot_holder.plot_idx + 1) plt.draw() if len(im.shape) >=3: @@ -557,7 +571,11 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, inver factor = get_units_factor(units) - if fig is None: + if isinstance(fig, plt.Axes): + ax = fig + fig = ax.get_figure() + plt.sca(ax) + elif fig is None: fig = plt.figure() ax = fig.add_subplot(111, **kwargs) else: @@ -614,7 +632,13 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr The figure object that was actually plotted to. """ - if fig is None: + ax_mode = isinstance(fig, plt.Axes) + if ax_mode: + ax = fig + fig = ax.get_figure() + ax.cla() + plt.sca(ax) + elif fig is None: fig = plt.figure() else: plt.figure(fig.number) @@ -622,7 +646,8 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr factor = get_units_factor(units) - bbox = fig.get_window_extent().transformed(fig.dpi_scale_trans.inverted()) + plot_area = ax if ax_mode else fig + bbox = plot_area.get_window_extent().transformed(fig.dpi_scale_trans.inverted()) if isinstance(translations, t.Tensor): trans = translations.detach().cpu().numpy() else: