Add tool to restrict object support, plot a nanomap, and include mask on SHARP-style initialization

CDTools/models/fancy_ptycho.py

@@ -17,7 +17,7 @@ class FancyPtycho(CDIModel):
                  probe_guess, obj_guess, min_translation = t.Tensor([0,0]),
                  background = None, translation_offsets=None, mask=None,
                  weights = None, translation_scale = 1, saturation=None,
-                 probe_support = None):
+                 probe_support = None, obj_support=None):
         
         super(FancyPtycho,self).__init__()
         self.wavelength = t.Tensor([wavelength])
@@ -76,10 +76,17 @@ class FancyPtycho(CDIModel):
             self.probe_support = probe_support
         else:
             self.probe_support = t.ones_like(self.probe[0])
+
+        if obj_support is not None:
+            self.obj_support = obj_support
+            self.obj.data = self.obj * obj_support
+        else:
+            self.obj_support = t.ones_like(self.obj)
         
         
     @classmethod
-    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None):
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1):
+        
         wavelength = dataset.wavelength
         det_basis = dataset.detector_geometry['basis']
         det_shape = dataset[0][1].shape
@@ -147,10 +154,19 @@ class FancyPtycho(CDIModel):
             probe_support[p_cent[0]-psr:p_cent[0]+psr,
                           p_cent[1]-psr:p_cent[1]+psr] = 1
         else:
-            probe_support = t.ones_like(probe[0].to(dtype=t.float32))
-        
+            probe_support = None;
 
-        return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, translation_offsets = translation_offsets, weights=weights, mask=mask, background=background, translation_scale=translation_scale, saturation=saturation, probe_support=probe_support)
+        if restrict_obj != -1:
+            ro = restrict_obj
+            os = np.array(obj_size)
+            ps = np.array(probe_shape)
+            obj_support = t.zeros_like(obj.to(dtype=t.float32))
+            obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
+                        ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
+        else:
+            obj_support = None
+
+        return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, translation_offsets = translation_offsets, weights=weights, mask=mask, background=background, translation_scale=translation_scale, saturation=saturation, probe_support=probe_support, obj_support=obj_support)
                    
     
     def interaction(self, index, translations):
@@ -168,7 +184,7 @@ class FancyPtycho(CDIModel):
             #                                                self.obj,
             #                                                pix_trans)
             exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr,
-                                                           self.obj,
+                                                           self.obj_support * self.obj,
                                                            pix_trans,
                                                            shift_probe=True)
             exit_waves = exit_waves * self.probe_support[...,:,:]
@@ -224,6 +240,7 @@ class FancyPtycho(CDIModel):
         self.probe_basis = self.probe_basis.to(*args,**kwargs)
         self.probe_norm = self.probe_norm.to(*args,**kwargs)
         self.probe_support = self.probe_support.to(*args,**kwargs)
+        self.obj_support = self.obj_support.to(*args,**kwargs)
         
 
     def sim_to_dataset(self, args_list):

CDTools/tools/initializers.py

@@ -236,10 +236,11 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None):
         propagatioin_distance (float) : Default is no propagation, an amount to propagate the guessed probe from it's focal point
     """
 
-    
+
+    # to use the mask or not?
     intensities = np.zeros(shape)
     for params, im in dataset:
-        intensities[det_slice] += im.cpu().numpy()
+        intensities[det_slice] += dataset.mask.cpu().numpy() * im.cpu().numpy()
     intensities /= len(dataset)
 
     # Subtract off a known background if it's stored
@@ -255,8 +256,7 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None):
     center = np.array(probe_guess.shape) // 2
 
     
-    # I had to remove this because it put some intensity outside of
-    # the detector region that caused issues
+    # I'm always divided on whether to use this modification:
     
     probe_guess[center[0], center[1]]=np.mean([
         probe_guess[center[0]-1, center[1]],

CDTools/tools/plotting.py

@@ -8,7 +8,8 @@ from matplotlib.colors import hsv_to_rgb
 
 
 __all__ = ['colorize','plot_1D','plot_amplitude','plot_phase',
-           'plot_colorized', 'plot_translations','get_units_factor']
+           'plot_colorized', 'plot_translations','get_units_factor',
+           'plot_nanomap']
 
 
 def colorize(z):
@@ -86,7 +87,7 @@ def plot_1D(arr, fig = None, **kwargs):
     plt.scatter(np.arange(arr.shape[-1]), arr)
 
     
-def plot_amplitude(im, fig = None, basis=None, units='um', **kwargs):
+def plot_amplitude(im, fig = None, basis=None, units='um', cmap='viridis', **kwargs):
     """ Plots the amplitude of a complex Tensor or numpy array with dimensions NxMx2.
     Args:
         im (t.Tensor) : An image with dimensions NxMx2.
