From 6a2ab44bb5e4c47cacc3291bfde20d4d1a227732 Mon Sep 17 00:00:00 2001 From: Stephanie Ribet Date: Mon, 20 Nov 2023 15:17:57 -0800 Subject: [PATCH] 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) ) )