Merge pull request #590 from smribet/STO_dev

Small plotting changes

Author: sezelt

py4DSTEM/datacube/virtualimage.py

@@ -681,6 +681,8 @@ def make_bragg_mask(
         origin,
         max_q,
         return_sum=True,
+        include_origin=True,
+        rotation_deg=0,
         **kwargs,
     ):
         """
@@ -696,12 +698,21 @@ def make_bragg_mask(
             return_sum (bool): if False, return a 3D array, where each
                 slice contains a single disk; if False, return a single
                 2D masks of all disks
+            include_origin (bool) : if False, removes origin disk
+            rotation_deg (float) : rotate g1 and g2 vectors
 
         Returns:
             (2 or 3D array) the mask
         """
-        nas = np.asarray
-        g1, g2, origin = nas(g1), nas(g2), nas(origin)
+        g1, g2, origin = np.asarray(g1), np.asarray(g2), np.asarray(origin)
+
+        rotation_rad = np.deg2rad(rotation_deg)
+        cost = np.cos(rotation_rad)
+        sint = np.sin(rotation_rad)
+        rotation_matrix = np.array(((cost, sint), (-sint, cost)))
+
+        g1 = np.dot(g1, rotation_matrix)
+        g2 = np.dot(g2, rotation_matrix)
 
         # Get N,M, the maximum indices to tile out to
         L1 = np.sqrt(np.sum(g1**2))
@@ -722,15 +733,18 @@ def make_bragg_mask(
         N = 0
         for h in range(-H, H + 1):
             for k in range(-K, K + 1):
-                v = h * g1 + k * g2
-                if np.sqrt(v.dot(v)) < max_q:
-                    center = origin + v
-                    mask[:, :, N] = self.make_detector(
-                        Qshape,
-                        mode="circle",
-                        geometry=(center, radius),
-                    )
-                    N += 1
+                if h == 0 and k == 0 and include_origin is False:
+                    continue
+                else:
+                    v = h * g1 + k * g2
+                    if np.sqrt(v.dot(v)) < max_q:
+                        center = origin + v
+                        mask[:, :, N] = self.make_detector(
+                            Qshape,
+                            mode="circle",
+                            geometry=(center, radius),
+                        )
+                        N += 1
 
         if return_sum:
             mask = np.sum(mask, axis=2)

py4DSTEM/process/strain/strain.py

@@ -6,14 +6,12 @@
 import matplotlib.pyplot as plt
 from matplotlib.patches import Circle
 from matplotlib.collections import PatchCollection
-from mpl_toolkits.axes_grid1 import make_axes_locatable
 import numpy as np
 from py4DSTEM import PointList, PointListArray, tqdmnd
 from py4DSTEM.braggvectors import BraggVectors
 from py4DSTEM.data import Data, RealSlice
 from py4DSTEM.preprocess.utils import get_maxima_2D
 from py4DSTEM.process.strain.latticevectors import (
-    add_indices_to_braggvectors,
     fit_lattice_vectors_all_DPs,
     get_reference_g1g2,
     get_rotated_strain_map,
@@ -25,8 +23,6 @@
     add_bragg_index_labels,
     add_pointlabels,
     add_vector,
-    ax_addaxes,
-    ax_addaxes_QtoR,
 )
 
 
@@ -583,6 +579,7 @@ def get_strain(
             amount, in degrees
         returncal : bool
             It True, returns rotated map
+        **kwargs: keywords passed to py4DSTEM show function
         """
         # confirm that the calstate hasn't changed
         assert (
@@ -681,6 +678,7 @@ def show_strain(
         layout="square",
         figsize=None,
         returnfig=False,
+        **kwargs,
     ):
         """
         Display a strain map, showing the 4 strain components
@@ -738,11 +736,12 @@ def show_strain(
             Scaling for the legend g-vectors relative to the coordinate axes
         layout : int
             Determines the layout of the grid which the strain components
-            will be plotted in.  Must be in (0,1,2).  0=(2x2), 1=(1x4), 2=(4x1).
+            will be plotted in.  Options are "square", "horizontal", "vertical."
         figsize : length 2 tuple of numbers
             Size of the figure
         returnfig : bool
             Toggles returning the figure
+        **kwargs: keywords passed to py4DSTEM show function
         """
 
         from py4DSTEM.visualize import show_strain
@@ -776,6 +775,7 @@ def show_strain(
             layout=layout,
             figsize=figsize,
             returnfig=True,
+            **kwargs,
         )
 
         # show/return

py4DSTEM/visualize/vis_special.py

@@ -926,6 +926,7 @@ def show_strain(
     layout="square",
     figsize=None,
     returnfig=False,
+    **kwargs,
 ):
     """
     Display a strain map, showing the 4 strain components
@@ -988,11 +989,12 @@ def show_strain(
         Scaling for the legend g-vectors relative to the coordinate axes
     layout : int
         Determines the layout of the grid which the strain components
-        will be plotted in.  Must be in (0,1,2).  0=(2x2), 1=(1x4), 2=(4x1).
+        will be plotted in.  Options are "square", "horizontal", "vertical."
     figsize : length 2 tuple of numbers
         Size of the figure
     returnfig : bool
         Toggles returning the figure
+    **kwargs: keywords passed to py4DSTEM show function
     """
     # Lookup table for different layouts
     assert layout in ("square", "horizontal", "vertical")
@@ -1099,6 +1101,7 @@ def show_strain(
         cmap=cmap,
         mask_color=mask_color,
         returncax=True,
+        **kwargs,
     )
     cax12 = show(
         e_yy,
@@ -1109,6 +1112,7 @@ def show_strain(
         cmap=cmap,
         mask_color=mask_color,
         returncax=True,
+        **kwargs,
     )
     cax21 = show(
         e_xy,
@@ -1119,6 +1123,7 @@ def show_strain(
         cmap=cmap,
         mask_color=mask_color,
         returncax=True,
+        **kwargs,
     )
     cax22 = show(
         theta,
@@ -1129,6 +1134,7 @@ def show_strain(
         cmap=cmap_theta,
         mask_color=mask_color,
         returncax=True,
+        **kwargs,
     )
     ax11.set_title(r"$\epsilon_{xx}$", size=titlesize)
     ax12.set_title(r"$\epsilon_{yy}$", size=titlesize)