diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 4d2c998..096a453 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -3,6 +3,9 @@ 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 +from matplotlib.widgets import Slider +from matplotlib import ticker +import numpy as np # # This is unrelated, but it will then be important to be able to save and load @@ -93,10 +96,6 @@ class CDIModel(t.nn.Module): def simulate_to_dataset(self, args_list): raise NotImplementedError() - - - def inspect(self): - raise NotImplementedError() def AD_optimize(self, iterations, data_loader, optimizer, scheduler=None): @@ -180,11 +179,9 @@ class CDIModel(t.nn.Module): the dataset (such as geometry or a comparison with measured data). Args: - dataset (torch.Dataset): Optional, a dataset matched to the model type + dataset (CDataset): 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 """ first_update = False if update and hasattr(self, 'figs') and self.figs: @@ -233,6 +230,102 @@ class CDIModel(t.nn.Module): plt.pause(0.05 * len(self.figs)) + def compare(self, dataset): + """Opens a tool for comparing simulated and measured diffraction patterns + + Args: + dataset (CDataset) : A dataset containing the simulated diffraction patterns to compare agains + """ + fig, axes = plt.subplots(1,3,figsize=(12,5.3)) + fig.tight_layout(rect=[0.02, 0.09, 0.98, 0.96]) + axslider = plt.axes([0.15,0.06,0.75,0.03]) + + + def update_colorbar(im): + # If the update brought the colorbar out of whack + # (say, from clicking back in the navbar) + # Holy fuck this was annoying. Sorry future for how + # crappy this solution is. + #if not np.allclose(im.colorbar.ax.get_xlim(), + # (np.min(im.get_array()), + # np.max(im.get_array()))): + if hasattr(im, 'norecurse') and im.norecurse: + im.norecurse=False + return + + im.norecurse=True + im.colorbar.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array())) + im.colorbar.ax.set_ylim(0,1) + im.colorbar.set_ticks(ticker.LinearLocator(numticks=5)) + im.colorbar.draw_all() + + + def update(idx): + idx = int(idx) % len(dataset) + fig.pattern_idx = idx + updating = True if len(axes[0].images) >= 1 else False + + inputs, output = dataset[idx] + sim_data = self.forward(*inputs).detach().cpu().numpy() + meas_data = output.detach().cpu().numpy() + + if not updating: + axes[0].set_title('Simulated') + axes[1].set_title('Measured') + axes[2].set_title('Difference') + + sim = axes[0].imshow(sim_data) + meas = axes[1].imshow(meas_data) + diff = axes[2].imshow(sim_data-meas_data) + + cb1 = plt.colorbar(sim, ax=axes[0], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1) + cb1.ax.tick_params(labelrotation=20) + cb1.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(sim)) + cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1) + cb2.ax.tick_params(labelrotation=20) + cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas)) + cb3 = plt.colorbar(diff, ax=axes[2], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1) + cb3.ax.tick_params(labelrotation=20) + cb3.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(diff)) + + else: + sim = axes[0].images[-1] + sim.set_data(sim_data) + update_colorbar(sim) + + meas = axes[1].images[-1] + meas.set_data(meas_data) + update_colorbar(meas) + + diff = axes[2].images[-1] + diff.set_data(sim_data-meas_data) + update_colorbar(diff) + + + # This is dumb but the slider doesn't work unless a reference to it is + # kept somewhere... + self.slider = Slider(axslider, 'Pattern #', 0, len(dataset)-1, valstep=1, valfmt="%d") + self.slider.on_changed(update) + + def on_action(event): + 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.pattern_idx - 1) + elif event.key == 'down' or event.button == 'down': + update(fig.pattern_idx + 1) + self.slider.set_val(fig.pattern_idx) + plt.draw() + + fig.canvas.mpl_connect('key_press_event',on_action) + fig.canvas.mpl_connect('scroll_event',on_action) + update(0) + + + pass diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index a448e95..8343bfc 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -47,5 +47,6 @@ for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100, lr=0.0 # pickle.dump(model.save_results(dataset),f) model.inspect(dataset) +model.compare(dataset) plt.show() exit()