From 53ac42f13a38e88fa67056490154199e5dfbb19d Mon Sep 17 00:00:00 2001 From: Stephanie Ribet Date: Mon, 13 Nov 2023 14:19:21 -0800 Subject: [PATCH 1/8] small plotting bug --- .../phase/iterative_mixedstate_multislice_ptychography.py | 4 +++- py4DSTEM/process/phase/iterative_multislice_ptychography.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py index f4c10cb13..fec5d2d93 100644 --- a/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py @@ -3313,7 +3313,9 @@ def show_depth( gaussian_filter_sigma /= self.sampling[0] rotated_object = gaussian_filter(rotated_object, gaussian_filter_sigma) - plot_im = rotated_object[:, 0, int(y1_0) : int(y2_0)] + plot_im = rotated_object[ + :, 0, np.max((0, int(y1_0))) : np.min((int(y2_0), rotated_object.shape[2])) + ] extent = [ 0, diff --git a/py4DSTEM/process/phase/iterative_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_multislice_ptychography.py index 93e32b079..8b3d1b531 100644 --- a/py4DSTEM/process/phase/iterative_multislice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_multislice_ptychography.py @@ -3145,7 +3145,9 @@ def show_depth( gaussian_filter_sigma /= self.sampling[0] rotated_object = gaussian_filter(rotated_object, gaussian_filter_sigma) - plot_im = rotated_object[:, 0, int(y1_0) : int(y2_0)] + plot_im = rotated_object[ + :, 0, np.max((0, int(y1_0))) : np.min((int(y2_0), rotated_object.shape[2])) + ] extent = [ 0, From 6a9f1cae8ee982455635d3894ceda04c61e6b7b3 Mon Sep 17 00:00:00 2001 From: Stephanie Ribet Date: Mon, 13 Nov 2023 17:26:57 -0800 Subject: [PATCH 2/8] parallax plotting bug --- py4DSTEM/process/phase/iterative_parallax.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/py4DSTEM/process/phase/iterative_parallax.py b/py4DSTEM/process/phase/iterative_parallax.py index 716e1d782..04366b8c2 100644 --- a/py4DSTEM/process/phase/iterative_parallax.py +++ b/py4DSTEM/process/phase/iterative_parallax.py @@ -603,7 +603,9 @@ def preprocess( 0.5 * (self._reciprocal_sampling[0] * self._dp_mask.shape[0]), -0.5 * (self._reciprocal_sampling[0] * self._dp_mask.shape[0]), ] - ax[1].imshow(self._dp_mask, extent=reciprocal_extent, cmap="gray") + ax[1].imshow( + self._asnumpy(self._dp_mask), extent=reciprocal_extent, cmap="gray" + ) ax[1].set_title("DP mask") ax[1].set_ylabel(r"$k_x$ [$A^{-1}$]") ax[1].set_xlabel(r"$k_y$ [$A^{-1}$]") From 45890580b79da94cd1d5e8777f493cab5672dcc5 Mon Sep 17 00:00:00 2001 From: Georgios Varnavides Date: Sat, 18 Nov 2023 17:42:11 -0800 Subject: [PATCH 3/8] fixing various positions_mask bugs --- .../process/phase/iterative_base_class.py | 14 ++++-- ...tive_mixedstate_multislice_ptychography.py | 14 +++--- .../iterative_mixedstate_ptychography.py | 13 ++--- .../iterative_multislice_ptychography.py | 13 ++--- .../iterative_overlap_magnetic_tomography.py | 47 +++++++++++++----- .../phase/iterative_overlap_tomography.py | 48 ++++++++++++++----- .../iterative_simultaneous_ptychography.py | 37 +++++++++----- .../iterative_singleslice_ptychography.py | 14 +++--- 8 files changed, 134 insertions(+), 66 deletions(-) diff --git a/py4DSTEM/process/phase/iterative_base_class.py b/py4DSTEM/process/phase/iterative_base_class.py index 767789df2..fdb0092b3 100644 --- a/py4DSTEM/process/phase/iterative_base_class.py +++ b/py4DSTEM/process/phase/iterative_base_class.py @@ -1168,7 +1168,7 @@ def _normalize_diffraction_intensities( diffraction_intensities = self._asnumpy(diffraction_intensities) if positions_mask is not None: - number_of_patterns = np.count_nonzero(self._positions_mask.ravel()) + number_of_patterns = np.count_nonzero(positions_mask.ravel()) else: number_of_patterns = np.prod(diffraction_intensities.shape[:2]) @@ -1217,7 +1217,7 @@ def _normalize_diffraction_intensities( for rx in range(diffraction_intensities.shape[0]): for ry in range(diffraction_intensities.shape[1]): if positions_mask is not None: - if not self._positions_mask[rx, ry]: + if not positions_mask[rx, ry]: continue intensities = get_shifted_ar( diffraction_intensities[rx, ry], @@ -1348,6 +1348,14 @@ def to_h5(self, group): data=metadata, ) + # saving multiple None positions_mask fix + if self._positions_mask is None: + positions_mask = None + elif self._positions_mask[0] is None: + positions_mask = None + else: + positions_mask = self._positions_mask + # preprocessing metadata self.metadata = Metadata( name="preprocess_metadata", @@ -1359,7 +1367,7 @@ def to_h5(self, group): "num_diffraction_patterns": self._num_diffraction_patterns, "sampling": self.sampling, "angular_sampling": self.angular_sampling, - "positions_mask": self._positions_mask, + "positions_mask": positions_mask, }, ) diff --git a/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py index fec5d2d93..6acf20620 100644 --- a/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py @@ -217,13 +217,6 @@ def __init__( f"object_type must be either 'potential' or 'complex', not {object_type}" ) - if positions_mask is not None and positions_mask.dtype != "bool": - warnings.warn( - ("`positions_mask` converted to `bool` array"), - UserWarning, - ) - positions_mask = np.asarray(positions_mask, dtype="bool") - self.set_save_defaults() # Data @@ -440,6 +433,13 @@ def preprocess( ) ) + if self._positions_mask is not None and self._positions_mask.dtype != "bool": + warnings.warn( + ("`positions_mask` converted to `bool` array"), + UserWarning, + ) + self._positions_mask = np.asarray(self._positions_mask, dtype="bool") + ( self._datacube, self._vacuum_probe_intensity, diff --git a/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py b/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py index d68291143..15a7438e7 100644 --- a/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py +++ b/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py @@ -164,12 +164,6 @@ def __init__( raise ValueError( f"object_type must be either 'potential' or 'complex', not {object_type}" ) - if positions_mask is not None and positions_mask.dtype != "bool": - warnings.warn( - ("`positions_mask` converted to `bool` array"), - UserWarning, - ) - positions_mask = np.asarray(positions_mask, dtype="bool") self.set_save_defaults() @@ -297,6 +291,13 @@ def preprocess( ) ) + if self._positions_mask is not None and self._positions_mask.dtype != "bool": + warnings.warn( + ("`positions_mask` converted to `bool` array"), + UserWarning, + ) + self._positions_mask = np.asarray(self._positions_mask, dtype="bool") + ( self._datacube, self._vacuum_probe_intensity, diff --git a/py4DSTEM/process/phase/iterative_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_multislice_ptychography.py index 8b3d1b531..f4e2bec79 100644 --- a/py4DSTEM/process/phase/iterative_multislice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_multislice_ptychography.py @@ -198,12 +198,6 @@ def __init__( raise ValueError( f"object_type must be either 'potential' or 'complex', not {object_type}" ) - if positions_mask is not None and positions_mask.dtype != "bool": - warnings.warn( - ("`positions_mask` converted to `bool` array"), - UserWarning, - ) - positions_mask = np.asarray(positions_mask, dtype="bool") self.set_save_defaults() @@ -420,6 +414,13 @@ def preprocess( ) ) + if self._positions_mask is not None and self._positions_mask.dtype != "bool": + warnings.warn( + ("`positions_mask` converted to `bool` array"), + UserWarning, + ) + self._positions_mask = np.asarray(self._positions_mask, dtype="bool") + ( self._datacube, self._vacuum_probe_intensity, diff --git a/py4DSTEM/process/phase/iterative_overlap_magnetic_tomography.py b/py4DSTEM/process/phase/iterative_overlap_magnetic_tomography.py index c49a1faac..bdd04b4b4 100644 --- a/py4DSTEM/process/phase/iterative_overlap_magnetic_tomography.py +++ b/py4DSTEM/process/phase/iterative_overlap_magnetic_tomography.py @@ -166,13 +166,6 @@ def __init__( if object_type != "potential": raise NotImplementedError() - if positions_mask is not None and positions_mask.dtype != "bool": - warnings.warn( - ("`positions_mask` converted to `bool` array"), - UserWarning, - ) - positions_mask = np.asarray(positions_mask, dtype="bool") - self.set_save_defaults() # Data @@ -512,15 +505,43 @@ def preprocess( ) ) + if self._positions_mask is not None: + self._positions_mask = np.asarray(self._positions_mask) + + if self._positions_mask.ndim == 2: + warnings.warn( + "2D `positions_mask` assumed the same for all measurements.", + UserWarning, + ) + self._positions_mask = np.tile( + self._positions_mask, (self._num_tilts, 1, 1) + ) + + if self._positions_mask.dtype != "bool": + warnings.warn( + ("`positions_mask` converted to `bool` array."), + UserWarning, + ) + self._positions_mask = self._positions_mask.astype("bool") + else: + self._positions_mask = [None] * self._num_tilts + # Prepopulate various arrays - num_probes_per_tilt = [0] - for dc in self._datacube: - rx, ry = dc.Rshape - num_probes_per_tilt.append(rx * ry) + if self._positions_mask[0] is None: + num_probes_per_tilt = [0] + for dc in self._datacube: + rx, ry = dc.Rshape + num_probes_per_tilt.append(rx * ry) + + num_probes_per_tilt = np.array(num_probes_per_tilt) + else: + num_probes_per_tilt = np.insert( + self._positions_mask.sum(axis=(-2, -1)), 0, 0 + ) - self._num_diffraction_patterns = sum(num_probes_per_tilt) - self._cum_probes_per_tilt = np.cumsum(np.array(num_probes_per_tilt)) + self._num_diffraction_patterns = num_probes_per_tilt.sum() + self._cum_probes_per_tilt = np.cumsum(num_probes_per_tilt) self._mean_diffraction_intensity = [] self._positions_px_all = np.empty((self._num_diffraction_patterns, 2)) diff --git a/py4DSTEM/process/phase/iterative_overlap_tomography.py b/py4DSTEM/process/phase/iterative_overlap_tomography.py index ddd13ac58..685f95257 100644 --- a/py4DSTEM/process/phase/iterative_overlap_tomography.py +++ b/py4DSTEM/process/phase/iterative_overlap_tomography.py @@ -175,13 +175,6 @@ def __init__( if object_type != "potential": raise NotImplementedError() - if positions_mask is not None and positions_mask.dtype != "bool": - warnings.warn( - ("`positions_mask` converted to `bool` array"), - UserWarning, - ) - positions_mask = np.asarray(positions_mask, dtype="bool") - self.set_save_defaults() # Data @@ -453,14 +446,43 @@ def preprocess( ) ) + if self._positions_mask is not None: + self._positions_mask = np.asarray(self._positions_mask) + + if self._positions_mask.ndim == 2: + warnings.warn( + "2D `positions_mask` assumed the same for all measurements.", + UserWarning, + ) + self._positions_mask = np.tile( + self._positions_mask, (self._num_tilts, 1, 1) + ) + + if self._positions_mask.dtype != "bool": + warnings.warn( + ("`positions_mask` converted to `bool` array."), + UserWarning, + ) + self._positions_mask = self._positions_mask.astype("bool") + else: + self._positions_mask = [None] * self._num_tilts + # Prepopulate various arrays - num_probes_per_tilt = [0] - for dc in self._datacube: - rx, ry = dc.Rshape - num_probes_per_tilt.append(rx * ry) - self._num_diffraction_patterns = sum(num_probes_per_tilt) - self._cum_probes_per_tilt = np.cumsum(np.array(num_probes_per_tilt)) + if self._positions_mask[0] is None: + num_probes_per_tilt = [0] + for dc in self._datacube: + rx, ry = dc.Rshape + num_probes_per_tilt.append(rx * ry) + + num_probes_per_tilt = np.array(num_probes_per_tilt) + else: + num_probes_per_tilt = np.insert( + self._positions_mask.sum(axis=(-2, -1)), 0, 0 + ) + + self._num_diffraction_patterns = num_probes_per_tilt.sum() + self._cum_probes_per_tilt = np.cumsum(num_probes_per_tilt) self._mean_diffraction_intensity = [] self._positions_px_all = np.empty((self._num_diffraction_patterns, 2)) diff --git a/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py b/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py index 233d34e45..637789108 100644 --- a/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py +++ b/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py @@ -153,12 +153,6 @@ def __init__( raise ValueError( f"object_type must be either 'potential' or 'complex', not {object_type}" ) - if positions_mask is not None and positions_mask.dtype != "bool": - warnings.warn( - ("`positions_mask` converted to `bool` array"), - UserWarning, - ) - positions_mask = np.asarray(positions_mask, dtype="bool") self.set_save_defaults() @@ -339,6 +333,27 @@ def preprocess( f"simultaneous_measurements_mode must be either '-+', '-0+', or '0+', not {self._simultaneous_measurements_mode}" ) + if self._positions_mask is not None: + self._positions_mask = np.asarray(self._positions_mask) + + if self._positions_mask.ndim == 2: + warnings.warn( + "2D `positions_mask` assumed the same for all measurements.", + UserWarning, + ) + self._positions_mask = np.tile( + self._positions_mask, (self._num_sim_measurements, 1, 1) + ) + + if self._positions_mask.dtype != "bool": + warnings.warn( + "`positions_mask` converted to `bool` array.", + UserWarning, + ) + self._positions_mask = self._positions_mask.astype("bool") + else: + self._positions_mask = [None] * self._num_sim_measurements + if force_com_shifts is None: force_com_shifts = [None, None, None] elif len(force_com_shifts) == self._num_sim_measurements: @@ -421,7 +436,7 @@ def preprocess( com_fitted_x_0, com_fitted_y_0, crop_patterns, - self._positions_mask, + self._positions_mask[0], ) # explicitly delete namescapes @@ -506,7 +521,7 @@ def preprocess( com_fitted_x_1, com_fitted_y_1, crop_patterns, - self._positions_mask, + self._positions_mask[1], ) # explicitly delete namescapes @@ -592,7 +607,7 @@ def preprocess( com_fitted_x_2, com_fitted_y_2, crop_patterns, - self._positions_mask, + self._positions_mask[2], ) # explicitly delete namescapes @@ -632,8 +647,8 @@ def preprocess( self._region_of_interest_shape = np.array(self._amplitudes[0].shape[-2:]) self._positions_px = self._calculate_scan_positions_in_pixels( - self._scan_positions, self._positions_mask - ) + self._scan_positions, self._positions_mask[0] + ) # TO-DO: generaltize to per-dataset probe positions # handle semiangle specified in pixels if self._semiangle_cutoff_pixels: diff --git a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py index 350d0a3cb..c89f7ca6f 100644 --- a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py @@ -150,13 +150,6 @@ def __init__( f"object_type must be either 'potential' or 'complex', not {object_type}" ) - if positions_mask is not None and positions_mask.dtype != "bool": - warnings.warn( - ("`positions_mask` converted to `bool` array"), - UserWarning, - ) - positions_mask = np.asarray(positions_mask, dtype="bool") - self.set_save_defaults() # Data @@ -282,6 +275,13 @@ def preprocess( ) ) + if self._positions_mask is not None and self._positions_mask.dtype != "bool": + warnings.warn( + ("`positions_mask` converted to `bool` array"), + UserWarning, + ) + self._positions_mask = np.asarray(self._positions_mask, dtype="bool") + ( self._datacube, self._vacuum_probe_intensity, From 49f946e4f49850c544a09b8133f8b0f5063f26a7 Mon Sep 17 00:00:00 2001 From: Georgios Varnavides Date: Sat, 18 Nov 2023 18:28:37 -0800 Subject: [PATCH 4/8] added force_rotation_deg to parallax --- py4DSTEM/process/phase/iterative_parallax.py | 43 +++++++++++++++----- 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/py4DSTEM/process/phase/iterative_parallax.py b/py4DSTEM/process/phase/iterative_parallax.py index 04366b8c2..c9dbb2fcf 100644 --- a/py4DSTEM/process/phase/iterative_parallax.py +++ b/py4DSTEM/process/phase/iterative_parallax.py @@ -587,7 +587,7 @@ def preprocess( self.recon_BF = asnumpy(self._recon_BF) if plot_average_bf: - figsize = kwargs.pop("figsize", (6, 12)) + figsize = kwargs.pop("figsize", (8, 4)) fig, ax = plt.subplots(1, 2, figsize=figsize) @@ -1335,6 +1335,7 @@ def aberration_fit( plot_BF_shifts_comparison: bool = None, upsampled: bool = True, force_transpose: bool = False, + force_rotation_deg: float = None, ): """ Fit aberrations to the measured image shifts. @@ -1365,7 +1366,9 @@ def aberration_fit( upsampled: bool If True, and upsampled BF is available, uses that for CTF FFT fitting. force_transpose: bool - If True, and fit_BF_shifts is True, flips the measured x and y shifts + If True, flips the measured x and y shifts. + force_rotation_deg: float + If not None, sets the rotation angle to value in degrees. """ xp = self._xp @@ -1381,23 +1384,41 @@ def aberration_fit( ) else: self._xy_shifts_Ang = self._xy_shifts * xp.array(self._scan_sampling) + self.transpose = force_transpose # Solve affine transformation m = asnumpy( xp.linalg.lstsq(self._probe_angles, self._xy_shifts_Ang, rcond=None)[0] ) - m_rotation, m_aberration = polar(m, side="right") - # Convert into rotation and aberration coefficients - self.rotation_Q_to_R_rads = -1 * np.arctan2(m_rotation[1, 0], m_rotation[0, 0]) - if np.abs(np.mod(self.rotation_Q_to_R_rads + np.pi, 2.0 * np.pi) - np.pi) > ( - np.pi * 0.5 - ): - self.rotation_Q_to_R_rads = ( - np.mod(self.rotation_Q_to_R_rads, 2.0 * np.pi) - np.pi + if force_rotation_deg is None: + m_rotation, m_aberration = polar(m, side="right") + + if force_transpose: + m_rotation = m_rotation.T + + # Convert into rotation and aberration coefficients + + self.rotation_Q_to_R_rads = -1 * np.arctan2( + m_rotation[1, 0], m_rotation[0, 0] ) - m_aberration = -1.0 * m_aberration + if np.abs( + np.mod(self.rotation_Q_to_R_rads + np.pi, 2.0 * np.pi) - np.pi + ) > (np.pi * 0.5): + self.rotation_Q_to_R_rads = ( + np.mod(self.rotation_Q_to_R_rads, 2.0 * np.pi) - np.pi + ) + m_aberration = -1.0 * m_aberration + else: + self.rotation_Q_to_R_rads = np.deg2rad(force_rotation_deg) + c, s = np.cos(self.rotation_Q_to_R_rads), np.sin(self.rotation_Q_to_R_rads) + + m_rotation = np.array([[c, -s], [s, c]]) + if force_transpose: + m_rotation = m_rotation.T + + m_aberration = m_rotation @ m self.aberration_C1 = (m_aberration[0, 0] + m_aberration[1, 1]) / 2.0 From 38234b1451fcdcf6ebc487df3c3525e7f9a07d7c Mon Sep 17 00:00:00 2001 From: Stephanie Ribet Date: Mon, 20 Nov 2023 14:06:42 -0800 Subject: [PATCH 5/8] single slice fourier plotting with initial probe subtracted --- .../iterative_singleslice_ptychography.py | 63 ++++++++++++++----- 1 file changed, 46 insertions(+), 17 deletions(-) diff --git a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py index c89f7ca6f..fb1f34c38 100644 --- a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py @@ -472,7 +472,13 @@ def preprocess( self._probe_initial = self._probe.copy() self._probe_initial_aperture = xp.abs(xp.fft.fft2(self._probe)) - + self._initial_aberrations_array = ComplexProbe( + energy=self._energy, + gpts=self._region_of_interest_shape, + sampling=self.sampling, + parameters=self._polar_parameters, + device=self._device, + )._evaluate_ctf() self._known_aberrations_array = ComplexProbe( energy=self._energy, gpts=self._region_of_interest_shape, @@ -1750,6 +1756,7 @@ def _visualize_last_iteration( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, padding: int, **kwargs, ): @@ -1768,16 +1775,18 @@ def _visualize_last_iteration( If true, the reconstructed complex probe is displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object """ + xp = self._xp + asnumpy = self._asnumpy figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) if self._object_type == "complex": obj = np.angle(self.object) @@ -1869,10 +1878,16 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: + probe_array = self.probe_fourier + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift(xp.conjugate(self._initial_aberrations_array)) + ) probe_array = Complex2RGB( - self.probe_fourier, + probe_array, chroma_boost=chroma_boost, ) + ax.set_title("Reconstructed Fourier probe") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -1940,6 +1955,7 @@ def _visualize_all_iterations( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, iterations_grid: Tuple[int, int], padding: int, **kwargs, @@ -1961,10 +1977,14 @@ def _visualize_all_iterations( If true, the reconstructed complex probe is displayed plot_fourier_probe: bool If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object """ asnumpy = self._asnumpy + xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError( @@ -2003,10 +2023,7 @@ def _visualize_all_iterations( figsize = kwargs.pop("figsize", auto_figsize) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) errors = np.array(self.error_iterations) @@ -2104,14 +2121,20 @@ def _visualize_all_iterations( for n, ax in enumerate(grid): if plot_fourier_probe: - probe_array = Complex2RGB( - asnumpy( - self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]] - ) - ), - chroma_boost=chroma_boost, + probe_array = asnumpy( + self._return_fourier_probe_from_centered_probe( + probes[grid_range[n]] + ) ) + + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift( + xp.conjugate(self._initial_aberrations_array) + ) + ) + + probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) ax.set_title(f"Iter: {grid_range[n]} Fourier probe") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -2158,6 +2181,7 @@ def visualize( plot_convergence: bool = True, plot_probe: bool = True, plot_fourier_probe: bool = False, + remove_initial_probe_aberrations: bool = False, cbar: bool = True, padding: int = 0, **kwargs, @@ -2179,6 +2203,9 @@ def visualize( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object @@ -2195,6 +2222,7 @@ def visualize( plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, cbar=cbar, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, padding=padding, **kwargs, ) @@ -2205,6 +2233,7 @@ def visualize( iterations_grid=iterations_grid, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, padding=padding, **kwargs, From 6a2ab44bb5e4c47cacc3291bfde20d4d1a227732 Mon Sep 17 00:00:00 2001 From: Stephanie Ribet Date: Mon, 20 Nov 2023 15:17:57 -0800 Subject: [PATCH 6/8] everyone gets new fourier probe plotting --- ...tive_mixedstate_multislice_ptychography.py | 57 +++++++++++++------ .../iterative_mixedstate_ptychography.py | 56 ++++++++++++------ .../iterative_multislice_ptychography.py | 56 ++++++++++++------ .../phase/iterative_overlap_tomography.py | 56 ++++++++++++------ .../iterative_simultaneous_ptychography.py | 29 ++++++++-- .../iterative_singleslice_ptychography.py | 12 +--- 6 files changed, 188 insertions(+), 78 deletions(-) diff --git a/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py index 6acf20620..e873d9199 100644 --- a/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py @@ -2544,6 +2544,7 @@ def _visualize_last_iteration( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, padding: int, **kwargs, ): @@ -2562,17 +2563,19 @@ def _visualize_last_iteration( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe padding : int, optional Pixels to pad by post rotating-cropping object """ + xp = self._xp + asnumpy = self._asnumpy + figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) if self._object_type == "complex": obj = np.angle(self.object) @@ -2667,9 +2670,16 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: + probe_array = self.probe_fourier[0] + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) + ) probe_array = Complex2RGB( - self.probe_fourier[0], chroma_boost=chroma_boost + probe_array, + chroma_boost=chroma_boost, ) + ax.set_title("Reconstructed Fourier probe[0]") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -2735,6 +2745,7 @@ def _visualize_all_iterations( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, iterations_grid: Tuple[int, int], padding: int, **kwargs, @@ -2756,10 +2767,14 @@ def _visualize_all_iterations( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object """ asnumpy = self._asnumpy + xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError( @@ -2798,10 +2813,7 @@ def _visualize_all_iterations( figsize = kwargs.pop("figsize", auto_figsize) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) errors = np.array(self.error_iterations) @@ -2900,14 +2912,21 @@ def _visualize_all_iterations( for n, ax in enumerate(grid): if plot_fourier_probe: - probe_array = Complex2RGB( - asnumpy( - self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]][0] - ) - ), - chroma_boost=chroma_boost, + probe_array = asnumpy( + self._return_fourier_probe_from_centered_probe( + probes[grid_range[n]][0] + ) ) + + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift( + xp.conjugate(self._known_aberrations_array) + ) + ) + + probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) + ax.set_title(f"Iter: {grid_range[n]} Fourier probe[0]") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -2953,6 +2972,7 @@ def visualize( plot_convergence: bool = True, plot_probe: bool = True, plot_fourier_probe: bool = False, + remove_initial_probe_aberrations: bool = False, cbar: bool = True, padding: int = 0, **kwargs, @@ -2974,6 +2994,9 @@ def visualize( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object @@ -2989,6 +3012,7 @@ def visualize( plot_convergence=plot_convergence, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, padding=padding, **kwargs, @@ -3000,6 +3024,7 @@ def visualize( iterations_grid=iterations_grid, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, padding=padding, **kwargs, diff --git a/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py b/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py index 15a7438e7..426216116 100644 --- a/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py +++ b/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py @@ -1842,6 +1842,7 @@ def _visualize_last_iteration( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, padding: int, **kwargs, ): @@ -1858,16 +1859,18 @@ def _visualize_last_iteration( If true, the reconstructed complex probe is displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe padding : int, optional Pixels to pad by post rotating-cropping object """ + xp = self._xp + asnumpy = self._asnumpy + figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) if self._object_type == "complex": obj = np.angle(self.object) @@ -1959,10 +1962,16 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: + probe_array = self.probe_fourier[0] + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) + ) probe_array = Complex2RGB( - self.probe_fourier[0], + probe_array, chroma_boost=chroma_boost, ) + ax.set_title("Reconstructed Fourier probe[0]") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -2030,6 +2039,7 @@ def _visualize_all_iterations( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, iterations_grid: Tuple[int, int], padding: int, **kwargs, @@ -2051,10 +2061,14 @@ def _visualize_all_iterations( If true, the reconstructed complex probe is displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object """ asnumpy = self._asnumpy + xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError( @@ -2093,10 +2107,7 @@ def _visualize_all_iterations( figsize = kwargs.pop("figsize", auto_figsize) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) errors = np.array(self.error_iterations) @@ -2193,14 +2204,21 @@ def _visualize_all_iterations( for n, ax in enumerate(grid): if plot_fourier_probe: - probe_array = Complex2RGB( - asnumpy( - self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]][0] - ) - ), - chroma_boost=chroma_boost, + probe_array = asnumpy( + self._return_fourier_probe_from_centered_probe( + probes[grid_range[n]][0] + ) ) + + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift( + xp.conjugate(self._known_aberrations_array) + ) + ) + + probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) + ax.set_title(f"Iter: {grid_range[n]} Fourier probe[0]") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -2246,6 +2264,7 @@ def visualize( plot_convergence: bool = True, plot_probe: bool = True, plot_fourier_probe: bool = False, + remove_initial_probe_aberrations: bool = False, cbar: bool = True, padding: int = 0, **kwargs, @@ -2267,6 +2286,9 @@ def visualize( If true, the reconstructed complex probe is displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object @@ -2282,6 +2304,7 @@ def visualize( plot_convergence=plot_convergence, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, padding=padding, **kwargs, @@ -2293,6 +2316,7 @@ def visualize( iterations_grid=iterations_grid, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, padding=padding, **kwargs, diff --git a/py4DSTEM/process/phase/iterative_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_multislice_ptychography.py index f4e2bec79..f2304f72b 100644 --- a/py4DSTEM/process/phase/iterative_multislice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_multislice_ptychography.py @@ -2421,6 +2421,7 @@ def _visualize_last_iteration( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, padding: int, **kwargs, ): @@ -2439,17 +2440,19 @@ def _visualize_last_iteration( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe padding : int, optional Pixels to pad by post rotating-cropping object """ + xp = self._xp + asnumpy = self._asnumpy + figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) if self._object_type == "complex": obj = np.angle(self.object) @@ -2544,10 +2547,16 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: + probe_array = self.probe_fourier + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) + ) probe_array = Complex2RGB( - self.probe_fourier, + probe_array, chroma_boost=chroma_boost, ) + ax.set_title("Reconstructed Fourier probe") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -2613,6 +2622,7 @@ def _visualize_all_iterations( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, iterations_grid: Tuple[int, int], padding: int, **kwargs, @@ -2634,10 +2644,14 @@ def _visualize_all_iterations( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object """ asnumpy = self._asnumpy + xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError( @@ -2676,10 +2690,7 @@ def _visualize_all_iterations( figsize = kwargs.pop("figsize", auto_figsize) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) errors = np.array(self.error_iterations) @@ -2778,14 +2789,21 @@ def _visualize_all_iterations( for n, ax in enumerate(grid): if plot_fourier_probe: - probe_array = Complex2RGB( - asnumpy( - self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]] - ) - ), - chroma_boost=chroma_boost, + probe_array = asnumpy( + self._return_fourier_probe_from_centered_probe( + probes[grid_range[n]] + ) ) + + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift( + xp.conjugate(self._known_aberrations_array) + ) + ) + + probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) + ax.set_title(f"Iter: {grid_range[n]} Fourier probe") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -2829,6 +2847,7 @@ def visualize( plot_convergence: bool = True, plot_probe: bool = True, plot_fourier_probe: bool = False, + remove_initial_probe_aberrations: bool = False, cbar: bool = True, padding: int = 0, **kwargs, @@ -2850,6 +2869,9 @@ def visualize( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object @@ -2865,6 +2887,7 @@ def visualize( plot_convergence=plot_convergence, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, padding=padding, **kwargs, @@ -2876,6 +2899,7 @@ def visualize( iterations_grid=iterations_grid, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, padding=padding, **kwargs, diff --git a/py4DSTEM/process/phase/iterative_overlap_tomography.py b/py4DSTEM/process/phase/iterative_overlap_tomography.py index 685f95257..4f4d6eb9e 100644 --- a/py4DSTEM/process/phase/iterative_overlap_tomography.py +++ b/py4DSTEM/process/phase/iterative_overlap_tomography.py @@ -2594,6 +2594,7 @@ def _visualize_last_iteration( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, projection_angle_deg: float, projection_axes: Tuple[int, int], x_lims: Tuple[int, int], @@ -2615,6 +2616,8 @@ def _visualize_last_iteration( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe projection_angle_deg: float Angle in degrees to rotate 3D array around prior to projection projection_axes: tuple(int,int) @@ -2624,13 +2627,13 @@ def _visualize_last_iteration( y_lims: tuple(float,float) min/max y indices """ + xp = self._xp + asnumpy = self._asnumpy + figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) asnumpy = self._asnumpy @@ -2731,10 +2734,16 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: + probe_array = self.probe_fourier + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) + ) probe_array = Complex2RGB( - self.probe_fourier, + probe_array, chroma_boost=chroma_boost, ) + ax.set_title("Reconstructed Fourier probe") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -2802,6 +2811,7 @@ def _visualize_all_iterations( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, iterations_grid: Tuple[int, int], projection_angle_deg: float, projection_axes: Tuple[int, int], @@ -2824,6 +2834,9 @@ def _visualize_all_iterations( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes iterations_grid: Tuple[int,int] Grid dimensions to plot reconstruction iterations projection_angle_deg: float @@ -2836,6 +2849,7 @@ def _visualize_all_iterations( min/max y indices """ asnumpy = self._asnumpy + xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError( @@ -2877,10 +2891,7 @@ def _visualize_all_iterations( figsize = kwargs.pop("figsize", auto_figsize) cmap = kwargs.pop("cmap", "magma") - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) errors = np.array(self.error_iterations) @@ -2989,14 +3000,21 @@ def _visualize_all_iterations( for n, ax in enumerate(grid): if plot_fourier_probe: - probe_array = Complex2RGB( - asnumpy( - self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]] - ) - ), - chroma_boost=chroma_boost, + probe_array = asnumpy( + self._return_fourier_probe_from_centered_probe( + probes[grid_range[n]] + ) ) + + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift( + xp.conjugate(self._known_aberrations_array) + ) + ) + + probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) + ax.set_title(f"Iter: {grid_range[n]} Fourier probe") ax.set_ylabel("kx [mrad]") ax.set_xlabel("ky [mrad]") @@ -3042,6 +3060,7 @@ def visualize( plot_convergence: bool = True, plot_probe: bool = True, plot_fourier_probe: bool = False, + remove_initial_probe_aberrations: bool = False, cbar: bool = True, projection_angle_deg: float = None, projection_axes: Tuple[int, int] = (0, 2), @@ -3064,6 +3083,9 @@ def visualize( If true, the reconstructed probe intensity is also displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes iterations_grid: Tuple[int,int] Grid dimensions to plot reconstruction iterations projection_angle_deg: float @@ -3087,6 +3109,7 @@ def visualize( plot_convergence=plot_convergence, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, projection_angle_deg=projection_angle_deg, projection_axes=projection_axes, @@ -3101,6 +3124,7 @@ def visualize( iterations_grid=iterations_grid, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, projection_angle_deg=projection_angle_deg, projection_axes=projection_axes, diff --git a/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py b/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py index 637789108..3aac72c3b 100644 --- a/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py +++ b/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py @@ -3071,6 +3071,7 @@ def _visualize_last_iteration( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, padding: int, **kwargs, ): @@ -3089,9 +3090,15 @@ def _visualize_last_iteration( If true, the reconstructed complex probe is displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object """ + xp = self._xp + asnumpy = self._asnumpy + figsize = kwargs.pop("figsize", (12, 5)) cmap_e = kwargs.pop("cmap_e", "magma") cmap_m = kwargs.pop("cmap_m", "PuOr") @@ -3116,10 +3123,7 @@ def _visualize_last_iteration( vmin_m = kwargs.pop("vmin_m", min_m) vmax_m = kwargs.pop("vmax_m", max_m) - if plot_fourier_probe: - chroma_boost = kwargs.pop("chroma_boost", 2) - else: - chroma_boost = kwargs.pop("chroma_boost", 1) + chroma_boost = kwargs.pop("chroma_boost", 1) extent = [ 0, @@ -3224,8 +3228,13 @@ def _visualize_last_iteration( # Probe ax = fig.add_subplot(spec[0, 2]) if plot_fourier_probe: + probe_array = self.probe_fourier + if remove_initial_probe_aberrations: + probe_array *= asnumpy( + xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) + ) probe_array = Complex2RGB( - self.probe_fourier, + probe_array, chroma_boost=chroma_boost, ) ax.set_title("Reconstructed Fourier probe") @@ -3311,6 +3320,7 @@ def _visualize_all_iterations( plot_convergence: bool, plot_probe: bool, plot_fourier_probe: bool, + remove_initial_probe_aberrations: bool, iterations_grid: Tuple[int, int], padding: int, **kwargs, @@ -3332,6 +3342,9 @@ def _visualize_all_iterations( If true, the reconstructed complex probe is displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object """ @@ -3344,6 +3357,7 @@ def visualize( plot_convergence: bool = True, plot_probe: bool = True, plot_fourier_probe: bool = False, + remove_initial_probe_aberrations: bool = False, cbar: bool = True, padding: int = 0, **kwargs, @@ -3365,6 +3379,9 @@ def visualize( If true, the reconstructed complex probe is displayed plot_fourier_probe: bool, optional If true, the reconstructed complex Fourier probe is displayed + remove_initial_probe_aberrations: bool, optional + If true, when plotting fourier probe, removes initial probe + to visualize changes padding : int, optional Pixels to pad by post rotating-cropping object @@ -3380,6 +3397,7 @@ def visualize( plot_convergence=plot_convergence, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, padding=padding, **kwargs, @@ -3391,6 +3409,7 @@ def visualize( iterations_grid=iterations_grid, plot_probe=plot_probe, plot_fourier_probe=plot_fourier_probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, cbar=cbar, padding=padding, **kwargs, diff --git a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py index fb1f34c38..31c27ac2c 100644 --- a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py @@ -472,13 +472,6 @@ def preprocess( self._probe_initial = self._probe.copy() self._probe_initial_aperture = xp.abs(xp.fft.fft2(self._probe)) - self._initial_aberrations_array = ComplexProbe( - energy=self._energy, - gpts=self._region_of_interest_shape, - sampling=self.sampling, - parameters=self._polar_parameters, - device=self._device, - )._evaluate_ctf() self._known_aberrations_array = ComplexProbe( energy=self._energy, gpts=self._region_of_interest_shape, @@ -1783,6 +1776,7 @@ def _visualize_last_iteration( """ xp = self._xp asnumpy = self._asnumpy + figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") @@ -1881,7 +1875,7 @@ def _visualize_last_iteration( probe_array = self.probe_fourier if remove_initial_probe_aberrations: probe_array *= asnumpy( - xp.fft.ifftshift(xp.conjugate(self._initial_aberrations_array)) + xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) ) probe_array = Complex2RGB( probe_array, @@ -2130,7 +2124,7 @@ def _visualize_all_iterations( if remove_initial_probe_aberrations: probe_array *= asnumpy( xp.fft.ifftshift( - xp.conjugate(self._initial_aberrations_array) + xp.conjugate(self._known_aberrations_array) ) ) From ed830e05fc805f1a4dec0c131233fddff4d83c93 Mon Sep 17 00:00:00 2001 From: Georgios Varnavides Date: Mon, 20 Nov 2023 16:01:12 -0800 Subject: [PATCH 7/8] cleaning up residual aberrations plotting logic --- .../process/phase/iterative_base_class.py | 45 ++++++++++++++++--- .../iterative_singleslice_ptychography.py | 17 +++---- 2 files changed, 43 insertions(+), 19 deletions(-) diff --git a/py4DSTEM/process/phase/iterative_base_class.py b/py4DSTEM/process/phase/iterative_base_class.py index fdb0092b3..d88e09551 100644 --- a/py4DSTEM/process/phase/iterative_base_class.py +++ b/py4DSTEM/process/phase/iterative_base_class.py @@ -2154,6 +2154,7 @@ def plot_position_correction( def _return_fourier_probe( self, probe=None, + remove_initial_probe_aberrations=False, ): """ Returns complex fourier probe shifted to center of array from @@ -2163,6 +2164,8 @@ def _return_fourier_probe( ---------- probe: complex array, optional if None is specified, uses self._probe + remove_initial_probe_aberrations: bool, optional + If True, removes initial probe aberrations from Fourier probe Returns ------- @@ -2176,11 +2179,17 @@ def _return_fourier_probe( else: probe = xp.asarray(probe, dtype=xp.complex64) - return xp.fft.fftshift(xp.fft.fft2(probe), axes=(-2, -1)) + fourier_probe = xp.fft.fft2(probe) + + if remove_initial_probe_aberrations: + fourier_probe *= xp.conjugate(self._known_aberrations_array) + + return xp.fft.fftshift(fourier_probe, axes=(-2, -1)) def _return_fourier_probe_from_centered_probe( self, probe=None, + remove_initial_probe_aberrations=False, ): """ Returns complex fourier probe shifted to center of array from @@ -2190,6 +2199,8 @@ def _return_fourier_probe_from_centered_probe( ---------- probe: complex array, optional if None is specified, uses self._probe + remove_initial_probe_aberrations: bool, optional + If True, removes initial probe aberrations from Fourier probe Returns ------- @@ -2197,7 +2208,10 @@ def _return_fourier_probe_from_centered_probe( Fourier-transformed and center-shifted probe. """ xp = self._xp - return self._return_fourier_probe(xp.fft.ifftshift(probe, axes=(-2, -1))) + return self._return_fourier_probe( + xp.fft.ifftshift(probe, axes=(-2, -1)), + remove_initial_probe_aberrations=remove_initial_probe_aberrations, + ) def _return_centered_probe( self, @@ -2490,6 +2504,7 @@ def show_uncertainty_visualization( def show_fourier_probe( self, probe=None, + remove_initial_probe_aberrations=False, cbar=True, scalebar=True, pixelsize=None, @@ -2503,6 +2518,8 @@ def show_fourier_probe( ---------- probe: complex array, optional if None is specified, uses the `probe_fourier` property + remove_initial_probe_aberrations: bool, optional + If True, removes initial probe aberrations from Fourier probe cbar: bool, optional if True, adds colorbar scalebar: bool, optional @@ -2514,10 +2531,11 @@ def show_fourier_probe( """ asnumpy = self._asnumpy - if probe is None: - probe = self.probe_fourier - else: - probe = asnumpy(self._return_fourier_probe(probe)) + probe = asnumpy( + self._return_fourier_probe( + probe, remove_initial_probe_aberrations=remove_initial_probe_aberrations + ) + ) if pixelsize is None: pixelsize = self._reciprocal_sampling[1] @@ -2525,7 +2543,7 @@ def show_fourier_probe( pixelunits = r"$\AA^{-1}$" figsize = kwargs.pop("figsize", (6, 6)) - chroma_boost = kwargs.pop("chroma_boost", 2) + chroma_boost = kwargs.pop("chroma_boost", 1) fig, ax = plt.subplots(figsize=figsize) show_complex( @@ -2578,6 +2596,19 @@ def probe_fourier(self): asnumpy = self._asnumpy return asnumpy(self._return_fourier_probe(self._probe)) + @property + def probe_fourier_residual(self): + """Current probe estimate in Fourier space""" + if not hasattr(self, "_probe"): + return None + + asnumpy = self._asnumpy + return asnumpy( + self._return_fourier_probe( + self._probe, remove_initial_probe_aberrations=True + ) + ) + @property def probe_centered(self): """Current probe estimate shifted to the center""" diff --git a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py index 31c27ac2c..fa280a4ae 100644 --- a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py @@ -1872,11 +1872,10 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: - probe_array = self.probe_fourier if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) - ) + probe_array = self.probe_fourier_residual + else: + probe_array = self.probe_fourier probe_array = Complex2RGB( probe_array, chroma_boost=chroma_boost, @@ -2117,17 +2116,11 @@ def _visualize_all_iterations( if plot_fourier_probe: probe_array = asnumpy( self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]] + probes[grid_range[n]], + remove_initial_probe_aberrations=remove_initial_probe_aberrations, ) ) - if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift( - xp.conjugate(self._known_aberrations_array) - ) - ) - probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) ax.set_title(f"Iter: {grid_range[n]} Fourier probe") ax.set_ylabel("kx [mrad]") From 69845943be6a3914f82fc2048e55aae82bead785 Mon Sep 17 00:00:00 2001 From: Georgios Varnavides Date: Mon, 20 Nov 2023 16:37:05 -0800 Subject: [PATCH 8/8] residual aberration fixes across all classes --- ...tive_mixedstate_multislice_ptychography.py | 63 ++++++++++++------- .../iterative_mixedstate_ptychography.py | 55 ++++++++++------ .../iterative_multislice_ptychography.py | 30 ++++----- .../phase/iterative_overlap_tomography.py | 20 ++---- .../iterative_simultaneous_ptychography.py | 11 ++-- .../iterative_singleslice_ptychography.py | 5 +- 6 files changed, 101 insertions(+), 83 deletions(-) diff --git a/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py index e873d9199..03483c04e 100644 --- a/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py @@ -2569,9 +2569,6 @@ def _visualize_last_iteration( Pixels to pad by post rotating-cropping object """ - xp = self._xp - asnumpy = self._asnumpy - figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") @@ -2670,11 +2667,11 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: - probe_array = self.probe_fourier[0] if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) - ) + probe_array = self.probe_fourier_residual[0] + else: + probe_array = self.probe_fourier[0] + probe_array = Complex2RGB( probe_array, chroma_boost=chroma_boost, @@ -2774,7 +2771,6 @@ def _visualize_all_iterations( Pixels to pad by post rotating-cropping object """ asnumpy = self._asnumpy - xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError( @@ -2914,17 +2910,11 @@ def _visualize_all_iterations( if plot_fourier_probe: probe_array = asnumpy( self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]][0] + probes[grid_range[n]][0], + remove_initial_probe_aberrations=remove_initial_probe_aberrations, ) ) - if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift( - xp.conjugate(self._known_aberrations_array) - ) - ) - probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) ax.set_title(f"Iter: {grid_range[n]} Fourier probe[0]") @@ -3032,7 +3022,14 @@ def visualize( return self def show_fourier_probe( - self, probe=None, scalebar=True, pixelsize=None, pixelunits=None, **kwargs + self, + probe=None, + remove_initial_probe_aberrations=False, + cbar=True, + scalebar=True, + pixelsize=None, + pixelunits=None, + **kwargs, ): """ Plot probe in fourier space @@ -3041,6 +3038,8 @@ def show_fourier_probe( ---------- probe: complex array, optional if None is specified, uses the `probe_fourier` property + remove_initial_probe_aberrations: bool, optional + If True, removes initial probe aberrations from Fourier probe scalebar: bool, optional if True, adds scalebar to probe pixelunits: str, optional @@ -3051,21 +3050,37 @@ def show_fourier_probe( asnumpy = self._asnumpy if probe is None: - probe = list(self.probe_fourier) + probe = list( + asnumpy( + self._return_fourier_probe( + probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, + ) + ) + ) else: if isinstance(probe, np.ndarray) and probe.ndim == 2: probe = [probe] - probe = [asnumpy(self._return_fourier_probe(pr)) for pr in probe] + probe = [ + asnumpy( + self._return_fourier_probe( + pr, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, + ) + ) + for pr in probe + ] if pixelsize is None: pixelsize = self._reciprocal_sampling[1] if pixelunits is None: pixelunits = r"$\AA^{-1}$" - chroma_boost = kwargs.pop("chroma_boost", 2) + chroma_boost = kwargs.pop("chroma_boost", 1) show_complex( probe if len(probe) > 1 else probe[0], + cbar=cbar, scalebar=scalebar, pixelsize=pixelsize, pixelunits=pixelunits, @@ -3077,6 +3092,7 @@ def show_fourier_probe( def show_transmitted_probe( self, plot_fourier_probe: bool = False, + remove_initial_probe_aberrations=False, **kwargs, ): """ @@ -3119,7 +3135,12 @@ def show_transmitted_probe( if plot_fourier_probe: bottom_row = [ - asnumpy(self._return_fourier_probe(probe)) + asnumpy( + self._return_fourier_probe( + probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, + ) + ) for probe in [ mean_transmitted, min_intensity_transmitted, diff --git a/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py b/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py index 426216116..f3deba614 100644 --- a/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py +++ b/py4DSTEM/process/phase/iterative_mixedstate_ptychography.py @@ -1864,9 +1864,6 @@ def _visualize_last_iteration( padding : int, optional Pixels to pad by post rotating-cropping object """ - xp = self._xp - asnumpy = self._asnumpy - figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") @@ -1962,11 +1959,11 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: - probe_array = self.probe_fourier[0] if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) - ) + probe_array = self.probe_fourier_residual[0] + else: + probe_array = self.probe_fourier[0] + probe_array = Complex2RGB( probe_array, chroma_boost=chroma_boost, @@ -2068,7 +2065,6 @@ def _visualize_all_iterations( Pixels to pad by post rotating-cropping object """ asnumpy = self._asnumpy - xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError( @@ -2206,17 +2202,11 @@ def _visualize_all_iterations( if plot_fourier_probe: probe_array = asnumpy( self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]][0] + probes[grid_range[n]][0], + remove_initial_probe_aberrations=remove_initial_probe_aberrations, ) ) - if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift( - xp.conjugate(self._known_aberrations_array) - ) - ) - probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) ax.set_title(f"Iter: {grid_range[n]} Fourier probe[0]") @@ -2325,7 +2315,14 @@ def visualize( return self def show_fourier_probe( - self, probe=None, scalebar=True, pixelsize=None, pixelunits=None, **kwargs + self, + probe=None, + remove_initial_probe_aberrations=False, + cbar=True, + scalebar=True, + pixelsize=None, + pixelunits=None, + **kwargs, ): """ Plot probe in fourier space @@ -2334,6 +2331,8 @@ def show_fourier_probe( ---------- probe: complex array, optional if None is specified, uses the `probe_fourier` property + remove_initial_probe_aberrations: bool, optional + If True, removes initial probe aberrations from Fourier probe scalebar: bool, optional if True, adds scalebar to probe pixelunits: str, optional @@ -2344,21 +2343,37 @@ def show_fourier_probe( asnumpy = self._asnumpy if probe is None: - probe = list(self.probe_fourier) + probe = list( + asnumpy( + self._return_fourier_probe( + probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, + ) + ) + ) else: if isinstance(probe, np.ndarray) and probe.ndim == 2: probe = [probe] - probe = [asnumpy(self._return_fourier_probe(pr)) for pr in probe] + probe = [ + asnumpy( + self._return_fourier_probe( + pr, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, + ) + ) + for pr in probe + ] if pixelsize is None: pixelsize = self._reciprocal_sampling[1] if pixelunits is None: pixelunits = r"$\AA^{-1}$" - chroma_boost = kwargs.pop("chroma_boost", 2) + chroma_boost = kwargs.pop("chroma_boost", 1) show_complex( probe if len(probe) > 1 else probe[0], + cbar=cbar, scalebar=scalebar, pixelsize=pixelsize, pixelunits=pixelunits, diff --git a/py4DSTEM/process/phase/iterative_multislice_ptychography.py b/py4DSTEM/process/phase/iterative_multislice_ptychography.py index f2304f72b..269919ddc 100644 --- a/py4DSTEM/process/phase/iterative_multislice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_multislice_ptychography.py @@ -2446,9 +2446,6 @@ def _visualize_last_iteration( Pixels to pad by post rotating-cropping object """ - xp = self._xp - asnumpy = self._asnumpy - figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") @@ -2547,11 +2544,11 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: - probe_array = self.probe_fourier if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) - ) + probe_array = self.probe_fourier_residual + else: + probe_array = self.probe_fourier + probe_array = Complex2RGB( probe_array, chroma_boost=chroma_boost, @@ -2651,7 +2648,6 @@ def _visualize_all_iterations( Pixels to pad by post rotating-cropping object """ asnumpy = self._asnumpy - xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError( @@ -2791,17 +2787,11 @@ def _visualize_all_iterations( if plot_fourier_probe: probe_array = asnumpy( self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]] + probes[grid_range[n]], + remove_initial_probe_aberrations=remove_initial_probe_aberrations, ) ) - if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift( - xp.conjugate(self._known_aberrations_array) - ) - ) - probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) ax.set_title(f"Iter: {grid_range[n]} Fourier probe") @@ -2909,6 +2899,7 @@ def visualize( def show_transmitted_probe( self, plot_fourier_probe: bool = False, + remove_initial_probe_aberrations=False, **kwargs, ): """ @@ -2951,7 +2942,12 @@ def show_transmitted_probe( if plot_fourier_probe: bottom_row = [ - asnumpy(self._return_fourier_probe(probe)) + asnumpy( + self._return_fourier_probe( + probe, + remove_initial_probe_aberrations=remove_initial_probe_aberrations, + ) + ) for probe in [ mean_transmitted, min_intensity_transmitted, diff --git a/py4DSTEM/process/phase/iterative_overlap_tomography.py b/py4DSTEM/process/phase/iterative_overlap_tomography.py index 4f4d6eb9e..e8639d469 100644 --- a/py4DSTEM/process/phase/iterative_overlap_tomography.py +++ b/py4DSTEM/process/phase/iterative_overlap_tomography.py @@ -2627,7 +2627,6 @@ def _visualize_last_iteration( y_lims: tuple(float,float) min/max y indices """ - xp = self._xp asnumpy = self._asnumpy figsize = kwargs.pop("figsize", (8, 5)) @@ -2734,11 +2733,11 @@ def _visualize_last_iteration( ax = fig.add_subplot(spec[0, 1]) if plot_fourier_probe: - probe_array = self.probe_fourier if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) - ) + probe_array = self.probe_fourier_residual + else: + probe_array = self.probe_fourier + probe_array = Complex2RGB( probe_array, chroma_boost=chroma_boost, @@ -2849,7 +2848,6 @@ def _visualize_all_iterations( min/max y indices """ asnumpy = self._asnumpy - xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError( @@ -3002,17 +3000,11 @@ def _visualize_all_iterations( if plot_fourier_probe: probe_array = asnumpy( self._return_fourier_probe_from_centered_probe( - probes[grid_range[n]] + probes[grid_range[n]], + remove_initial_probe_aberrations=remove_initial_probe_aberrations, ) ) - if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift( - xp.conjugate(self._known_aberrations_array) - ) - ) - probe_array = Complex2RGB(probe_array, chroma_boost=chroma_boost) ax.set_title(f"Iter: {grid_range[n]} Fourier probe") diff --git a/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py b/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py index 3aac72c3b..6e31707c1 100644 --- a/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py +++ b/py4DSTEM/process/phase/iterative_simultaneous_ptychography.py @@ -3096,9 +3096,6 @@ def _visualize_last_iteration( padding : int, optional Pixels to pad by post rotating-cropping object """ - xp = self._xp - asnumpy = self._asnumpy - figsize = kwargs.pop("figsize", (12, 5)) cmap_e = kwargs.pop("cmap_e", "magma") cmap_m = kwargs.pop("cmap_m", "PuOr") @@ -3228,11 +3225,11 @@ def _visualize_last_iteration( # Probe ax = fig.add_subplot(spec[0, 2]) if plot_fourier_probe: - probe_array = self.probe_fourier if remove_initial_probe_aberrations: - probe_array *= asnumpy( - xp.fft.ifftshift(xp.conjugate(self._known_aberrations_array)) - ) + probe_array = self.probe_fourier_residual + else: + probe_array = self.probe_fourier + probe_array = Complex2RGB( probe_array, chroma_boost=chroma_boost, diff --git a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py index fa280a4ae..2c5e506e3 100644 --- a/py4DSTEM/process/phase/iterative_singleslice_ptychography.py +++ b/py4DSTEM/process/phase/iterative_singleslice_ptychography.py @@ -1774,9 +1774,6 @@ def _visualize_last_iteration( padding : int, optional Pixels to pad by post rotating-cropping object """ - xp = self._xp - asnumpy = self._asnumpy - figsize = kwargs.pop("figsize", (8, 5)) cmap = kwargs.pop("cmap", "magma") @@ -1876,6 +1873,7 @@ def _visualize_last_iteration( probe_array = self.probe_fourier_residual else: probe_array = self.probe_fourier + probe_array = Complex2RGB( probe_array, chroma_boost=chroma_boost, @@ -1977,7 +1975,6 @@ def _visualize_all_iterations( Pixels to pad by post rotating-cropping object """ asnumpy = self._asnumpy - xp = self._xp if not hasattr(self, "object_iterations"): raise ValueError(