@@ -94,6 +95,7 @@ def plot_amplitude(im, fig = None, basis=None, units='um', **kwargs):
         a new figure is created with an Axes subplot at 111.
         basis (numpy array) : Optional, the 3x2 probe basis, used to put the axis labels in real space units.
         units (str) : The units to convert the basis to
+        cmap (str) : Default is 'viridis', the colormap to plot with
         **kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class
         (see https://matplotlib.org/api/axes_api.html#the-axes-class)
     """
@@ -120,8 +122,10 @@ def plot_amplitude(im, fig = None, basis=None, units='um', **kwargs):
     else:
         extent=None
         
-    plt.imshow(absolute, cmap = 'viridis', extent = extent)
-    plt.colorbar()
+    plt.imshow(absolute, cmap = cmap, extent = extent)
+    cbar = plt.colorbar()
+    cbar.set_label('Amplitude (a.u.)')
+
     if basis is not None:
         plt.xlabel('X (' + units + ')')
         plt.ylabel('Y (' + units + ')')
@@ -132,13 +136,14 @@ def plot_amplitude(im, fig = None, basis=None, units='um', **kwargs):
     return fig
 
 
-def plot_phase(im, fig=None, basis=None, units='um', **kwargs):
+def plot_phase(im, fig=None, basis=None, units='um', cmap='auto', **kwargs):
     """ Plots the phase of a complex Tensor or numpy array with dimensions NxMx2.
     Args:
         im (t.Tensor) : An image with dimensions NxMx2.
         fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None,
         a new figure is created with an Axes subplot at 111.
         basis (numpy array) : Optional, the 3x2 probe basis, used to put the axis labels in real space units.
+        cmap (str) : Default is 'auto', which chooses between twilight and hsv based on availability.
         **kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class
         (see https://matplotlib.org/api/axes_api.html#the-axes-class)
     """
@@ -165,13 +170,17 @@ def plot_phase(im, fig=None, basis=None, units='um', **kwargs):
     else:
         extent=None
 
-    try:
-        plt.imshow(phase, cmap = 'twilight', extent=extent)
-    except:
-        plt.imshow(phase, cmap = 'hsv', extent=extent)
-
-    plt.colorbar()
-
+    if cmap == 'auto':
+        try:
+            plt.imshow(phase, cmap = 'twilight', extent=extent)
+        except:
+            plt.imshow(phase, cmap = 'hsv', extent=extent)
+    else:
+        plt.imshow(phase, cmap = cmap, extent=extent)
+        
+    cbar = plt.colorbar()
+    cbar.set_label('Phase (rad)')
+    
     if basis is not None:
         plt.xlabel('X (' + units + ')')
         plt.ylabel('Y (' + units + ')')
@@ -262,3 +271,54 @@ def plot_translations(translations, fig=None, units='um', lines=True):
     plt.xlabel('X (' + units + ')')
     plt.ylabel('Y (' + units + ')')
 
+
+    
+def plot_nanomap(translations, values, fig=None, units='um', convention='probe'):
+    """Plots a set of nanomap data in a flexible way
+    
+    Args:
+        translations : An Nx2 or Nx3 set of translations in real space
+        values : a length-N object of values associated with the translations
+        fig : Optional, a figure to plot into
+        units : Default is um, units to report in (assuming input in m)
+        lines : Whether to plot the lines indicating the path
+        convention : 'probe' if the translations refer to probe translations, 'obj' if they refer to object translations
+
+    Returns:
+        None
+    """
+
+    if fig is None:
+        fig = plt.figure()
+    else:
+        plt.figure(fig.number)
+        plt.gcf().clear()
+
+    factor = get_units_factor(units)
+
+    bbox = fig.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+    if isinstance(translations, t.Tensor):
+        trans = translations.detach().cpu().numpy()
+    else:
+        trans = np.array(translations)
+
+    if isinstance(values, t.Tensor):
+        values = values.detach().cpu().numpy()
+    else:
+        values = np.array(values)
+
+    if convention.lower() != 'probe':
+        trans = trans * -1
+        
+    s = bbox.width * bbox.height / trans.shape[0] * 72**2 #72 is points per inch
+    s /= 4 # A rough value to make the size work out
+    
+    plt.scatter(factor * trans[:,0],factor * trans[:,1],s=s,c=values)
+    
+    plt.gca().invert_xaxis()
+    plt.gca().set_facecolor('k')
+    plt.xlabel('Translation x (' + units + ')')
+    plt.ylabel('Translation y (' + units + ')')
+    plt.colorbar()
+
+

CDTools/tools/propagators.py

@@ -144,3 +144,8 @@ def inverse_near_field(wavefront, angular_spectrum_propagator):
 
 
 
+# I think it would be worthwhile to implement an FFT-DI based strategy as
+# well, especially for probe initialization where the propagation distance
+# can be large relative to what the angular spectrum method can reliably handle
+
+