From bab740671565f471474602c98fff1fc4f63251dc Mon Sep 17 00:00:00 2001
From: Stephanie Ribet <stephanie.ribet@gmail.com>
Date: Sun, 5 Nov 2023 09:17:23 -0800
Subject: [PATCH] more parallax plotting fun(ctionality)

---
 py4DSTEM/process/phase/iterative_parallax.py | 23 +++++++++++++-----
 py4DSTEM/visualize/vis_special.py            | 25 ++++++++++++++++++++
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/py4DSTEM/process/phase/iterative_parallax.py b/py4DSTEM/process/phase/iterative_parallax.py
index 9f690c434..716e1d782 100644
--- a/py4DSTEM/process/phase/iterative_parallax.py
+++ b/py4DSTEM/process/phase/iterative_parallax.py
@@ -587,16 +587,27 @@ def preprocess(
         self.recon_BF = asnumpy(self._recon_BF)
 
         if plot_average_bf:
-            figsize = kwargs.pop("figsize", (6, 6))
+            figsize = kwargs.pop("figsize", (6, 12))
 
-            fig, ax = plt.subplots(figsize=figsize)
+            fig, ax = plt.subplots(1, 2, figsize=figsize)
 
-            self._visualize_figax(fig, ax, **kwargs)
+            self._visualize_figax(fig, ax[0], **kwargs)
 
-            ax.set_ylabel("x [A]")
-            ax.set_xlabel("y [A]")
-            ax.set_title("Average Bright Field Image")
+            ax[0].set_ylabel("x [A]")
+            ax[0].set_xlabel("y [A]")
+            ax[0].set_title("Average Bright Field Image")
 
+            reciprocal_extent = [
+                -0.5 * (self._reciprocal_sampling[1] * self._dp_mask.shape[1]),
+                0.5 * (self._reciprocal_sampling[1] * self._dp_mask.shape[1]),
+                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].set_title("DP mask")
+            ax[1].set_ylabel(r"$k_x$ [$A^{-1}$]")
+            ax[1].set_xlabel(r"$k_y$ [$A^{-1}$]")
+            plt.tight_layout()
         self._preprocessed = True
 
         if self._device == "gpu":
diff --git a/py4DSTEM/visualize/vis_special.py b/py4DSTEM/visualize/vis_special.py
index da501c746..1d46ebf44 100644
--- a/py4DSTEM/visualize/vis_special.py
+++ b/py4DSTEM/visualize/vis_special.py
@@ -842,6 +842,31 @@ def show_complex(
 
 
 def return_scaled_histogram_ordering(array, vmin=None, vmax=None, normalize=False):
+    """
+    Utility function for calculating min and max values for plotting array
+    based on distribution of pixel values
+
+    Parameters
+    ----------
+    array: np.array
+        array to be plotted
+    vmin: float 
+        lower fraction cut off of pixel values 
+    vmax: float
+        upper fraction cut off of pixel values 
+    normalize: bool
+        if True, rescales from 0 to 1
+
+    Returns
+    ----------
+    scaled_array: np.array
+        array clipped outside vmin and vmax
+    vmin: float
+        lower value to be plotted
+    vmax: float
+        upper value to be plotted
+    """
+
     if vmin is None:
         vmin = 0.02
     if vmax is None: