diff --git a/py4DSTEM/datacube/virtualimage.py b/py4DSTEM/datacube/virtualimage.py index 87aeae8b1..627223d23 100644 --- a/py4DSTEM/datacube/virtualimage.py +++ b/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) diff --git a/py4DSTEM/process/strain/strain.py b/py4DSTEM/process/strain/strain.py index 621fc820e..25162180f 100644 --- a/py4DSTEM/process/strain/strain.py +++ b/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 diff --git a/py4DSTEM/visualize/vis_special.py b/py4DSTEM/visualize/vis_special.py index 125a2ce67..88b7d7815 100644 --- a/py4DSTEM/visualize/vis_special.py +++ b/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)