Skip to content

Commit

Permalink
everyone gets new fourier probe plotting
Browse files Browse the repository at this point in the history
  • Loading branch information
@smribet
smribet committed Nov 20, 2023
1 parent 38234b1 commit 6a2ab44
Show file tree
Hide file tree
Showing 6 changed files with 188 additions and 78 deletions.
57 changes: 41 additions & 16 deletions py4DSTEM/process/phase/iterative_mixedstate_multislice_ptychography.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
):
Expand All @@ -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)
Expand Down Expand Up @@ -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]")
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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]")
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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,
Expand Down
56 changes: 40 additions & 16 deletions py4DSTEM/process/phase/iterative_mixedstate_ptychography.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
):
Expand All @@ -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)
Expand Down Expand Up @@ -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]")
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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]")
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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,
Expand Down
Loading

0 comments on commit 6a2ab44

Please sign in to comment.