diff --git a/py4DSTEM/__init__.py b/py4DSTEM/__init__.py
index 0f4490eb2..e26f24c78 100644
--- a/py4DSTEM/__init__.py
+++ b/py4DSTEM/__init__.py
@@ -25,7 +25,6 @@
 from py4DSTEM import io
 from py4DSTEM.io import import_file, read, save
 
-
 ### basic data classes
 
 # data
@@ -40,7 +39,6 @@
 # datacube
 from py4DSTEM.datacube import DataCube, VirtualImage, VirtualDiffraction
 
-
 ### visualization
 
 from py4DSTEM import visualize
@@ -64,15 +62,12 @@
 # diffraction
 from py4DSTEM.process.diffraction import Crystal, Orientation
 
-
 # ptycho
 from py4DSTEM.process import phase
 
-
 # polar
 from py4DSTEM.process.polar import PolarDatacube
 
-
 # strain
 from py4DSTEM.process.strain.strain import StrainMap
 
@@ -89,7 +84,6 @@
 from py4DSTEM import preprocess
 from py4DSTEM import process
 
-
 ### utilities
 
 # config
diff --git a/py4DSTEM/braggvectors/diskdetection_parallel_new.py b/py4DSTEM/braggvectors/diskdetection_parallel_new.py
index dccc0dd4b..637237a46 100644
--- a/py4DSTEM/braggvectors/diskdetection_parallel_new.py
+++ b/py4DSTEM/braggvectors/diskdetection_parallel_new.py
@@ -21,7 +21,6 @@
 from emdfile import PointListArray, PointList
 from py4DSTEM.braggvectors.diskdetection import _find_Bragg_disks_single_DP_FK
 
-
 #### SERIALISERS ####
 # Define Serialiser
 # these are functions which allow the hdf5 objects to be passed. May not be required anymore
diff --git a/py4DSTEM/data/propagating_calibration.py b/py4DSTEM/data/propagating_calibration.py
index 4de0c8d96..774b6ab86 100644
--- a/py4DSTEM/data/propagating_calibration.py
+++ b/py4DSTEM/data/propagating_calibration.py
@@ -3,7 +3,6 @@
 
 import warnings
 
-
 # This is the abstract pattern:
 
 
diff --git a/py4DSTEM/datacube/virtualimage.py b/py4DSTEM/datacube/virtualimage.py
index d4fe15241..98c71f3d9 100644
--- a/py4DSTEM/datacube/virtualimage.py
+++ b/py4DSTEM/datacube/virtualimage.py
@@ -13,7 +13,6 @@
 from py4DSTEM.preprocess import get_shifted_ar
 from py4DSTEM.visualize import show
 
-
 # Virtual image container class
 
 
diff --git a/py4DSTEM/io/__init__.py b/py4DSTEM/io/__init__.py
index fa7cd099e..a8cdcd700 100644
--- a/py4DSTEM/io/__init__.py
+++ b/py4DSTEM/io/__init__.py
@@ -3,6 +3,5 @@
 from py4DSTEM.io.read import read
 from py4DSTEM.io.save import save
 
-
 # google downloader
 from py4DSTEM.io.google_drive_downloader import gdrive_download, get_sample_file_ids
diff --git a/py4DSTEM/io/google_drive_downloader.py b/py4DSTEM/io/google_drive_downloader.py
index a50fb2c09..bc41e764d 100644
--- a/py4DSTEM/io/google_drive_downloader.py
+++ b/py4DSTEM/io/google_drive_downloader.py
@@ -2,7 +2,6 @@
 import os
 import warnings
 
-
 ### File IDs
 
 # single files
diff --git a/py4DSTEM/io/legacy/legacy13/v13_emd_classes/io.py b/py4DSTEM/io/legacy/legacy13/v13_emd_classes/io.py
index 0452a5854..4cdab267a 100644
--- a/py4DSTEM/io/legacy/legacy13/v13_emd_classes/io.py
+++ b/py4DSTEM/io/legacy/legacy13/v13_emd_classes/io.py
@@ -6,7 +6,6 @@
 from numbers import Number
 from emdfile import tqdmnd
 
-
 # Define the EMD group types
 
 EMD_group_types = {
diff --git a/py4DSTEM/io/legacy/legacy13/v13_py4dstem_classes/io.py b/py4DSTEM/io/legacy/legacy13/v13_py4dstem_classes/io.py
index 2556ebe8f..d50c4e872 100644
--- a/py4DSTEM/io/legacy/legacy13/v13_py4dstem_classes/io.py
+++ b/py4DSTEM/io/legacy/legacy13/v13_py4dstem_classes/io.py
@@ -18,7 +18,6 @@
     _read_metadata,
 )
 
-
 # Calibration
 
 
diff --git a/py4DSTEM/io/legacy/legacy13/v13_to_14.py b/py4DSTEM/io/legacy/legacy13/v13_to_14.py
index 650529b22..444293daa 100644
--- a/py4DSTEM/io/legacy/legacy13/v13_to_14.py
+++ b/py4DSTEM/io/legacy/legacy13/v13_to_14.py
@@ -3,7 +3,6 @@
 import numpy as np
 from emdfile import tqdmnd
 
-
 # v13 imports
 
 from py4DSTEM.io.legacy.legacy13.v13_emd_classes import (
@@ -25,7 +24,6 @@
     BraggVectors as BraggVectors13,
 )
 
-
 # v14 imports
 
 from emdfile import Root, Metadata, Array, PointList, PointListArray
diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py
index f7a4a20db..c1fc46c9e 100644
--- a/py4DSTEM/process/diffraction/crystal_ACOM.py
+++ b/py4DSTEM/process/diffraction/crystal_ACOM.py
@@ -335,7 +335,8 @@ def orientation_plan(
     )
     self.orientation_zone_axis_steps = (
         np.round(step / self.orientation_refine_ratio) * self.orientation_refine_ratio
-    ).astype(np.integer)
+    ).astype(np.int32)
+    # ).astype(np.integer)
 
     if self.orientation_fiber and self.orientation_fiber_angles[0] == 0:
         self.orientation_num_zones = int(1)
@@ -370,7 +371,8 @@ def orientation_plan(
             (self.orientation_zone_axis_steps + 1)
             * (self.orientation_zone_axis_steps + 2)
             / 2
-        ).astype(np.integer)
+        ).astype(np.int32)
+        # ).astype(np.integer)
         self.orientation_vecs = np.zeros((self.orientation_num_zones, 3))
         self.orientation_vecs[0, :] = self.orientation_zone_axis_range[0, :]
         self.orientation_inds = np.zeros((self.orientation_num_zones, 3), dtype="int")
