Loveplotting!

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

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

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)