Add a tool to look at raw data

CDTools/datasets.py

@@ -4,7 +4,11 @@ import torch as t
 from copy import copy
 
 from CDTools.tools import data as cdtdata
+from CDTools.tools import plotting
 from torch.utils import data as torchdata
+from matplotlib import pyplot as plt
+from matplotlib.widgets import Slider
+from matplotlib import ticker
 
 __all__ = ['CDataset', 'Ptycho_2D_Dataset']
 
@@ -178,3 +182,115 @@ class Ptycho_2D_Dataset(CDataset):
         super(Ptycho_2D_Dataset,self).to_cxi(cxi_file)
         cdtdata.add_data(cxi_file, self.patterns, axes=self.axes)
         cdtdata.add_ptycho_translations(cxi_file, self.translations)
+
+
+    def inspect(self):
+        fig, axes = plt.subplots(1,2,figsize=(8,5.3))
+        fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96])
+        axslider = plt.axes([0.15,0.06,0.75,0.03])
+
+        translations = self.translations.detach().cpu().numpy()
+        nanomap_values = self.patterns.sum(dim=(1,2)).detach().cpu().numpy()
+        
+        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 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(self)
+            fig.pattern_idx = idx
+            updating = True if len(axes[1].images) >= 1 else False
+            
+            inputs, output = self[idx]
+            meas_data = output.detach().cpu().numpy()
+            if hasattr(self, 'mask') and self.mask is not None:
+                mask = self.mask.detach().cpu().numpy()
+            else:
+                mask = 1
+                
+            if not updating:
+                print('hi')
+                axes[0].set_title('Nanomap')
+                axes[1].set_title('Pattern')
+
+                bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+                
+                s = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
+                s /= 4 # A rough value to make the size work out
+                s = np.ones(len(nanomap_values)) * s
+
+                s[idx] *= 4
+    
+                nanomap = axes[0].scatter(1e6 * translations[:,0],1e6 * translations[:,1],s=s,c=nanomap_values)
+    
+                axes[0].invert_xaxis()
+                axes[0].set_facecolor('k')
+                axes[0].set_xlabel('Translation x (um)')
+                axes[0].set_ylabel('Translation y (um)')
+                cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.15,fraction=0.1)
+                cb1.ax.tick_params(labelrotation=20)
+    
+                meas = axes[1].imshow(meas_data * mask)
+
+                cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.15,fraction=0.1)
+                cb2.ax.tick_params(labelrotation=20)
+                cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
+                
+            else:
+                axes[0].set_title('Nanomap')
+                bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+                
+                s = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
+                s /= 4 # A rough value to make the size work out
+                s = np.ones(len(nanomap_values)) * s
+                s[idx] *= 4
+    
+                axes[0].clear()
+                nanomap = axes[0].scatter(1e6 * translations[:,0],1e6 * translations[:,1],s=s,c=nanomap_values)
+                axes[0].invert_xaxis()
+                axes[0].set_facecolor('k')
+                axes[0].set_xlabel('Translation x (um)')
+                axes[0].set_ylabel('Translation y (um)')
+
+
+                
+                meas = axes[1].images[-1]
+                meas.set_data(meas_data * mask)
+                update_colorbar(meas)
+
+                
+        # This is dumb but the slider doesn't work unless a reference to it is
+        # kept somewhere...
+        self.slider = Slider(axslider, 'Pattern #', 0, len(self)-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)
+        
+