diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 3bf347c..490b421 100644 --- a/CDTools/models/__init__.py +++ b/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 diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 0b398c4..e214f27 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/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() diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index 48d9388..4ba4f68 100644 --- a/CDTools/models/simple_ptycho.py +++ b/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): diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py index e93869b..f37581c 100644 --- a/tests/tools/test_image_processing.py +++ b/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) + +