diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 490b421..4d2c998 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -168,10 +168,12 @@ class CDIModel(t.nn.Module): plot_list = [] - def inspect(self, dataset=None): + def inspect(self, dataset=None, update=True): """Plots all the plots defined in the model's plot_list attribute - It will plot all the registered plots for this model in new figures. + 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 Optionally, a dataset can be passed, which then will plot any registered plots which need to incorporate some information from @@ -179,8 +181,29 @@ class CDIModel(t.nn.Module): Args: dataset (torch.Dataset): Optional, a dataset matched to the model type + update (bool) : Whether to update existing plots or plot new ones + + Returns: + list : A list of figure numbers noting where the plots were plotted """ - for plots in self.plot_list: + first_update = False + if update and hasattr(self, 'figs') and self.figs: + figs = self.figs + elif update: + figs = None + self.figs = [] + first_update = True + else: + figs = None + self.figs = [] + + for idx, plots in enumerate(self.plot_list): + if figs is None: + fig = plt.figure() + self.figs.append(fig) + else: + fig = figs[idx] + name = plots[0] plotter = plots[1] # If a conditional is included in the plot @@ -191,18 +214,24 @@ class CDIModel(t.nn.Module): if len(plots) >= 3 and not plots[2](self, dataset): continue try: - plotter(self) + plotter(self,fig) plt.title(name) except TypeError as e: if dataset is not None: try: - plotter(self, dataset) + plotter(self, fig, dataset) plt.title(name) except (IndexError, KeyError, AttributeError) as e: pass except (IndexError, KeyError, AttributeError) as e: pass - + if update: + plt.draw() + fig.canvas.start_event_loop(0.001) + + if first_update: + plt.pause(0.05 * len(self.figs)) + diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index e214f27..5030ed7 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -276,23 +276,23 @@ class FancyPtycho(CDIModel): # Needs to be updated to allow for plotting to an existing figure plot_list = [ ('Dominant Probe Amplitude', - lambda self: p.plot_amplitude(self.probe[0], basis=self.probe_basis)), + lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)), ('Dominant Probe Phase', - lambda self: p.plot_phase(self.probe[0], basis=self.probe_basis)), + lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)), ('Subdominant Probe Amplitude', - lambda self: p.plot_amplitude(self.probe[1], basis=self.probe_basis), + lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis), lambda self: len(self.probe) >=1), ('Subdominant Probe Phase', - lambda self: p.plot_phase(self.probe[1], basis=self.probe_basis), + lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), lambda self: len(self.probe) >=1), ('Object Amplitude', - lambda self: p.plot_amplitude(self.obj, basis=self.probe_basis)), + lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), ('Object Phase', - lambda self: p.plot_phase(self.obj, basis=self.probe_basis)), + lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)), ('Corrected Translations', - lambda self, dataset: p.plot_translations(self.corrected_translations(dataset))), + lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)), ('Background', - lambda self: plt.figure() and plt.imshow(self.background.detach().cpu().numpy()**2)) + lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)) ] diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 79c3783..a448e95 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -31,12 +31,15 @@ dataset.get_as(device='cuda') #model.translation_offsets.requires_grad = False for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)): + model.inspect(dataset) print(i,loss) for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100, lr=0.001)): + model.inspect(dataset) print(i,loss) for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100, lr=0.0001)): + model.inspect(dataset) print(i,loss)