@@ -379,7 +381,8 @@ def orientation_plan(
         # or circular arc SLERP for fiber texture
         for a0 in np.arange(1, self.orientation_zone_axis_steps + 1):
             inds = np.arange(a0 * (a0 + 1) / 2, a0 * (a0 + 1) / 2 + a0 + 1).astype(
-                np.integer
+                np.int32
+                # np.integer
             )
 
             p0 = pv[a0, :]
@@ -617,7 +620,8 @@ def orientation_plan(
 
     # Solve for number of angular steps along in-plane rotation direction
     self.orientation_in_plane_steps = np.round(360 / angle_step_in_plane).astype(
-        np.integer
+        np.int32
+        # np.integer
     )
 
     # Calculate -z angles (Euler angle 3)
@@ -2208,8 +2212,6 @@ def calculate_strain(
     deformation tensor which transforms the simulated diffraction pattern
     into the experimental pattern, for all probe positons.
 
-    TODO: add robust fitting?
-
     Parameters
     ----------
     bragg_peaks_array (PointListArray):
@@ -2330,71 +2332,75 @@ def calculate_strain(
                     inds_match[a0] = ind_min
                     keep[a0] = True
 
-            # Get all paired peaks
-            qxy = np.vstack((p.data["qx"][keep], p.data["qy"][keep])).T
-            qxy_ref = np.vstack(
-                (p_ref.data["qx"][inds_match[keep]], p_ref.data["qy"][inds_match[keep]])
-            ).T
+            if np.sum(keep) >= min_num_peaks:
+                # Get all paired peaks
+                qxy = np.vstack((p.data["qx"][keep], p.data["qy"][keep])).T
+                qxy_ref = np.vstack(
+                    (
+                        p_ref.data["qx"][inds_match[keep]],
+                        p_ref.data["qy"][inds_match[keep]],
+                    )
+                ).T
 
-            # Fit transformation matrix
-            # Note - not sure about transpose here
-            # (though it might not matter if rotation isn't included)
-            if intensity_weighting:
-                weights = np.sqrt(p.data["intensity"][keep, None]) * 0 + 1
-                m = lstsq(
-                    qxy_ref * weights,
-                    qxy * weights,
-                    rcond=None,
-                )[0].T
-            else:
-                m = lstsq(
-                    qxy_ref,
-                    qxy,
-                    rcond=None,
-                )[0].T
-
-            # Robust fitting
-            if robust:
-                for a0 in range(5):
-                    # calculate new weights
-                    qxy_fit = qxy_ref @ m
-                    diff2 = np.sum((qxy_fit - qxy) ** 2, axis=1)
-
-                    weights = np.exp(
-                        diff2 / ((-2 * robust_thresh**2) * np.median(diff2))
-                    )[:, None]
-                    if intensity_weighting:
-                        weights *= np.sqrt(p.data["intensity"][keep, None])
-
-                    # calculate new fits
+                # Fit transformation matrix
+                # Note - not sure about transpose here
+                # (though it might not matter if rotation isn't included)
+                if intensity_weighting:
+                    weights = np.sqrt(p.data["intensity"][keep, None]) * 0 + 1
                     m = lstsq(
                         qxy_ref * weights,
                         qxy * weights,
                         rcond=None,
                     )[0].T
+                else:
+                    m = lstsq(
+                        qxy_ref,
+                        qxy,
+                        rcond=None,
+                    )[0].T
 
-            # Set values into the infinitesimal strain matrix
-            strain_map.get_slice("e_xx").data[rx, ry] = 1 - m[0, 0]
-            strain_map.get_slice("e_yy").data[rx, ry] = 1 - m[1, 1]
-            strain_map.get_slice("e_xy").data[rx, ry] = -(m[0, 1] + m[1, 0]) / 2.0
-            strain_map.get_slice("theta").data[rx, ry] = (m[0, 1] - m[1, 0]) / 2.0
-
-            # Add finite rotation from ACOM orientation map.
-            # I am not sure about the relative signs here.
-            # Also, maybe I need to add in the mirror operator?
-            if orientation_map.mirror[rx, ry, 0]:
-                strain_map.get_slice("theta").data[rx, ry] += (
-                    orientation_map.angles[rx, ry, 0, 0]
-                    + orientation_map.angles[rx, ry, 0, 2]
-                )
-            else:
-                strain_map.get_slice("theta").data[rx, ry] -= (
-                    orientation_map.angles[rx, ry, 0, 0]
-                    + orientation_map.angles[rx, ry, 0, 2]
-                )
+                # Robust fitting
+                if robust:
+                    for a0 in range(5):
+                        # calculate new weights
+                        qxy_fit = qxy_ref @ m
+                        diff2 = np.sum((qxy_fit - qxy) ** 2, axis=1)
+
+                        weights = np.exp(
+                            diff2 / ((-2 * robust_thresh**2) * np.median(diff2))
+                        )[:, None]
+                        if intensity_weighting:
+                            weights *= np.sqrt(p.data["intensity"][keep, None])
+
+                        # calculate new fits
+                        m = lstsq(
+                            qxy_ref * weights,
+                            qxy * weights,
+                            rcond=None,
+                        )[0].T
+
+                # Set values into the infinitesimal strain matrix
+                strain_map.get_slice("e_xx").data[rx, ry] = 1 - m[0, 0]
+                strain_map.get_slice("e_yy").data[rx, ry] = 1 - m[1, 1]
+                strain_map.get_slice("e_xy").data[rx, ry] = -(m[0, 1] + m[1, 0]) / 2.0
+                strain_map.get_slice("theta").data[rx, ry] = (m[0, 1] - m[1, 0]) / 2.0
+
+                # Add finite rotation from ACOM orientation map.
+                # I am not sure about the relative signs here.
+                # Also, maybe I need to add in the mirror operator?
+                if orientation_map.mirror[rx, ry, 0]:
+                    strain_map.get_slice("theta").data[rx, ry] += (
+                        orientation_map.angles[rx, ry, 0, 0]
+                        + orientation_map.angles[rx, ry, 0, 2]
+                    )
+                else:
+                    strain_map.get_slice("theta").data[rx, ry] -= (
+                        orientation_map.angles[rx, ry, 0, 0]
+                        + orientation_map.angles[rx, ry, 0, 2]
+                    )
 
-        else:
-            strain_map.get_slice("mask").data[rx, ry] = 0.0
+            else:
+                strain_map.get_slice("mask").data[rx, ry] = 0.0
 
     if rotation_range is not None:
         strain_map.get_slice("theta").data[:] = np.mod(
diff --git a/py4DSTEM/process/diffraction/crystal_viz.py b/py4DSTEM/process/diffraction/crystal_viz.py
index 45de2461b..9532aa41a 100644
--- a/py4DSTEM/process/diffraction/crystal_viz.py
+++ b/py4DSTEM/process/diffraction/crystal_viz.py
@@ -1121,6 +1121,8 @@ def plot_orientation_maps(
     corr_normalize: bool = True,
     show_legend: bool = True,
     scale_legend: bool = None,
+    legend_hist: bool = False,
+    scale_hist: float = 2,
     figsize: Union[list, tuple, np.ndarray] = (16, 5),
     figbound: Union[list, tuple, np.ndarray] = (0.01, 0.005),
     show_axes: bool = True,
@@ -1128,6 +1130,8 @@ def plot_orientation_maps(
     plot_limit=None,
     plot_layout=0,
     swap_axes_xy_limits=False,
+    crop_range=None,
+    figax=None,
     returnfig: bool = False,
     progress_bar=False,
 ):
@@ -1144,6 +1148,8 @@ def plot_orientation_maps(
         corr_normalize (bool):              If true, set mean correlation to 1.
         show_legend (bool):                 Show the legend
         scale_legend (float):               2 elements, x and y scaling of legend panel
+        legend_hist(bool):                  If True, adds histogram of detected orientations to legend
+        scale_hist (float):                 Amount to scale histogram values by
         figsize (array):                    2 elements defining figure size
         figbound (array):                   2 elements defining figure boundary
         show_axes (bool):                   Flag setting whether orienation map axes are visible.
@@ -1152,6 +1158,8 @@ def plot_orientation_maps(
         plot_layout (int):                  subplot layout: 0 - 1 row, 3 col
                                                             1 - 3 row, 1 col
         swap_axes_xy_limits (bool):         swap x and y boundaries for legend (not sure why we need this in some cases)
+        crop_range (4-tuple):               region to crop for making orientation maps
+        figax (matplotlib figure):          figure and axes for plotting (3 axes needed)
         returnfig (bool):                   set to True to return figure and axes handles
         progress_bar (bool):                Enable progressbar when calculating orientation images.
 
@@ -1251,6 +1259,8 @@ def plot_orientation_maps(
         mask = (corr - corr_range[0]) / (corr_range[1] - corr_range[0])
         mask = np.clip(mask, 0, 1)
 
+    self.mask_corr = mask
+
     # Generate images
     for rx, ry in tqdmnd(
         orientation_map.num_x,
@@ -1289,11 +1299,14 @@ def plot_orientation_maps(
                 )
     basis_x = np.clip(basis_x, 0, 1)
     basis_z = np.clip(basis_z, 0, 1)
+    self.basis_z = basis_z
+    self.A = A
 
     # Convert to RGB images
     basis_x_max = np.max(basis_x, axis=2)
     sub = basis_x_max > 0
     basis_x_scale = basis_x * mask[:, :, None]
+    basis_x_scale_binary = basis_x * (mask > 0)[:, :, None]
     if self.pointgroup.get_crystal_system() == "monoclinic":
         rgb_x = (
             basis_x_scale[:, :, 0][:, :, None] * np.array((1, 1, 1))[None, None, :]
@@ -1307,6 +1320,8 @@ def plot_orientation_maps(
         for a0 in range(3):
             basis_x_scale[:, :, a0][sub] /= basis_x_max[sub]
             basis_x_scale[:, :, a0][np.logical_not(sub)] = 0
+            basis_x_scale_binary[:, :, a0][sub] /= basis_x_max[sub]
+            basis_x_scale_binary[:, :, a0][np.logical_not(sub)] = 0
         rgb_x = (
             basis_x_scale[:, :, 0][:, :, None] * color_basis[0, :][None, None, :]
             + basis_x_scale[:, :, 1][:, :, None] * color_basis[1, :][None, None, :]
@@ -1316,6 +1331,7 @@ def plot_orientation_maps(
     basis_z_max = np.max(basis_z, axis=2)
     sub = basis_z_max > 0
     basis_z_scale = basis_z * mask[:, :, None]
+    basis_z_scale_binary = basis_z * (mask > 0)[:, :, None]
     if self.pointgroup.get_crystal_system() == "monoclinic":
         rgb_z = (
             basis_z_scale[:, :, 0][:, :, None] * np.array((1, 1, 1))[None, None, :]
@@ -1329,6 +1345,8 @@ def plot_orientation_maps(
         for a0 in range(3):
             basis_z_scale[:, :, a0][sub] /= basis_z_max[sub]
             basis_z_scale[:, :, a0][np.logical_not(sub)] = 0
+            basis_z_scale_binary[:, :, a0][sub] /= basis_z_max[sub]
+            basis_z_scale_binary[:, :, a0][np.logical_not(sub)] = 0
         rgb_z = (
             basis_z_scale[:, :, 0][:, :, None] * color_basis[0, :][None, None, :]
             + basis_z_scale[:, :, 1][:, :, None] * color_basis[1, :][None, None, :]
@@ -1337,6 +1355,7 @@ def plot_orientation_maps(
 
     rgb_x = np.clip(rgb_x, 0, 1)
     rgb_z = np.clip(rgb_z, 0, 1)
+    self.basis_z_scale = basis_z_scale
 
     # if np.abs(self.cell[4] - 120.0) < 1e-6 or np.abs(self.cell[5] - 120.0) or np.abs(self.cell[6] - 120.0):
     #     label_0 = self.rational_ind(
@@ -1558,26 +1577,35 @@ def plot_orientation_maps(
 
     # plotting frame
     # fig, ax = plt.subplots(1, 3, figsize=figsize)
-    fig = plt.figure(figsize=figsize)
-    if plot_layout == 0:
-        ax_x = fig.add_axes([0.0 + figbound[0], 0.0, 0.4 - 2 * +figbound[0], 1.0])
-        ax_z = fig.add_axes([0.4 + figbound[0], 0.0, 0.4 - 2 * +figbound[0], 1.0])
-        ax_l = fig.add_axes(
-            [0.8 + figbound[0], 0.0, 0.2 - 2 * +figbound[0], 1.0],
-            # projection="3d",
-            # elev=el,
-            # azim=az,
-        )
-    elif plot_layout == 1:
-        ax_x = fig.add_axes([0.0, 0.0 + figbound[0], 1.0, 0.4 - 2 * +figbound[0]])
-        ax_z = fig.add_axes([0.0, 0.4 + figbound[0], 1.0, 0.4 - 2 * +figbound[0]])
-        ax_l = fig.add_axes(
-            [0.0, 0.8 + figbound[0], 1.0, 0.2 - 2 * +figbound[0]],
-            # projection="3d",
-            # elev=el,
-            # azim=az,
-        )
-
+    if figax is None:
+        fig = plt.figure(figsize=figsize)
+        if plot_layout == 0:
+            ax_x = fig.add_axes([0.0 + figbound[0], 0.0, 0.4 - 2 * +figbound[0], 1.0])
+            ax_z = fig.add_axes([0.4 + figbound[0], 0.0, 0.4 - 2 * +figbound[0], 1.0])
+            ax_l = fig.add_axes(
+                [0.8 + figbound[0], 0.0, 0.2 - 2 * +figbound[0], 1.0],
+                # projection="3d",
+                # elev=el,
+                # azim=az,
+            )
+        elif plot_layout == 1:
+            ax_x = fig.add_axes([0.0, 0.0 + figbound[0], 1.0, 0.4 - 2 * +figbound[0]])
+            ax_z = fig.add_axes([0.0, 0.4 + figbound[0], 1.0, 0.4 - 2 * +figbound[0]])
+            ax_l = fig.add_axes(
+                [0.0, 0.8 + figbound[0], 1.0, 0.2 - 2 * +figbound[0]],
+                # projection="3d",
+                # elev=el,
+                # azim=az,
+            )
+    else:
+        fig = figax[0]
+        ax_x = figax[1]
+        ax_z = figax[2]
+        ax_l = figax[3]
+
+    if crop_range is not None:
+        rgb_x = rgb_x[crop_range[0] : crop_range[1], crop_range[2] : crop_range[3]]
+        rgb_z = rgb_z[crop_range[0] : crop_range[1], crop_range[2] : crop_range[3]]
     # orientation images
     if self.pymatgen_available:
         ax_x.imshow(rgb_x)
@@ -1615,391 +1643,359 @@ def plot_orientation_maps(
 
     # Legend
     if show_legend:
-        ax_l.imshow(rgb_leg)
+        if legend_hist is True:
+            A = np.array([0, 0])
+            B = np.array([1, 0])
+            C = np.array([0.5, np.sqrt(3) / 2])
+            vertices = np.array([A, B, C])
+            ax_l.scatter(vertices[:, 0], vertices[:, 1], c=color_basis)
+            points = []
+            colors = []
+
+            for a0 in range(basis_z_scale.shape[0]):
+                for a1 in range(basis_z_scale.shape[1]):
+                    if mask[a0, a1] > 0:
+                        w = basis_z_scale_binary[a0, a1]
+                        rgb_z = (
+                            w * color_basis[0, :]
+                            + w * color_basis[1, :]
+                            + w * color_basis[2, :]
+                        )
+                        colors.append(rgb_z)
+                        w = w / np.sum(w)
+                        pt = w[0] * A + w[1] * B + w[2] * C
+                        points.append(pt)
 
-        # Add text labels
-        text_scale_pos = 0.1
-        text_params = {
-            "va": "center",
-            "family": "sans-serif",
-            "fontweight": "normal",
-            "color": "k",
-            "size": 12,
-        }
-        format_labels = "{0:.2g}"
+            points = np.asarray(points)
+            colors = np.asarray(colors)
 
-        bound = num_points * 0.25
-        shift = num_points * 0.10
-        if self.pointgroup.get_crystal_system() == "monoclinic":
-            p0 = np.array((1, 1)) * num_points
-            p1 = np.array((0, 1)) * num_points
-            p2 = np.array((1, 2)) * num_points
-            p3 = np.array((1, 0)) * num_points
-        else:
-            p0 = np.array((1, 0)) * num_points
-            p1 = np.array((1, 1)) * num_points
-            p2 = (
-                np.array((1 - np.cos(np.pi / 2 - tspan), np.sin(np.pi / 2 - tspan)))
-                * num_points
+            points_and_colors = np.hstack([points, colors])
+            points_and_colors_unique, counts = np.unique(
+                points_and_colors, axis=0, return_counts=True
             )
 
-        if self.pointgroup.get_crystal_system() == "monoclinic":
-            v = np.double(self.orientation_zone_axis_range[0, :]).copy()
-            v /= np.max(np.abs(v))
-            v = np.round(v, 2)
+            ax_l.scatter(
+                points_and_colors_unique[:, 0],
+                points_and_colors_unique[:, 1],
+                c=points_and_colors_unique[:, 2:],
+                s=counts * scale_hist,
+            )
+
+            v0 = self.orientation_zone_axis_range[0, :].copy()
+            v0 /= np.max(np.abs(v0))
+            v0 = np.round(v0, 2)
+
+            v1 = self.orientation_zone_axis_range[1, :].copy()
+            v1 /= np.max(np.abs(v1))
+            v1 = np.round(v1, 2)
+
+            v2 = self.orientation_zone_axis_range[2, :].copy()
+            v2 /= np.max(np.abs(v2))
+            v2 = np.round(v2, 2)
+
+            text_params = {
+                "va": "center",
+                "family": "sans-serif",
+                "fontweight": "normal",
+                "color": "k",
+                "size": 12,
+            }
+            shift = 0.15
             ax_l.text(
-                p0[1],
-                p0[0] + shift,
-                "["
-                + format_labels.format(v[0])
-                + " "
-                + format_labels.format(v[1])
-                + " "
-                + format_labels.format(v[2])
-                + "]",
-                None,
-                zorder=11,
+                A[0],
+                A[1] - shift,
+                v0,
                 ha="center",
                 **text_params,
             )
-            v = np.double(self.orientation_zone_axis_range[1, :]).copy()
-            v /= np.max(np.abs(v))
-            v = np.round(v, 2)
+
             ax_l.text(
-                p1[1],
-                p1[0] - shift,
-                "["
-                + format_labels.format(v[0])
-                + " "
-                + format_labels.format(v[1])
-                + " "
-                + format_labels.format(v[2])
-                + "]",
-                None,
-                zorder=11,
+                B[0],
+                B[1] - shift,
+                v1,
                 ha="center",
                 **text_params,
             )
-            v = np.double(self.orientation_zone_axis_range[2, :]).copy()
-            v /= np.max(np.abs(v))
-            v = np.round(v, 2)
+
             ax_l.text(
-                p2[1],
-                p2[0] + shift,
-                "["
-                + format_labels.format(v[0])
-                + " "
-                + format_labels.format(v[1])
-                + " "
-                + format_labels.format(v[2])
-                + "]",
-                None,
-                zorder=11,
+                C[0] + shift * 2,
+                C[1],
+                v2,
                 ha="center",
                 **text_params,
             )
-            v = -1 * np.double(self.orientation_zone_axis_range[2, :].copy()) + 0
-            v /= np.max(np.abs(v))
-            v = np.round(v, 2)
+
             ax_l.text(
-                p3[1],
-                p3[0] + shift,
-                "["
-                + format_labels.format(v[0])
-                + " "
-                + format_labels.format(v[1])
-                + " "
-                + format_labels.format(v[2])
-                + "]",
-                None,
-                zorder=11,
+                C[0],
+                C[1] + 0.1,
+                "out-of-plane orientation",
                 ha="center",
                 **text_params,
             )
 
-        else:
-            v = self.orientation_zone_axis_range[0, :].copy()
-            v /= np.max(np.abs(v))
-            v = np.round(v, 2)
+            A = np.array([0, 1.5])
+            B = np.array([1, 1.5])
+            C = np.array([0.5, np.sqrt(3) / 2 + 1.5])
+            vertices = np.array([A, B, C])
+            points = []
+            colors = []
+            ax_l.scatter(vertices[:, 0], vertices[:, 1], c=color_basis)
+
+            for a0 in range(basis_x_scale_binary.shape[0]):
+                for a1 in range(basis_x_scale_binary.shape[1]):
+                    if mask[a0, a1] > 0:
+                        w = basis_x_scale_binary[a0, a1]
+                        rgb_z = (
+                            w * color_basis[0, :]
+                            + w * color_basis[1, :]
+                            + w * color_basis[2, :]
+                        )
+                        colors.append(rgb_z)
+                        w = w / np.sum(w)
+                        pt = w[0] * A + w[1] * B + w[2] * C
+                        points.append(pt)
+
+            points = np.asarray(points)
+            colors = np.asarray(colors)
+
+            points_and_colors = np.hstack([points, colors])
+            points_and_colors_unique, counts = np.unique(
+                points_and_colors, axis=0, return_counts=True
+            )
+
+            ax_l.scatter(
+                points_and_colors_unique[:, 0],
+                points_and_colors_unique[:, 1],
+                c=points_and_colors_unique[:, 2:],
+                s=counts * scale_hist,
+            )
+
+            v0 = self.orientation_zone_axis_range[0, :].copy()
+            v0 /= np.max(np.abs(v0))
+            v0 = np.round(v0, 2)
+
+            v1 = self.orientation_zone_axis_range[1, :].copy()
+            v1 /= np.max(np.abs(v1))
+            v1 = np.round(v1, 2)
+
+            v2 = self.orientation_zone_axis_range[2, :].copy()
+            v2 /= np.max(np.abs(v2))
+            v2 = np.round(v2, 2)
+
+            text_params = {
+                "va": "center",
+                "family": "sans-serif",
+                "fontweight": "normal",
+                "color": "k",
+                "size": 12,
+            }
             ax_l.text(
-                p0[1],
-                p0[0] + shift,
-                "["
-                + format_labels.format(v[0])
-                + " "
-                + format_labels.format(v[1])
-                + " "
-                + format_labels.format(v[2])
-                + "]",
-                None,
-                zorder=11,
+                A[0],
+                A[1] - shift,
+                v0,
                 ha="center",
                 **text_params,
             )
-            v = self.orientation_zone_axis_range[1, :].copy()
-            v /= np.max(np.abs(v))
-            v = np.round(v, 2)
+
             ax_l.text(
-                p1[1],
-                p1[0] + shift,
-                "["
-                + format_labels.format(v[0])
-                + " "
-                + format_labels.format(v[1])
-                + " "
-                + format_labels.format(v[2])
-                + "]",
-                None,
-                zorder=11,
+                B[0],
+                B[1] - shift,
+                v1,
                 ha="center",
                 **text_params,
             )
-            v = self.orientation_zone_axis_range[2, :].copy()
-            v /= np.max(np.abs(v))
-            v = np.round(v, 2)
+
             ax_l.text(
-                p2[1],
-                p2[0] - shift,
-                "["
-                + format_labels.format(v[0])
-                + " "
-                + format_labels.format(v[1])
-                + " "
-                + format_labels.format(v[2])
-                + "]",
-                None,
-                zorder=11,
+                C[0] + shift * 2,
+                C[1],
+                v2,
                 ha="center",
                 **text_params,
             )
 
-        if self.pointgroup.get_crystal_system() == "monoclinic":
-            ax_l.set_xlim((-bound, num_points * 2 - 1 + bound))
+            ax_l.text(
+                C[0],
+                C[1] + 0.1,
+                "in-plane orientation",
+                ha="center",
+                **text_params,
+            )
+
+            ax_l.set_aspect("equal")
+            ax_l.axis("off")
+
+            ax_l.set_xlim([-0.3, 1.3])
+            # ax_l.set_ylim([-0.3, C[1] + 0.3])
+
         else:
-            ax_l.set_xlim((-bound, num_points + bound))
-        ax_l.set_ylim((num_points + bound, -bound))
-        ax_l.axis("off")
-
-        # if np.abs(self.cell[5] - 120.0) > 1e-6:
-        #     ax_l.text(
-        #         self.orientation_vecs[inds_legend[0], 1] + vec[1] * text_scale_pos,
-        #         self.orientation_vecs[inds_legend[0], 0] + vec[0] * text_scale_pos,
-        #         self.orientation_vecs[inds_legend[0], 2] + vec[2] * text_scale_pos,
-        #         "["
-        #         + format_labels.format(label_0[0])
-        #         + " "
-        #         + format_labels.format(label_0[1])
-        #         + " "
-        #         + format_labels.format(label_0[2])
-        #         + "]",
-        #         None,
-        #         zorder=11,
-        #         ha="center",
-        #         **text_params,
-        #     )
-        # else:
-        #     ax_l.text(
-        #         self.orientation_vecs[inds_legend[0], 1] + vec[1] * text_scale_pos,
-        #         self.orientation_vecs[inds_legend[0], 0] + vec[0] * text_scale_pos,
-        #         self.orientation_vecs[inds_legend[0], 2] + vec[2] * text_scale_pos,
-        #         "["
-        #         + format_labels.format(label_0[0])
-        #         + " "
-        #         + format_labels.format(label_0[1])
-        #         + " "
-        #         + format_labels.format(label_0[2])
-        #         + " "
-        #         + format_labels.format(label_0[3])
-        #         + "]",
-        #         None,
-        #         zorder=11,
-        #         ha="center",
-        #         **text_params,
-        #     )
+            ax_l.imshow(rgb_leg)
+
+            # Add text labels
+            text_scale_pos = 0.1
+            text_params = {
+                "va": "center",
+                "family": "sans-serif",
+                "fontweight": "normal",
+                "color": "k",
+                "size": 12,
+            }
+            format_labels = "{0:.2g}"
+
+            bound = num_points * 0.25
+            shift = num_points * 0.10
+            if self.pointgroup.get_crystal_system() == "monoclinic":
+                p0 = np.array((1, 1)) * num_points
+                p1 = np.array((0, 1)) * num_points
+                p2 = np.array((1, 2)) * num_points
+                p3 = np.array((1, 0)) * num_points
+            else:
+                p0 = np.array((1, 0)) * num_points
+                p1 = np.array((1, 1)) * num_points
+                p2 = (
+                    np.array((1 - np.cos(np.pi / 2 - tspan), np.sin(np.pi / 2 - tspan)))
+                    * num_points
+                )
 
-    # # Triangulate faces
-    # p = self.orientation_vecs[:, (1, 0, 2)]
-    # tri = mtri.Triangulation(
-    #     self.orientation_inds[:, 1] - self.orientation_inds[:, 0] * 1e-3,
-    #     self.orientation_inds[:, 0] - self.orientation_inds[:, 1] * 1e-3,
-    # )
-    # convert rgb values of pixels to faces
-    # rgb_faces = (
-    #     rgb_legend[tri.triangles[:, 0], :]
-    #     + rgb_legend[tri.triangles[:, 1], :]
-    #     + rgb_legend[tri.triangles[:, 2], :]
-    # ) / 3
-    # Add triangulated surface plot to axes
-    # pc = art3d.Poly3DCollection(
-    #     p[tri.triangles],
-    #     facecolors=rgb_faces,
-    #     alpha=1,
-    # )
-    # pc.set_antialiased(False)
-    # ax_l.add_collection(pc)
-
-    # if plot_limit is None:
-    #     plot_limit = np.array(
-    #         [
-    #             [np.min(p[:, 0]), np.min(p[:, 1]), np.min(p[:, 2])],
-    #             [np.max(p[:, 0]), np.max(p[:, 1]), np.max(p[:, 2])],
-    #         ]
-    #     )
-    #     # plot_limit = (plot_limit - np.mean(plot_limit, axis=0)) * 1.5 + np.mean(
-    #     #     plot_limit, axis=0
-    #     # )
-    #     plot_limit[:, 0] = (
-    #         plot_limit[:, 0] - np.mean(plot_limit[:, 0])
-    #     ) * 1.5 + np.mean(plot_limit[:, 0])
-    #     plot_limit[:, 1] = (
-    #         plot_limit[:, 2] - np.mean(plot_limit[:, 1])
-    #     ) * 1.5 + np.mean(plot_limit[:, 1])
-    #     plot_limit[:, 2] = (
-    #         plot_limit[:, 1] - np.mean(plot_limit[:, 2])
-    #     ) * 1.1 + np.mean(plot_limit[:, 2])
-
-    # # ax_l.view_init(elev=el, azim=az)
-    # # Appearance
-    # ax_l.invert_yaxis()
-    # if swap_axes_xy_limits:
-    #     ax_l.axes.set_xlim3d(left=plot_limit[0, 0], right=plot_limit[1, 0])
-    #     ax_l.axes.set_ylim3d(bottom=plot_limit[0, 1], top=plot_limit[1, 1])
-    #     ax_l.axes.set_zlim3d(bottom=plot_limit[0, 2], top=plot_limit[1, 2])
-    # else:
-    #     ax_l.axes.set_xlim3d(left=plot_limit[0, 1], right=plot_limit[1, 1])
-    #     ax_l.axes.set_ylim3d(bottom=plot_limit[0, 0], top=plot_limit[1, 0])
-    #     ax_l.axes.set_zlim3d(bottom=plot_limit[0, 2], top=plot_limit[1, 2])
-    # axisEqual3D(ax_l)
-    # if camera_dist is not None:
-    #     ax_l.dist = camera_dist
-    # ax_l.axis("off")
-
-    # # Add text labels
-    # text_scale_pos = 0.1
-    # text_params = {
-    #     "va": "center",
-    #     "family": "sans-serif",
-    #     "fontweight": "normal",
-    #     "color": "k",
-    #     "size": 14,
-    # }
-    # format_labels = "{0:.2g}"
-    # vec = self.orientation_vecs[inds_legend[0], :] - cam_dir
-    # vec = vec / np.linalg.norm(vec)
-    # if np.abs(self.cell[5] - 120.0) > 1e-6:
-    #     ax_l.text(
-    #         self.orientation_vecs[inds_legend[0], 1] + vec[1] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[0], 0] + vec[0] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[0], 2] + vec[2] * text_scale_pos,
-    #         "["
-    #         + format_labels.format(label_0[0])
-    #         + " "
-    #         + format_labels.format(label_0[1])
-    #         + " "
-    #         + format_labels.format(label_0[2])
-    #         + "]",
-    #         None,
-    #         zorder=11,
-    #         ha="center",
-    #         **text_params,
-    #     )
-    # else:
-    #     ax_l.text(
-    #         self.orientation_vecs[inds_legend[0], 1] + vec[1] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[0], 0] + vec[0] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[0], 2] + vec[2] * text_scale_pos,
-    #         "["
-    #         + format_labels.format(label_0[0])
-    #         + " "
-    #         + format_labels.format(label_0[1])
-    #         + " "
-    #         + format_labels.format(label_0[2])
-    #         + " "
-    #         + format_labels.format(label_0[3])
-    #         + "]",
-    #         None,
-    #         zorder=11,
-    #         ha="center",
-    #         **text_params,
-    #     )
-    # vec = self.orientation_vecs[inds_legend[1], :] - cam_dir
-    # vec = vec / np.linalg.norm(vec)
-    # if np.abs(self.cell[5] - 120.0) > 1e-6:
-    #     ax_l.text(
-    #         self.orientation_vecs[inds_legend[1], 1] + vec[1] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[1], 0] + vec[0] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[1], 2] + vec[2] * text_scale_pos,
-    #         "["
-    #         + format_labels.format(label_1[0])
-    #         + " "
-    #         + format_labels.format(label_1[1])
-    #         + " "
-    #         + format_labels.format(label_1[2])
-    #         + "]",
-    #         None,
-    #         zorder=12,
-    #         ha=ha_1,
-    #         **text_params,
-    #     )
-    # else:
-    #     ax_l.text(
-    #         self.orientation_vecs[inds_legend[1], 1] + vec[1] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[1], 0] + vec[0] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[1], 2] + vec[2] * text_scale_pos,
-    #         "["
-    #         + format_labels.format(label_1[0])
-    #         + " "
-    #         + format_labels.format(label_1[1])
-    #         + " "
-    #         + format_labels.format(label_1[2])
-    #         + " "
-    #         + format_labels.format(label_1[3])
-    #         + "]",
-    #         None,
-    #         zorder=12,
-    #         ha=ha_1,
-    #         **text_params,
-    #     )
-    # vec = self.orientation_vecs[inds_legend[2], :] - cam_dir
-    # vec = vec / np.linalg.norm(vec)
-    # if np.abs(self.cell[5] - 120.0) > 1e-6:
-    #     ax_l.text(
-    #         self.orientation_vecs[inds_legend[2], 1] + vec[1] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[2], 0] + vec[0] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[2], 2] + vec[2] * text_scale_pos,
-    #         "["
-    #         + format_labels.format(label_2[0])
-    #         + " "
-    #         + format_labels.format(label_2[1])
-    #         + " "
-    #         + format_labels.format(label_2[2])
-    #         + "]",
-    #         None,
-    #         zorder=13,
-    #         ha=ha_2,
-    #         **text_params,
-    #     )
-    # else:
-    #     ax_l.text(
-    #         self.orientation_vecs[inds_legend[2], 1] + vec[1] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[2], 0] + vec[0] * text_scale_pos,
-    #         self.orientation_vecs[inds_legend[2], 2] + vec[2] * text_scale_pos,
-    #         "["
-    #         + format_labels.format(label_2[0])
-    #         + " "
-    #         + format_labels.format(label_2[1])
-    #         + " "
-    #         + format_labels.format(label_2[2])
-    #         + " "
-    #         + format_labels.format(label_2[3])
-    #         + "]",
-    #         None,
-    #         zorder=13,
-    #         ha=ha_2,
-    #         **text_params,
-    #     )
+            if self.pointgroup.get_crystal_system() == "monoclinic":
+                v = np.double(self.orientation_zone_axis_range[0, :]).copy()
+                v /= np.max(np.abs(v))
+                v = np.round(v, 2)
+                ax_l.text(
+                    p0[1],
+                    p0[0] + shift,
+                    "["
+                    + format_labels.format(v[0])
+                    + " "
+                    + format_labels.format(v[1])
+                    + " "
+                    + format_labels.format(v[2])
+                    + "]",
+                    None,
+                    zorder=11,
+                    ha="center",
+                    **text_params,
+                )
+                v = np.double(self.orientation_zone_axis_range[1, :]).copy()
+                v /= np.max(np.abs(v))
+                v = np.round(v, 2)
+                ax_l.text(
+                    p1[1],
+                    p1[0] - shift,
+                    "["
+                    + format_labels.format(v[0])
+                    + " "
+                    + format_labels.format(v[1])
+                    + " "
+                    + format_labels.format(v[2])
+                    + "]",
+                    None,
+                    zorder=11,
+                    ha="center",
+                    **text_params,
+                )
+                v = np.double(self.orientation_zone_axis_range[2, :]).copy()
+                v /= np.max(np.abs(v))
+                v = np.round(v, 2)
+                ax_l.text(
+                    p2[1],
+                    p2[0] + shift,
+                    "["
+                    + format_labels.format(v[0])
+                    + " "
+                    + format_labels.format(v[1])
+                    + " "
+                    + format_labels.format(v[2])
+                    + "]",
+                    None,
+                    zorder=11,
+                    ha="center",
+                    **text_params,
+                )
+                v = -1 * np.double(self.orientation_zone_axis_range[2, :].copy()) + 0
+                v /= np.max(np.abs(v))
+                v = np.round(v, 2)
+                ax_l.text(
+                    p3[1],
+                    p3[0] + shift,
+                    "["
+                    + format_labels.format(v[0])
+                    + " "
+                    + format_labels.format(v[1])
+                    + " "
+                    + format_labels.format(v[2])
+                    + "]",
+                    None,
+                    zorder=11,
+                    ha="center",
+                    **text_params,
+                )
+                plt.tight_layout()
+            else:
+                v = self.orientation_zone_axis_range[0, :].copy()
+                v /= np.max(np.abs(v))
+                v = np.round(v, 2)
+                ax_l.text(
+                    p0[1],
+                    p0[0] + shift,
+                    "["
+                    + format_labels.format(v[0])
+                    + " "
+                    + format_labels.format(v[1])
+                    + " "
+                    + format_labels.format(v[2])
+                    + "]",
+                    None,
+                    zorder=11,
+                    ha="center",
+                    **text_params,
+                )
+                v = self.orientation_zone_axis_range[1, :].copy()
+                v /= np.max(np.abs(v))
+                v = np.round(v, 2)
+                ax_l.text(
+                    p1[1],
+                    p1[0] + shift,
+                    "["
+                    + format_labels.format(v[0])
+                    + " "
+                    + format_labels.format(v[1])
+                    + " "
+                    + format_labels.format(v[2])
+                    + "]",
+                    None,
+                    zorder=11,
+                    ha="center",
+                    **text_params,
+                )
+                v = self.orientation_zone_axis_range[2, :].copy()
+                v /= np.max(np.abs(v))
+                v = np.round(v, 2)
+                ax_l.text(
+                    p2[1],
+                    p2[0] - shift,
+                    "["
+                    + format_labels.format(v[0])
+                    + " "
+                    + format_labels.format(v[1])
+                    + " "
+                    + format_labels.format(v[2])
+                    + "]",
+                    None,
+                    zorder=11,
+                    ha="center",
+                    **text_params,
+                )
 
-    images_orientation = np.zeros((orientation_map.num_x, orientation_map.num_y, 3, 2))
+            if self.pointgroup.get_crystal_system() == "monoclinic":
+                ax_l.set_xlim((-bound, num_points * 2 - 1 + bound))
+            else:
+                ax_l.set_xlim((-bound, num_points + bound))
+            ax_l.set_ylim((num_points + bound, -bound))
+            ax_l.axis("off")
+    else:
+        ax_l.remove()
+
+    images_orientation = np.zeros((rgb_x.shape[0], rgb_x.shape[1], 3, 2))
     if self.pymatgen_available:
         images_orientation[:, :, :, 0] = rgb_x
     images_orientation[:, :, :, 1] = rgb_z
diff --git a/py4DSTEM/process/utils/cluster.py b/py4DSTEM/process/utils/cluster.py
index aed0d213c..8916bccf1 100644
--- a/py4DSTEM/process/utils/cluster.py
+++ b/py4DSTEM/process/utils/cluster.py
@@ -8,28 +8,55 @@
 
 class Cluster:
     """
-    Clustering 4D data
-
+    Class for clustering data in 4D-STEM DataCube based on
+    similarity of neighboring diffraction patterns.
     """
 
     def __init__(
         self,
         datacube,
+        r_space_mask=None,
     ):
         """
-        Args:
-            datacube (py4DSTEM.DataCube):            4D-STEM data
-
-
+        Parameters
+        ----------
+        datacube: DataCube
+            4D-STEM data
+        r_space_mask: np.ndarray
+            Mask in real space to apply background thresholding on the similarity array.
         """
-
         self.datacube = datacube
+        self.r_space_mask = r_space_mask
+        self.similarity = None
+        self.similarity_raw = None
+
+    def _apply_bg_mask(self, similarity):
+        if self.r_space_mask is None:
+            return similarity
+        return similarity * self.r_space_mask[..., None]
 
     def find_similarity(
-        self,
-        mask=None,  # by default
+        self, q_space_mask=None, smooth_sigma=0, return_similarity=False
     ):
-        # Which neighbors to search
+        """
+        Find similarity to neighboring pixels
+
+        Parameters
+        ----------
+        q_space_mask : np.ndarray, optional
+            boolean q_space_mask to apply on the diffraction patterns
+        smooth_sigma : float, optional
+            sigma for Gaussian smoothing of the diffraction patterns
+            before calculating similarity
+        return_similarity : bool, optinal
+            if True, return the similarity array
+
+        Returns
+        --------
+        similarity: np.ndarray
+            similarity scores for each pixel
+        """
+        # List of neighbors to search
         # (-1,-1) will be equivalent to (1,1)
         self.dxy = np.array(
             (
@@ -54,10 +81,17 @@ def find_similarity(
             range(self.datacube.shape[0]),
             range(self.datacube.shape[1]),
         ):
-            if mask is None:
-                diff_ref = self.datacube[rx, ry]
-            else:
-                diff_ref = self.datacube[rx, ry][mask]
+            diff_ref = self.datacube[rx, ry].copy().astype("float")
+            diff_ref -= diff_ref.mean()
+
+            if smooth_sigma > 0:
+                diff_ref = gaussian_filter(diff_ref, smooth_sigma)
+
+            if q_space_mask is not None:
+                diff_ref = diff_ref[q_space_mask]
+
+            norm_diff_ref = np.sqrt(np.sum(diff_ref * diff_ref))
+            # diff_ref_mean = np.mean(diff_ref)
 
             # loop over neighbors
             for ind in range(self.dxy.shape[0]):
@@ -69,26 +103,28 @@ def find_similarity(
                     and x_ind < self.datacube.shape[0]
                     and y_ind < self.datacube.shape[1]
                 ):
+                    diff = self.datacube[x_ind, y_ind].copy().astype("float")
+                    diff -= diff.mean()
 
-                    if mask is None:
-                        diff = self.datacube[x_ind, y_ind]
-                    else:
-                        diff = self.datacube[x_ind, y_ind][mask]
+                    if smooth_sigma > 0:
+                        diff = gaussian_filter(diff, smooth_sigma)
 
-                    # # image self.similarity with mean abs difference
-                    # self.similarity[rx,ry,ind] = np.mean(
-                    #     np.abs(
-                    #         diff - diff_ref
-                    #     )
-                    # )
+                    if q_space_mask is not None:
+                        diff = diff[q_space_mask]
 
-                    # image self.similarity with normalized corr: cosine self.similarity?
+                    # image self.similarity with normalized cosine correlation
                     self.similarity[rx, ry, ind] = (
                         np.sum(diff * diff_ref)
                         / np.sqrt(np.sum(diff * diff))
-                        / np.sqrt(np.sum(diff_ref * diff_ref))
+                        / norm_diff_ref
                     )
 
+        self.similarity_raw = self.similarity.copy()
+        self.similarity = self._apply_bg_mask(self.similarity)
+
+        if return_similarity:
+            return self.similarity
+
     # Create a function to map cluster index to color
     def get_color(self, cluster_index):
         colors = [
@@ -108,31 +144,34 @@ def get_color(self, cluster_index):
     # Find the pixel with the highest self.similarity and start the clustering from there
     def indexing_clusters_all(
         self,
-        mask,
         threshold,
     ):
+        """
+        Index all pixsl in a cluster
 
-        self.dxy = np.array(
-            (
-                (-1, -1),
-                (-1, 0),
-                (-1, 1),
-                (0, -1),
-                (1, 1),
-                (1, 0),
-                (1, -1),
-                (0, 1),
-            )
-        )
+        Parameters
+        ----------
+        threshold: float
+            similarity score threshold to consider pixels as part
+            of the same cluster
+        """
 
         sim_averaged = np.mean(self.similarity, axis=2)
 
+        # Assigning the background as 'counted'
+        if self.r_space_mask is not None:
+            sim_averaged[~self.r_space_mask] = -1.0
+
         # color the pixels with the cluster index
-        # map_cluster = np.zeros((sim_averaged.shape[0],sim_averaged.shape[1]))
-        self.cluster_map = np.zeros(
+        self.cluster_map = -1 * np.ones(
+            (sim_averaged.shape[0], sim_averaged.shape[1]), dtype=np.float64
+        )
+        self.cluster_map_rgb = np.zeros(
             (sim_averaged.shape[0], sim_averaged.shape[1], 4), dtype=np.float64
         )
 
+        self.cluster_map_rgb[..., 3] = 1.0  # start as opaque black
+
         # store arrays of cluster_indices in a list
         self.cluster_list = []
 
@@ -147,16 +186,21 @@ def indexing_clusters_all(
             # finding the pixel that has the highest self.similarity among the pixel that hasn't been clustered yet
             # this will be the 'starting pixel' of a new cluster
             rx0, ry0 = np.unravel_index(sim_averaged.argmax(), sim_averaged.shape)
-            # print(rx0, ry0)
+
+            # Guarding to check if the seed is background
+            if self.r_space_mask is not None and not self.r_space_mask[rx0, ry0]:
+                sim_averaged[rx0, ry0] = -1  # mark processed so we don't pick it again
+                continue
 
             cluster_indices = np.empty((0, 2))
             cluster_indices = (np.append(cluster_indices, [[rx0, ry0]], axis=0)).astype(
                 np.int32
             )
 
-            # map_cluster[rx0, ry0] = cluster_count_ind+1
+            self.cluster_map[rx0, ry0] = cluster_count_ind
+
             color = self.get_color(cluster_count_ind + 1)
-            self.cluster_map[rx0, ry0] = plt.cm.colors.to_rgba(color)
+            self.cluster_map_rgb[rx0, ry0] = plt.cm.colors.to_rgba(color)
 
             # Clustering: one cluster per while loop(until it breaks)
             # Marching algorithm: find a new position and search the nearest neighbor
@@ -171,52 +215,76 @@ def indexing_clusters_all(
                         # counter to check if pixel in the cluster are checked for NN
                         counting_added_pixel += 1
 
-                        # set to -1 as its NN will be checked
+                        # set to -1 since now its NN will be checked
                         sim_averaged[rx0, ry0] = -1
 
                         for ind in range(self.dxy.shape[0]):
                             x_ind = rx0 + self.dxy[ind, 0]
                             y_ind = ry0 + self.dxy[ind, 1]
 
-                            # add if the neighbor is similar, but don't add if the neighbor is already in a cluster
-                            if self.similarity[
-                                rx0, ry0, ind
-                            ] > threshold and np.array_equal(
-                                self.cluster_map[x_ind, y_ind], [0, 0, 0, 0]
+                            if (
+                                x_ind > 1
+                                and y_ind > 1
+                                and x_ind < self.similarity.shape[0] - 2
+                                and y_ind < self.similarity.shape[1] - 2
                             ):
 
-                                cluster_indices = np.append(
-                                    cluster_indices, [[x_ind, y_ind]], axis=0
-                                )
-                                # self.cluster_map[x_ind, y_ind] = cluster_count_ind+1
-                                color = self.get_color(cluster_count_ind + 1)
-                                self.cluster_map[x_ind, y_ind] = plt.cm.colors.to_rgba(
-                                    color
+                                r_ok = (
+                                    True
+                                    if self.r_space_mask is None
+                                    else bool(self.r_space_mask[x_ind, y_ind])
                                 )
 
+                                # add if the neighbor is similar, but don't add if the neighbor is already in a cluster
+                                if (
+                                    self.similarity[rx0, ry0, ind] >= threshold
+                                    and self.cluster_map[x_ind, y_ind] == -1
+                                    and r_ok
+                                ):
+
+                                    cluster_indices = np.append(
+                                        cluster_indices, [[x_ind, y_ind]], axis=0
+                                    )
+
+                                    self.cluster_map[x_ind, y_ind] = cluster_count_ind
+
+                                    color = self.get_color(cluster_count_ind + 1)
+                                    self.cluster_map_rgb[x_ind, y_ind] = (
+                                        plt.cm.colors.to_rgba(color)
+                                    )
+
                 # if no new pixel is checked for NN then break
                 if counting_added_pixel == 0:
                     break
 
-            # single pixel cluster
-            if cluster_indices.shape[0] == 1:
-                self.cluster_map[cluster_indices[0, 0], cluster_indices[0, 1]] = [
-                    0,
-                    0,
-                    0,
-                    1,
-                ]
-
             self.cluster_list.append(cluster_indices)
             cluster_count_ind += 1
 
-        # return cluster_count_ind, self.cluster_list, map_cluster, sim_averaged
-
     def create_cluster_cube(
         self,
         min_cluster_size,
         return_cluster_datacube=False,
     ):
+        """
+        Create dataset (N, 1, qx, qy), where N is the number of clusters
+        that contains diffraction patterns that are averaged across pixels
+        in each cluster
+
+        Parameters
+        ----------
+        min_cluster_size: int
+            minimum size for a clsuter to be included in dataset
+        return_cluster_datacube: bool
+            if True, returns clustered dataset and list of indicies
+            of clusters
+
+        Returns
+        --------
+        cluster_cube: np.ndarray
+            dataset with clsutered diffraction patterns
+        filtered_cluster_list: list
+            list of indicies in real space of each pixel of each cluster
+        """
 
         self.filtered_cluster_list = [
             arr for arr in self.cluster_list if arr.shape[0] >= min_cluster_size
diff --git a/test/gettestdata.py b/test/gettestdata.py
index 0cb6cb964..376b5737b 100644
--- a/test/gettestdata.py
+++ b/test/gettestdata.py
@@ -8,7 +8,6 @@
 from py4DSTEM import _TESTPATH as testpath
 from py4DSTEM.io import gdrive_download as download
 
-
 # Make the argument parser
 parser = argparse.ArgumentParser(
     description="A command line tool for downloading data to run the py4DSTEM test suite"
diff --git a/test/test_native_io/test_realslice_read.py b/test/test_native_io/test_realslice_read.py
index eeee5217d..18727129c 100644
--- a/test/test_native_io/test_realslice_read.py
+++ b/test/test_native_io/test_realslice_read.py
@@ -4,7 +4,6 @@
 import py4DSTEM
 from os.path import join
 
-
 # Set filepaths
 filepath = join(py4DSTEM._TESTPATH, "test_io/test_realslice_io.h5")
 
diff --git a/test/test_native_io/test_v0_13.py b/test/test_native_io/test_v0_13.py
index e1d91bba4..c9e851f72 100644
--- a/test/test_native_io/test_v0_13.py
+++ b/test/test_native_io/test_v0_13.py
@@ -1,7 +1,6 @@
 from py4DSTEM import read, print_h5_tree, _TESTPATH
 from os.path import join
 
-
 # Set filepaths
 filepath = join(_TESTPATH, "test_io/legacy_v0.13.h5")
 
diff --git a/test/test_native_io/test_v0_14.py b/test/test_native_io/test_v0_14.py
index c71638c8e..77ca8ca7e 100644
--- a/test/test_native_io/test_v0_14.py
+++ b/test/test_native_io/test_v0_14.py
@@ -1,7 +1,6 @@
 import py4DSTEM
 from os.path import join, exists
 
-
 path = join(py4DSTEM._TESTPATH, "test_io/legacy_v0.14.h5")
 
 
diff --git a/test/test_nonnative_io/test_arina.py b/test/test_nonnative_io/test_arina.py
index c02964cf8..556fd1a13 100644
--- a/test/test_nonnative_io/test_arina.py
+++ b/test/test_nonnative_io/test_arina.py
@@ -2,7 +2,6 @@
 import emdfile
 from os.path import join
 
-
 # Set filepaths
 filepath = join(py4DSTEM._TESTPATH, "test_arina/STO_STEM_bench_20us_master.h5")
 
diff --git a/test/test_nonnative_io/test_dm.py b/test/test_nonnative_io/test_dm.py
index ee6f1b2eb..82964d547 100644
--- a/test/test_nonnative_io/test_dm.py
+++ b/test/test_nonnative_io/test_dm.py
@@ -2,7 +2,6 @@
 import emdfile
 from os.path import join
 
-
 # Set filepaths
 filepath_dm4_datacube = join(py4DSTEM._TESTPATH, "small_datacube.dm4")
 filepath_dm3_3Dstack = join(py4DSTEM._TESTPATH, "test_io/small_dm3_3Dstack.dm3")
diff --git a/test/test_strain.py b/test/test_strain.py
index d309c5ffb..1144d1c17 100644
--- a/test/test_strain.py
+++ b/test/test_strain.py
@@ -3,7 +3,6 @@
 from os.path import join
 from numpy import zeros
 
-
 # set filepath
 path = join(py4DSTEM._TESTPATH, "strain/downsample_Si_SiGe_analysis_braggdisks_cal.h5")