Merge pull request #671 from py4dstem/phase_contrast

June: where the days are long and commits non-existent

README.md

@@ -1,7 +1,7 @@
 
 > :warning: **py4DSTEM version 0.14 update** :warning: Warning: this is a major update and we expect some workflows to break.  You can still install previous versions of py4DSTEM [as discussed here](#legacyinstall)
 
-> :warning: **Phase retrieval refactor version 0.14.9** :warning: Warning: The phase-retrieval modules in py4DSTEM (DPC, parallax, and ptychography) underwent a major refactor in version 0.14.9 and as such older tutorial notebooks will not work as expected. Notably, class names have been pruned to remove the trailing "Reconstruction" (`DPCReconstruction` -> `DPC` etc.), and regularization functions have dropped the `_iter` suffix (and are instead specified as boolean flags). We are working on updating the tutorial notebooks to reflect these changes. In the meantime, there's some more information in the relevant pull request [here](https://github.com/py4dstem/py4DSTEM/pull/597#issuecomment-1890325568).
+> :warning: **Phase retrieval refactor version 0.14.9** :warning: Warning: The phase-retrieval modules in py4DSTEM (DPC, parallax, and ptychography) underwent a major refactor in version 0.14.9 and as such older tutorial notebooks will not work as expected. Notably, class names have been pruned to remove the trailing "Reconstruction" (`DPCReconstruction` -> `DPC` etc.), and regularization functions have dropped the `_iter` suffix (and are instead specified as boolean flags). See the [updated tutorials](https://github.com/py4dstem/py4DSTEM_tutorials) for more information.
 
 ![py4DSTEM logo](/images/py4DSTEM_logo.png)
 

py4DSTEM/process/phase/__init__.py

@@ -12,5 +12,6 @@
 from py4DSTEM.process.phase.ptychographic_tomography import PtychographicTomography
 from py4DSTEM.process.phase.singleslice_ptychography import SingleslicePtychography
 from py4DSTEM.process.phase.parameter_optimize import OptimizationParameter, PtychographyOptimizer
+from py4DSTEM.process.phase.xray_magnetic_ptychography import XRayMagneticPtychography
 
 # fmt: on

py4DSTEM/process/phase/magnetic_ptychography.py

@@ -101,6 +101,9 @@ class MagneticPtychography(
     initial_scan_positions: np.ndarray, optional
         Probe positions in Å for each diffraction intensity
         If None, initialized to a grid scan
+    object_fov_ang: Tuple[int,int], optional
+        Fixed object field of view in Å. If None, the fov is initialized using the
+        probe positions and object_padding_px
     positions_offset_ang: np.ndarray, optional
         Offset of positions in A
     verbose: bool, optional
@@ -138,6 +141,7 @@ def __init__(
         initial_object_guess: np.ndarray = None,
         initial_probe_guess: np.ndarray = None,
         initial_scan_positions: np.ndarray = None,
+        object_fov_ang: Tuple[float, float] = None,
         positions_offset_ang: np.ndarray = None,
         object_type: str = "complex",
         verbose: bool = True,
@@ -189,6 +193,7 @@ def __init__(
         self._rolloff = rolloff
         self._object_type = object_type
         self._object_padding_px = object_padding_px
+        self._object_fov_ang = object_fov_ang
         self._positions_mask = positions_mask
         self._verbose = verbose
         self._preprocessed = False
@@ -219,6 +224,7 @@ def preprocess(
         progress_bar: bool = True,
         object_fov_mask: np.ndarray = True,
         crop_patterns: bool = False,
+        center_positions_in_fov: bool = True,
         store_initial_arrays: bool = True,
         device: str = None,
         clear_fft_cache: bool = None,
@@ -286,6 +292,8 @@ def preprocess(
             If None, probe_overlap intensity is thresholded
         crop_patterns: bool
             if True, crop patterns to avoid wrap around of patterns when centering
+        center_positions_in_fov: bool
+            If True (default), probe positions are centered in the fov.
         store_initial_arrays: bool
             If True, preprocesed object and probe arrays are stored allowing reset=True in reconstruct.
         device: str, optional
@@ -610,14 +618,17 @@ def preprocess(
             self._positions_px_all, dtype=xp_storage.float32
         )
 
-        for index in range(self._num_measurements):
-            idx_start = self._cum_probes_per_measurement[index]
-            idx_end = self._cum_probes_per_measurement[index + 1]
+        if center_positions_in_fov:
+            for index in range(self._num_measurements):
+                idx_start = self._cum_probes_per_measurement[index]
+                idx_end = self._cum_probes_per_measurement[index + 1]
 
-            positions_px = self._positions_px_all[idx_start:idx_end]
-            positions_px_com = positions_px.mean(0)
-            positions_px -= positions_px_com - xp_storage.array(self._object_shape) / 2
-            self._positions_px_all[idx_start:idx_end] = positions_px.copy()
+                positions_px = self._positions_px_all[idx_start:idx_end]
+                positions_px_com = positions_px.mean(0)
+                positions_px -= (
+                    positions_px_com - xp_storage.array(self._object_shape) / 2
+                )
+                self._positions_px_all[idx_start:idx_end] = positions_px.copy()
 
         self._positions_px_initial_all = self._positions_px_all.copy()
         self._positions_initial_all = self._positions_px_initial_all.copy()
@@ -1785,55 +1796,55 @@ def _visualize_last_iteration(
         if fig is None:
             fig = plt.figure(figsize=figsize)
 
-        if plot_probe or plot_fourier_probe:
-            # Object_e
-            ax = fig.add_subplot(spec[0, 0])
-            im = ax.imshow(
-                obj[0],
-                extent=extent,
-                cmap=cmap_e,
-                vmin=vmin_e,
-                vmax=vmax_e,
-                **kwargs,
-            )
-            ax.set_ylabel("x [A]")
-            ax.set_xlabel("y [A]")
-
-            if self._object_type == "potential":
-                ax.set_title("Electrostatic potential")
-            elif self._object_type == "complex":
-                ax.set_title("Electrostatic phase")
+        # Object_e
+        ax = fig.add_subplot(spec[0, 0])
+        im = ax.imshow(
+            obj[0],
+            extent=extent,
+            cmap=cmap_e,
+            vmin=vmin_e,
+            vmax=vmax_e,
+            **kwargs,
+        )
+        ax.set_ylabel("x [A]")
+        ax.set_xlabel("y [A]")
 
-            if cbar:
-                divider = make_axes_locatable(ax)
-                ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
-                fig.add_axes(ax_cb)
-                fig.colorbar(im, cax=ax_cb)
+        if self._object_type == "potential":
+            ax.set_title("Electrostatic potential")
+        elif self._object_type == "complex":
+            ax.set_title("Electrostatic phase")
+
+        if cbar:
+            divider = make_axes_locatable(ax)
+            ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
+            fig.add_axes(ax_cb)
+            fig.colorbar(im, cax=ax_cb)
+
+        # Object_m
+        ax = fig.add_subplot(spec[0, 1])
+        im = ax.imshow(
+            obj[1],
+            extent=extent,
+            cmap=cmap_m,
+            vmin=vmin_m,
+            vmax=vmax_m,
+            **kwargs,
+        )
+        ax.set_ylabel("x [A]")
+        ax.set_xlabel("y [A]")
 
-            # Object_m
-            ax = fig.add_subplot(spec[0, 1])
-            im = ax.imshow(
-                obj[1],
-                extent=extent,
-                cmap=cmap_m,
-                vmin=vmin_m,
-                vmax=vmax_m,
-                **kwargs,
-            )
-            ax.set_ylabel("x [A]")
-            ax.set_xlabel("y [A]")
+        if self._object_type == "potential":
+            ax.set_title("Magnetic potential")
+        elif self._object_type == "complex":
+            ax.set_title("Magnetic phase")
 
-            if self._object_type == "potential":
-                ax.set_title("Magnetic potential")
-            elif self._object_type == "complex":
-                ax.set_title("Magnetic phase")
-
-            if cbar:
-                divider = make_axes_locatable(ax)
-                ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
-                fig.add_axes(ax_cb)
-                fig.colorbar(im, cax=ax_cb)
+        if cbar:
+            divider = make_axes_locatable(ax)
+            ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
+            fig.add_axes(ax_cb)
+            fig.colorbar(im, cax=ax_cb)
 
+        if plot_probe or plot_fourier_probe:
             # Probe
             ax = fig.add_subplot(spec[0, 2])
             if plot_fourier_probe:
@@ -1872,55 +1883,6 @@ def _visualize_last_iteration(
                 ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
                 add_colorbar_arg(ax_cb, chroma_boost=chroma_boost)
 
-        else:
-            # Object_e
-            ax = fig.add_subplot(spec[0, 0])
-            im = ax.imshow(
-                obj[0],
-                extent=extent,
-                cmap=cmap_e,
-                vmin=vmin_e,
-                vmax=vmax_e,
-                **kwargs,
-            )
-            ax.set_ylabel("x [A]")
-            ax.set_xlabel("y [A]")
-
-            if self._object_type == "potential":
-                ax.set_title("Electrostatic potential")
-            elif self._object_type == "complex":
-                ax.set_title("Electrostatic phase")
-
-            if cbar:
-                divider = make_axes_locatable(ax)
-                ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
-                fig.add_axes(ax_cb)
-                fig.colorbar(im, cax=ax_cb)
-
-            # Object_e
-            ax = fig.add_subplot(spec[0, 1])
-            im = ax.imshow(
-                obj[1],
-                extent=extent,
-                cmap=cmap_m,
-                vmin=vmin_m,
-                vmax=vmax_m,
-                **kwargs,
-            )
-            ax.set_ylabel("x [A]")
-            ax.set_xlabel("y [A]")
-
-            if self._object_type == "potential":
-                ax.set_title("Magnetic potential")
-            elif self._object_type == "complex":
-                ax.set_title("Magnetic phase")
-
-            if cbar:
-                divider = make_axes_locatable(ax)
-                ax_cb = divider.append_axes("right", size="5%", pad="2.5%")
-                fig.add_axes(ax_cb)
-                fig.colorbar(im, cax=ax_cb)
-
         if plot_convergence and hasattr(self, "error_iterations"):
             errors = np.array(self.error_iterations)
 

py4DSTEM/process/phase/mixedstate_multislice_ptychography.py

@@ -107,6 +107,9 @@ class MixedstateMultislicePtychography(
     initial_scan_positions: np.ndarray, optional
         Probe positions in Å for each diffraction intensity
         If None, initialized to a grid scan
+    object_fov_ang: Tuple[int,int], optional
+        Fixed object field of view in Å. If None, the fov is initialized using the
+        probe positions and object_padding_px
     positions_offset_ang: np.ndarray, optional
         Offset of positions in A
     theta_x: float
@@ -159,6 +162,7 @@ def __init__(
         initial_object_guess: np.ndarray = None,
         initial_probe_guess: np.ndarray = None,
         initial_scan_positions: np.ndarray = None,
+        object_fov_ang: Tuple[float, float] = None,
         positions_offset_ang: np.ndarray = None,
         theta_x: float = 0,
         theta_y: float = 0,
@@ -245,6 +249,7 @@ def __init__(
         self._object_type = object_type
         self._positions_mask = positions_mask
         self._object_padding_px = object_padding_px
+        self._object_fov_ang = object_fov_ang
         self._verbose = verbose
         self._preprocessed = False
 
@@ -278,6 +283,7 @@ def preprocess(
         force_reciprocal_sampling: float = None,
         object_fov_mask: np.ndarray = None,
         crop_patterns: bool = False,
+        center_positions_in_fov: bool = True,
         store_initial_arrays: bool = True,
         device: str = None,
         clear_fft_cache: bool = None,
@@ -348,6 +354,8 @@ def preprocess(
             If None, probe_overlap intensity is thresholded
         crop_patterns: bool
             if True, crop patterns to avoid wrap around of patterns when centering
+        center_positions_in_fov: bool
+            If True (default), probe positions are centered in the fov.
         store_initial_arrays: bool
             If True, preprocesed object and probe arrays are stored allowing reset=True in reconstruct.
         device: str, optional
@@ -524,15 +532,18 @@ def preprocess(
             self._object_type_initial = self._object_type
         self._object_shape = self._object.shape[-2:]
 
-        # center probe positions
         self._positions_px = xp_storage.asarray(
             self._positions_px, dtype=xp_storage.float32
         )
         self._positions_px_initial_com = self._positions_px.mean(0)
-        self._positions_px -= (
-            self._positions_px_initial_com - xp_storage.array(self._object_shape) / 2
-        )
-        self._positions_px_initial_com = self._positions_px.mean(0)
+
+        # center probe positions
+        if center_positions_in_fov:
+            self._positions_px -= (
+                self._positions_px_initial_com
+                - xp_storage.array(self._object_shape) / 2
+            )
+            self._positions_px_initial_com = self._positions_px.mean(0)
 
         self._positions_px_initial = self._positions_px.copy()
         self._positions_initial = self._positions_px_initial.copy()

py4DSTEM/process/phase/mixedstate_ptychography.py

@@ -96,6 +96,9 @@ class MixedstatePtychography(
     initial_scan_positions: np.ndarray, optional
         Probe positions in Å for each diffraction intensity
         If None, initialized to a grid scan
+    object_fov_ang: Tuple[int,int], optional
+        Fixed object field of view in Å. If None, the fov is initialized using the
+        probe positions and object_padding_px
     positions_offset_ang: np.ndarray, optional
         Offset of positions in A
     positions_mask: np.ndarray, optional
@@ -127,6 +130,7 @@ def __init__(
         initial_object_guess: np.ndarray = None,
         initial_probe_guess: np.ndarray = None,
         initial_scan_positions: np.ndarray = None,
+        object_fov_ang: Tuple[float, float] = None,
         positions_offset_ang: np.ndarray = None,
         object_type: str = "complex",
         positions_mask: np.ndarray = None,
@@ -194,6 +198,7 @@ def __init__(
         self._rolloff = rolloff
         self._object_type = object_type
         self._object_padding_px = object_padding_px
+        self._object_fov_ang = object_fov_ang
         self._positions_mask = positions_mask
         self._verbose = verbose
         self._preprocessed = False
@@ -224,6 +229,7 @@ def preprocess(
         force_reciprocal_sampling: float = None,
         object_fov_mask: np.ndarray = None,
         crop_patterns: bool = False,
+        center_positions_in_fov: bool = True,
         store_initial_arrays: bool = True,
         device: str = None,
         clear_fft_cache: bool = None,
@@ -294,6 +300,8 @@ def preprocess(
             If None, probe_overlap intensity is thresholded
         crop_patterns: bool
             if True, crop patterns to avoid wrap around of patterns when centering
+        center_positions_in_fov: bool
+            If True (default), probe positions are centered in the fov.
         store_initial_arrays: bool
             If True, preprocesed object and probe arrays are stored allowing reset=True in reconstruct.
         device: str, optional
@@ -469,15 +477,18 @@ def preprocess(
             self._object_type_initial = self._object_type
         self._object_shape = self._object.shape
 
-        # center probe positions
         self._positions_px = xp_storage.asarray(
             self._positions_px, dtype=xp_storage.float32
         )
         self._positions_px_initial_com = self._positions_px.mean(0)
-        self._positions_px -= (
-            self._positions_px_initial_com - xp_storage.array(self._object_shape) / 2
-        )
-        self._positions_px_initial_com = self._positions_px.mean(0)
+
+        # center probe positions
+        if center_positions_in_fov:
+            self._positions_px -= (
+                self._positions_px_initial_com
+                - xp_storage.array(self._object_shape) / 2
+            )
+            self._positions_px_initial_com = self._positions_px.mean(0)
 
         self._positions_px_initial = self._positions_px.copy()
         self._positions_initial = self._positions_px_initial.copy()