Change the way plots are defined for inspect in preparation for adding liveplotting

CDTools/models/__init__.py

@@ -2,7 +2,7 @@ from __future__ import division, print_function, absolute_import
 
 import torch as t
 from torch.utils import data as torchdata
-
+from matplotlib import pyplot as plt
 
 #
 # This is unrelated, but it will then be important to be able to save and load
@@ -164,6 +164,47 @@ class CDIModel(t.nn.Module):
         return self.AD_optimize(iterations, data_loader, optimizer)
 
 
+    # By default, the plot_list is empty
+    plot_list = []
+    
+    
+    def inspect(self, dataset=None):
+        """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.
+        
+        Optionally, a dataset can be passed, which then will plot any
+        registered plots which need to incorporate some information from
+        the dataset (such as geometry or a comparison with measured data).
+        
+        Args:
+            dataset (torch.Dataset): Optional, a dataset matched to the model type
+        """
+        for plots in self.plot_list:
+            name = plots[0]
+            plotter = plots[1]
+            # If a conditional is included in the plot
+            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
+            try:
+                plotter(self)
+                plt.title(name)
+            except TypeError as e:
+                if dataset is not None:
+                    try:
+                        plotter(self, dataset)
+                        plt.title(name)
+                    except (IndexError, KeyError, AttributeError) as e:
+                        pass
+            except (IndexError, KeyError, AttributeError) as e:
+                pass
+        
+
+
 
 
 from CDTools.models.simple_ptycho import SimplePtycho

CDTools/models/fancy_ptycho.py

@@ -271,33 +271,31 @@ class FancyPtycho(CDIModel):
         translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
         t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
         return translations + t_offset
+
+
+    # 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)),
+        ('Dominant Probe Phase',
+         lambda self: p.plot_phase(self.probe[0], basis=self.probe_basis)),
+        ('Subdominant Probe Amplitude',
+         lambda self: p.plot_amplitude(self.probe[1], 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: len(self.probe) >=1),
+        ('Object Amplitude', 
+         lambda self: p.plot_amplitude(self.obj, basis=self.probe_basis)),
+        ('Object Phase',
+         lambda self: p.plot_phase(self.obj, basis=self.probe_basis)),
+        ('Corrected Translations',
+         lambda self, dataset: p.plot_translations(self.corrected_translations(dataset))),
+        ('Background',
+         lambda self: plt.figure() and plt.imshow(self.background.detach().cpu().numpy()**2))
+    ]
+
     
-
-    def inspect(self, dataset=None):
-        p.plot_amplitude(self.probe[0], basis=self.probe_basis)
-        plt.title('Dominant Probe Amplitude')
-        p.plot_phase(self.probe[0], basis=self.probe_basis)
-        plt.title('Dominant Probe Phase')
-        
-        if len(self.probe) >=2:
-            p.plot_amplitude(self.probe[1], basis=self.probe_basis)
-            plt.title('Subdominant Probe Amplitude')
-            p.plot_phase(self.probe[1], basis=self.probe_basis)
-            plt.title('Subdominant Probe Phase')
-
-        p.plot_amplitude(self.obj, basis=self.probe_basis)
-        plt.title('Object Amplitude')
-        p.plot_phase(self.obj, basis=self.probe_basis)
-        plt.title('Object Phase')
-
-        if dataset is not None:
-            p.plot_translations(self.corrected_translations(dataset))
-
-        plt.figure()
-        plt.imshow(self.background.detach().cpu().numpy()**2)
-        plt.title('Background')
-        
-
     def save_results(self, dataset):
         basis = self.probe_basis.detach().cpu().numpy()
         translations = self.corrected_translations(dataset).detach().cpu().numpy()

CDTools/models/simple_ptycho.py

@@ -138,11 +138,16 @@ class SimplePtycho(CDIModel):
         raise NotImplementedError()
 
 
-    def inspect(self):
-        p.plot_amplitude(self.probe, basis=self.probe_basis, title = 'Probe Amplitude')
-        p.plot_phase(self.probe, basis=self.probe_basis, title = 'Probe Phase')
-        p.plot_amplitude(self.obj, basis=self.probe_basis, title = 'Object Amplitude')
-        p.plot_phase(self.obj, basis=self.probe_basis, title = 'Object Phase')
+    plot_list = [
+        ('Probe Amplitude',
+         lambda self: p.plot_amplitude(self.probe, basis=self.probe_basis)),
+        ('Probe Phase',
+         lambda self: p.plot_phase(self.probe, basis=self.probe_basis)),
+        ('Object Amplitude',
+         lambda self: p.plot_amplitude(self.obj, basis=self.probe_basis)),
+        ('Object Phase',
+         lambda self: p.plot_phase(self.obj, basis=self.probe_basis))
+    ]
 
 
     def save_results(self):

tests/tools/test_image_processing.py

@@ -138,3 +138,5 @@ def test_convolve_1d():
 
     np_result = np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0)
     assert np.allclose(convolved,np_result)
+
+