From 2a2b6ee9602defd8b85975615d4bc2e01f4c918d Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 22 May 2019 13:00:52 -0400 Subject: [PATCH] Add tool to restrict object support, plot a nanomap, and include mask on SHARP-style initialization --- CDTools/models/fancy_ptycho.py | 29 +++++++++--- CDTools/tools/initializers.py | 8 ++-- CDTools/tools/plotting.py | 84 +++++++++++++++++++++++++++++----- CDTools/tools/propagators.py | 5 ++ 4 files changed, 104 insertions(+), 22 deletions(-) diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index f5ba210..9d83724 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/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): diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py index 111778c..5ff497e 100644 --- a/CDTools/tools/initializers.py +++ b/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]], diff --git a/CDTools/tools/plotting.py b/CDTools/tools/plotting.py index f6c2639..97d0930 100644 --- a/CDTools/tools/plotting.py +++ b/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() + + diff --git a/CDTools/tools/propagators.py b/CDTools/tools/propagators.py index 2e39cb6..23c919a 100644 --- a/CDTools/tools/propagators.py +++ b/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 + +