From 19b5d629a8d7f09d7d16f773f42decb9c69b92b4 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Tue, 7 Nov 2023 16:02:22 -0800
Subject: [PATCH 01/18] changing h,k for g1_ind, g2_ind

---
 py4DSTEM/process/strain/latticevectors.py | 54 +++++++++++------------
 1 file changed, 27 insertions(+), 27 deletions(-)

diff --git a/py4DSTEM/process/strain/latticevectors.py b/py4DSTEM/process/strain/latticevectors.py
index ba9bb4fcf..65b5f6c69 100644
--- a/py4DSTEM/process/strain/latticevectors.py
+++ b/py4DSTEM/process/strain/latticevectors.py
@@ -33,7 +33,7 @@ def index_bragg_directions(x0, y0, gx, gy, g1, g2):
             * **k**: *(ndarray of ints)* second index of the bragg directions
             * **bragg_directions**: *(PointList)* a 4-coordinate PointList with the
               indexed bragg directions; coords 'qx' and 'qy' contain bragg_x and bragg_y
-              coords 'h' and 'k' contain h and k.
+              coords 'g1_ind' and 'g2_ind' contain g1_ind and g2_ind.
     """
     # Get beta, the matrix of lattice vectors
     beta = np.array([[g1[0], g2[0]], [g1[1], g2[1]]])
@@ -45,39 +45,39 @@ def index_bragg_directions(x0, y0, gx, gy, g1, g2):
     M = lstsq(beta, alpha, rcond=None)[0].T
     M = np.round(M).astype(int)
 
-    # Get h,k
-    h = M[:, 0]
-    k = M[:, 1]
+    # Get g1_ind,g2_ind
+    g1_ind = M[:, 0]
+    g2_ind  = M[:, 1]
 
     # Store in a PointList
-    coords = [("qx", float), ("qy", float), ("h", int), ("k", int)]
+    coords = [("qx", float), ("qy", float), ("g1_ind", int), ("g2_ind", int)]
     temp_array = np.zeros([], dtype=coords)
     bragg_directions = PointList(data=temp_array)
-    bragg_directions.add_data_by_field((gx, gy, h, k))
+    bragg_directions.add_data_by_field((gx, gy, g1_ind, g2_ind))
     mask = np.zeros(bragg_directions["qx"].shape[0])
     mask[0] = 1
     bragg_directions.remove(mask)
 
-    return h, k, bragg_directions
+    return g1_ind, g2_ind, bragg_directions
 
 
 def add_indices_to_braggvectors(
     braggpeaks, lattice, maxPeakSpacing, qx_shift=0, qy_shift=0, mask=None
 ):
     """
-    Using the peak positions (qx,qy) and indices (h,k) in the PointList lattice,
+    Using the peak positions (qx,qy) and indices (g1_ind,g2_ind) in the PointList lattice,
     identify the indices for each peak in the PointListArray braggpeaks.
     Return a new braggpeaks_indexed PointListArray, containing a copy of braggpeaks plus
-    three additional data columns -- 'h','k', and 'index_mask' -- specifying the peak
-    indices with the ints (h,k) and indicating whether the peak was successfully indexed
+    three additional data columns -- 'g1_ind','g2_ind', and 'index_mask' -- specifying the peak
+    indices with the ints (g1_ind,g2_ind) and indicating whether the peak was successfully indexed
     or not with the bool index_mask. If `mask` is specified, only the locations where
     mask is True are indexed.
 
     Args:
         braggpeaks (PointListArray): the braggpeaks to index. Must contain
             the coordinates 'qx', 'qy', and 'intensity'
-        lattice (PointList): the positions (qx,qy) of the (h,k) lattice points.
-            Must contain the coordinates 'qx', 'qy', 'h', and 'k'
+        lattice (PointList): the positions (qx,qy) of the (g1_ind,g2_ind) lattice points.
+            Must contain the coordinates 'qx', 'qy', 'g1_ind', and 'g2_ind'
         maxPeakSpacing (float): Maximum distance from the ideal lattice points
             to include a peak for indexing
         qx_shift,qy_shift (number): the shift of the origin in the `lattice` PointList
@@ -88,7 +88,7 @@ def add_indices_to_braggvectors(
 
     Returns:
         (PointListArray): The original braggpeaks pointlistarray, with new coordinates
-        'h', 'k', containing the indices of each indexable peak.
+        'g1_ind', 'g2_ind', containing the indices of each indexable peak.
     """
 
     # assert isinstance(braggpeaks,BraggVectors)
@@ -107,8 +107,8 @@ def add_indices_to_braggvectors(
         ("qx", float),
         ("qy", float),
         ("intensity", float),
-        ("h", int),
-        ("k", int),
+        ("g1_ind", int),
+        ("g2_ind", int),
     ]
 
     indexed_braggpeaks = PointListArray(
@@ -140,8 +140,8 @@ def add_indices_to_braggvectors(
                             pl.data["qx"][i],
                             pl.data["qy"][i],
                             pl.data["intensity"][i],
-                            lattice.data["h"][ind],
-                            lattice.data["k"][ind],
+                            lattice.data["g1_ind"][ind],
+                            lattice.data["g2_ind"][ind],
                         )
                     )
 
@@ -150,12 +150,12 @@ def add_indices_to_braggvectors(
 
 def fit_lattice_vectors(braggpeaks, x0=0, y0=0, minNumPeaks=5):
     """
-    Fits lattice vectors g1,g2 to braggpeaks given some known (h,k) indexing.
+    Fits lattice vectors g1,g2 to braggpeaks given some known (g1_ind,g2_ind) indexing.
 
     Args:
         braggpeaks (PointList): A 6 coordinate PointList containing the data to fit.
             Coords are 'qx','qy' (the bragg peak positions), 'intensity' (used as a
-            weighting factor when fitting), 'h','k' (indexing). May optionally also
+            weighting factor when fitting), 'g1_ind','g2_ind' (indexing). May optionally also
             contain 'index_mask' (bool), indicating which peaks have been successfully
             indixed and should be used.
         x0 (float): x-coord of the origin
@@ -176,22 +176,22 @@ def fit_lattice_vectors(braggpeaks, x0=0, y0=0, minNumPeaks=5):
     """
     assert isinstance(braggpeaks, PointList)
     assert np.all(
-        [name in braggpeaks.dtype.names for name in ("qx", "qy", "intensity", "h", "k")]
+        [name in braggpeaks.dtype.names for name in ("qx", "qy", "intensity", "g1_ind", "g2_ind")]
     )
     braggpeaks = braggpeaks.copy()
 
     # Remove unindexed peaks
     if "index_mask" in braggpeaks.dtype.names:
-        deletemask = braggpeaks.data["index_mask"] == False
+        deletemask = braggpeaks.data["index_mask"] is False
         braggpeaks.remove(deletemask)
 
     # Check to ensure enough peaks are present
     if braggpeaks.length < minNumPeaks:
         return None, None, None, None, None, None, None
 
-    # Get M, the matrix of (h,k) indices
-    h, k = braggpeaks.data["h"], braggpeaks.data["k"]
-    M = np.vstack((np.ones_like(h, dtype=int), h, k)).T
+    # Get M, the matrix of (g1_ind,g2_ind) indices
+    g1_ind, g2_ind = braggpeaks.data["g1_ind"], braggpeaks.data["g2_ind"]
+    M = np.vstack((np.ones_like(g1_ind, dtype=int), g1_ind, g2_ind)).T
 
     # Get alpha, the matrix of measured Bragg peak positions
     alpha = np.vstack((braggpeaks.data["qx"] - x0, braggpeaks.data["qy"] - y0)).T
@@ -223,7 +223,7 @@ def fit_lattice_vectors_all_DPs(braggpeaks, x0=0, y0=0, minNumPeaks=5):
     Args:
         braggpeaks (PointList): A 6 coordinate PointList containing the data to fit.
             Coords are 'qx','qy' (the bragg peak positions), 'intensity' (used as a
-            weighting factor when fitting), 'h','k' (indexing). May optionally also
+            weighting factor when fitting), 'g1_ind','g2_ind' (indexing). May optionally also
             contain 'index_mask' (bool), indicating which peaks have been successfully
             indixed and should be used.
         x0 (float): x-coord of the origin
@@ -246,7 +246,7 @@ def fit_lattice_vectors_all_DPs(braggpeaks, x0=0, y0=0, minNumPeaks=5):
     """
     assert isinstance(braggpeaks, PointListArray)
     assert np.all(
-        [name in braggpeaks.dtype.names for name in ("qx", "qy", "intensity", "h", "k")]
+        [name in braggpeaks.dtype.names for name in ("qx", "qy", "intensity", "g1_ind", "g2_ind")]
     )
 
     # Make RealSlice to contain outputs
@@ -272,7 +272,7 @@ def fit_lattice_vectors_all_DPs(braggpeaks, x0=0, y0=0, minNumPeaks=5):
         # Store data
         if g1x is not None:
             g1g2_map.get_slice("x0").data[Rx, Ry] = qx0
-            g1g2_map.get_slice("y0").data[Rx, Ry] = qx0
+            g1g2_map.get_slice("y0").data[Rx, Ry] = qy0 # Assume this is a correct change
             g1g2_map.get_slice("g1x").data[Rx, Ry] = g1x
             g1g2_map.get_slice("g1y").data[Rx, Ry] = g1y
             g1g2_map.get_slice("g2x").data[Rx, Ry] = g2x

From ac55ee266f1b091bf5f0545e58f73c5515bc2966 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Tue, 7 Nov 2023 17:21:22 -0800
Subject: [PATCH 02/18] adding g1_ind g2_ind  fields to overlays

---
 py4DSTEM/visualize/overlay.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/py4DSTEM/visualize/overlay.py b/py4DSTEM/visualize/overlay.py
index 32baff443..991cc0485 100644
--- a/py4DSTEM/visualize/overlay.py
+++ b/py4DSTEM/visualize/overlay.py
@@ -553,8 +553,15 @@ def add_bragg_index_labels(ax, d):
     assert "bragg_directions" in d.keys()
     bragg_directions = d["bragg_directions"]
     assert isinstance(bragg_directions, PointList)
-    for k in ("qx", "qy", "h", "k"):
-        assert k in bragg_directions.data.dtype.fields
+    # check pointlist has ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields
+    assert (
+        all(key in bragg_directions.data.dtype.names for key in ("qx", "qy", "h", "k")) or
+        all(key in bragg_directions.data.dtype.names for key in ("qx", "qy", "g1_ind", "g2_ind"))
+    ), 'pointlist must contain ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields'
+    
+
+    # for k in ("qx", "qy", "h", "k"):
+    #     assert k in bragg_directions.data.dtype.fields
     include_l = True if "l" in bragg_directions.data.dtype.fields else False
     # offsets
     hoffset = d["hoffset"] if "hoffset" in d.keys() else 0
@@ -586,10 +593,13 @@ def add_bragg_index_labels(ax, d):
         x, y = bragg_directions.data["qx"][i], bragg_directions.data["qy"][i]
         x -= voffset
         y += hoffset
-        h, k = bragg_directions.data["h"][i], bragg_directions.data["k"][i]
+        # bit of a hack to get either "h" or "g1_ind"
+        h = bragg_directions.data["h"][i] if not ValueError else bragg_directions.data["g1_ind"][i]
+        k = bragg_directions.data["k"][i] if not ValueError else bragg_directions.data["g2_ind"][i]
         h = str(h) if h >= 0 else r"$\overline{{{}}}$".format(np.abs(h))
         k = str(k) if k >= 0 else r"$\overline{{{}}}$".format(np.abs(k))
         s = h + "," + k
+        # TODO might need to to add g3_ind check
         if include_l:
             l = bragg_directions.data["l"][i]
             l = str(l) if l >= 0 else r"$\overline{{{}}}$".format(np.abs(l))

From 6523aafb03ca3157d385611eb4713a7ab5d29d30 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Tue, 7 Nov 2023 17:23:55 -0800
Subject: [PATCH 03/18] adding g1_ind, g2_ind to plotting

---
 py4DSTEM/process/strain/strain.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/py4DSTEM/process/strain/strain.py b/py4DSTEM/process/strain/strain.py
index ab8a46a9a..d559f35d4 100644
--- a/py4DSTEM/process/strain/strain.py
+++ b/py4DSTEM/process/strain/strain.py
@@ -505,8 +505,8 @@ def fit_basis_vectors(
                 ("qx", float),
                 ("qy", float),
                 ("intensity", float),
-                ("h", int),
-                ("k", int),
+                ("g1_ind", int),
+                ("g2_ind", int),
             ],
             shape=self.braggvectors.Rshape,
         )
@@ -542,8 +542,8 @@ def fit_basis_vectors(
                                 pl.data["qx"][i],
                                 pl.data["qy"][i],
                                 pl.data["intensity"][i],
-                                self.braggdirections.data["h"][ind],
-                                self.braggdirections.data["k"][ind],
+                                self.braggdirections.data["g1_ind"][ind],
+                                self.braggdirections.data["g2_ind"][ind],
                             )
                         )
         self.bragg_vectors_indexed = indexed_braggpeaks
@@ -1510,11 +1510,14 @@ def show_bragg_indexing(
             The display image
         bragg_directions : PointList
             The Bragg scattering directions.  Must have coordinates
-            'qx','qy','h', and 'k'. Optionally may also have 'l'.
+            ('qx','qy','h', and 'k') or ('qx','qy','g1_ind', and 'g2_ind'. Optionally may also have 'l'.
         """
         assert isinstance(bragg_directions, PointList)
-        for k in ("qx", "qy", "h", "k"):
-            assert k in bragg_directions.data.dtype.fields
+        # checking if it has h, k or g1_ind, g2_ind
+        assert (
+                all(key in bragg_directions.data.dtype.names for key in ("qx", "qy", "h", "k")) or
+                all(key in bragg_directions.data.dtype.names for key in ("qx", "qy", "g1_ind", "g2_ind"))
+            ), 'pointlist must contain ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields'
 
         if figax is None:
             fig, ax = show(ar, returnfig=True, **kwargs)
@@ -1533,6 +1536,7 @@ def show_bragg_indexing(
             "pointsize": pointsize,
             "pointcolor": pointcolor,
         }
+        # this can take ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields 
         add_bragg_index_labels(ax, d)
 
         if returnfig:

From a1444ac2f8dbe4602abeb65e3e664365fa5ae86b Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Tue, 7 Nov 2023 22:55:13 -0800
Subject: [PATCH 04/18] adding set_hkl func

---
 py4DSTEM/process/strain/strain.py | 72 +++++++++++++++++++++++++++++++
 1 file changed, 72 insertions(+)

diff --git a/py4DSTEM/process/strain/strain.py b/py4DSTEM/process/strain/strain.py
index d559f35d4..facbc0b59 100644
--- a/py4DSTEM/process/strain/strain.py
+++ b/py4DSTEM/process/strain/strain.py
@@ -390,6 +390,78 @@ def choose_basis_vectors(
         else:
             return
 
+    def set_hkl(
+            self,
+            g1_hkl,
+            g2_hkl,
+    ):
+        
+        g1_hkl = np.array(g1_hkl)
+        g2_hkl = np.array(g2_hkl)
+
+        # TODO check this is correct 
+        # calculate the direct beam 
+        direct_beam = np.cross(g1_hkl, g2_hkl)
+        direct_beam = direct_beam / np.gcd.reduce(direct_beam)
+
+        # Initialize a PLA
+        bvs_hkl = PointListArray(
+            shape = self.shape,
+            dtype = [
+                    ('qx',float),
+                    ('qy',float),
+                    ('intensity',float),
+                    ('h',int),
+                    ('k',int),
+                    ('l',int),
+                ]
+            )   
+        # loop over the probe posistions 
+        for Rx, Ry in tqdmnd(
+            self.shape[0],
+            self.shape[1],
+            desc="Converting (g1_ind,g2_ind) to (h,k,l)",
+            unit="DP",
+            unit_scale=True,
+        ):
+            # get a single indexed 
+            braggvectors_indexed_dp = self.bragg_vectors_indexed[Rx,Ry]
+            
+            # make a Pointlsit
+            bvs_hkl_curr= PointList(
+            data = np.empty(
+                len(braggvectors_indexed_dp),
+                dtype = bvs_hkl.dtype
+                )
+            )
+            # populate qx, qy and intensity fields
+            bvs_hkl_curr.data['qx'] = braggvectors_indexed_dp['qx']
+            bvs_hkl_curr.data['qy'] = braggvectors_indexed_dp['qy']
+            bvs_hkl_curr.data['intensity'] = braggvectors_indexed_dp['intensity']
+
+            # calcuate the hkl vectors
+            vectors_hkl = g1_hkl[:,np.newaxis]*braggvectors_indexed_dp['g1_ind'] +\
+            g2_hkl[:,np.newaxis]*braggvectors_indexed_dp['g2_ind'] 
+            # self.vectors_hkl = vectors_hkl
+            # I think 0,0,0 reflection will always be at index 0 but this will find it 
+            index = np.where((vectors_hkl == 0).all(axis=0))[0][0] # returns tuple and we want the zeroth index
+            # print(vectors_hkl.shape)
+            # print(direct_beam.shape)
+            vectors_hkl[:,index] = direct_beam
+            
+            # populate h,k,l fields
+            # print(vectors_hkl.shape)
+            # bvs_hkl_curr.data['h'] = vectors_hkl[0,:]
+            # bvs_hkl_curr.data['k'] = vectors_hkl[1,:]
+            # bvs_hkl_curr.data['l'] = vectors_hkl[2,:]
+            bvs_hkl_curr.data['h'],bvs_hkl_curr.data['k'],bvs_hkl_curr.data['l'] = np.vsplit(vectors_hkl,3)
+            
+            # add to the PLA
+            bvs_hkl[Rx,Ry] += bvs_hkl_curr
+        
+        # add the PLA to the Strainmap object
+        self.bragg_vectors_indexed_hkl = bvs_hkl
+
     def set_max_peak_spacing(
         self,
         max_peak_spacing,

From eb2e7af48e1b09a30aec99a1f78abd0969ab496c Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Tue, 7 Nov 2023 23:09:37 -0800
Subject: [PATCH 05/18] adding docstring

---
 py4DSTEM/process/strain/strain.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/py4DSTEM/process/strain/strain.py b/py4DSTEM/process/strain/strain.py
index facbc0b59..403ef5e4d 100644
--- a/py4DSTEM/process/strain/strain.py
+++ b/py4DSTEM/process/strain/strain.py
@@ -392,9 +392,16 @@ def choose_basis_vectors(
 
     def set_hkl(
             self,
-            g1_hkl,
-            g2_hkl,
+            g1_hkl:list[int,int,int] | tuple[int,int,int] | np.ndarray[int],
+            g2_hkl:list[int,int,int] | tuple[int,int,int] | np.ndarray[int],
     ):
+        """
+        calculate the [h,k,l] reflections from the `g1_ind`,`g2_ind` from known 'g1_hkl` and 'g2_hkl' reflections.
+        Creates 'bragg_vectors_indexed_hkl' attribute
+        Args:
+            g1_hkl (list[int] | tuple[int,int,int] | np.ndarray[int]): known [h,k,l] reflection for g1_vector 
+            g2_hkl (list[int] | tuple[int,int,int] | np.ndarray[int]): known [h,k,l] reflection for g1_vector 
+        """
         
         g1_hkl = np.array(g1_hkl)
         g2_hkl = np.array(g2_hkl)

From 03a4c9d0e7997a388cd3d5984d5e096db2fe7ea8 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Tue, 7 Nov 2023 23:10:54 -0800
Subject: [PATCH 06/18] black

---
 py4DSTEM/visualize/overlay.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/py4DSTEM/visualize/overlay.py b/py4DSTEM/visualize/overlay.py
index 991cc0485..8c6c06d9f 100644
--- a/py4DSTEM/visualize/overlay.py
+++ b/py4DSTEM/visualize/overlay.py
@@ -554,11 +554,12 @@ def add_bragg_index_labels(ax, d):
     bragg_directions = d["bragg_directions"]
     assert isinstance(bragg_directions, PointList)
     # check pointlist has ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields
-    assert (
-        all(key in bragg_directions.data.dtype.names for key in ("qx", "qy", "h", "k")) or
-        all(key in bragg_directions.data.dtype.names for key in ("qx", "qy", "g1_ind", "g2_ind"))
+    assert all(
+        key in bragg_directions.data.dtype.names for key in ("qx", "qy", "h", "k")
+    ) or all(
+        key in bragg_directions.data.dtype.names
+        for key in ("qx", "qy", "g1_ind", "g2_ind")
     ), 'pointlist must contain ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields'
-    
 
     # for k in ("qx", "qy", "h", "k"):
     #     assert k in bragg_directions.data.dtype.fields
@@ -594,8 +595,16 @@ def add_bragg_index_labels(ax, d):
         x -= voffset
         y += hoffset
         # bit of a hack to get either "h" or "g1_ind"
-        h = bragg_directions.data["h"][i] if not ValueError else bragg_directions.data["g1_ind"][i]
-        k = bragg_directions.data["k"][i] if not ValueError else bragg_directions.data["g2_ind"][i]
+        h = (
+            bragg_directions.data["h"][i]
+            if not ValueError
+            else bragg_directions.data["g1_ind"][i]
+        )
+        k = (
+            bragg_directions.data["k"][i]
+            if not ValueError
+            else bragg_directions.data["g2_ind"][i]
+        )
         h = str(h) if h >= 0 else r"$\overline{{{}}}$".format(np.abs(h))
         k = str(k) if k >= 0 else r"$\overline{{{}}}$".format(np.abs(k))
         s = h + "," + k

From 724587cd006d222fc325dd465492a0cb417a7444 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Tue, 7 Nov 2023 23:11:31 -0800
Subject: [PATCH 07/18] docstring and black

---
 py4DSTEM/process/strain/strain.py | 97 +++++++++++++++++--------------
 1 file changed, 52 insertions(+), 45 deletions(-)

diff --git a/py4DSTEM/process/strain/strain.py b/py4DSTEM/process/strain/strain.py
index 403ef5e4d..f622f59ae 100644
--- a/py4DSTEM/process/strain/strain.py
+++ b/py4DSTEM/process/strain/strain.py
@@ -391,39 +391,39 @@ def choose_basis_vectors(
             return
 
     def set_hkl(
-            self,
-            g1_hkl:list[int,int,int] | tuple[int,int,int] | np.ndarray[int],
-            g2_hkl:list[int,int,int] | tuple[int,int,int] | np.ndarray[int],
+        self,
+        g1_hkl: list[int, int, int] | tuple[int, int, int] | np.ndarray[int],
+        g2_hkl: list[int, int, int] | tuple[int, int, int] | np.ndarray[int],
     ):
         """
         calculate the [h,k,l] reflections from the `g1_ind`,`g2_ind` from known 'g1_hkl` and 'g2_hkl' reflections.
         Creates 'bragg_vectors_indexed_hkl' attribute
         Args:
-            g1_hkl (list[int] | tuple[int,int,int] | np.ndarray[int]): known [h,k,l] reflection for g1_vector 
-            g2_hkl (list[int] | tuple[int,int,int] | np.ndarray[int]): known [h,k,l] reflection for g1_vector 
+            g1_hkl (list[int] | tuple[int,int,int] | np.ndarray[int]): known [h,k,l] reflection for g1_vector
+            g2_hkl (list[int] | tuple[int,int,int] | np.ndarray[int]): known [h,k,l] reflection for g1_vector
         """
-        
+
         g1_hkl = np.array(g1_hkl)
         g2_hkl = np.array(g2_hkl)
 
-        # TODO check this is correct 
-        # calculate the direct beam 
+        # TODO check this is correct
+        # calculate the direct beam
         direct_beam = np.cross(g1_hkl, g2_hkl)
         direct_beam = direct_beam / np.gcd.reduce(direct_beam)
 
         # Initialize a PLA
         bvs_hkl = PointListArray(
-            shape = self.shape,
-            dtype = [
-                    ('qx',float),
-                    ('qy',float),
-                    ('intensity',float),
-                    ('h',int),
-                    ('k',int),
-                    ('l',int),
-                ]
-            )   
-        # loop over the probe posistions 
+            shape=self.shape,
+            dtype=[
+                ("qx", float),
+                ("qy", float),
+                ("intensity", float),
+                ("h", int),
+                ("k", int),
+                ("l", int),
+            ],
+        )
+        # loop over the probe posistions
         for Rx, Ry in tqdmnd(
             self.shape[0],
             self.shape[1],
@@ -431,41 +431,46 @@ def set_hkl(
             unit="DP",
             unit_scale=True,
         ):
-            # get a single indexed 
-            braggvectors_indexed_dp = self.bragg_vectors_indexed[Rx,Ry]
-            
+            # get a single indexed
+            braggvectors_indexed_dp = self.bragg_vectors_indexed[Rx, Ry]
+
             # make a Pointlsit
-            bvs_hkl_curr= PointList(
-            data = np.empty(
-                len(braggvectors_indexed_dp),
-                dtype = bvs_hkl.dtype
-                )
+            bvs_hkl_curr = PointList(
+                data=np.empty(len(braggvectors_indexed_dp), dtype=bvs_hkl.dtype)
             )
             # populate qx, qy and intensity fields
-            bvs_hkl_curr.data['qx'] = braggvectors_indexed_dp['qx']
-            bvs_hkl_curr.data['qy'] = braggvectors_indexed_dp['qy']
-            bvs_hkl_curr.data['intensity'] = braggvectors_indexed_dp['intensity']
+            bvs_hkl_curr.data["qx"] = braggvectors_indexed_dp["qx"]
+            bvs_hkl_curr.data["qy"] = braggvectors_indexed_dp["qy"]
+            bvs_hkl_curr.data["intensity"] = braggvectors_indexed_dp["intensity"]
 
             # calcuate the hkl vectors
-            vectors_hkl = g1_hkl[:,np.newaxis]*braggvectors_indexed_dp['g1_ind'] +\
-            g2_hkl[:,np.newaxis]*braggvectors_indexed_dp['g2_ind'] 
+            vectors_hkl = (
+                g1_hkl[:, np.newaxis] * braggvectors_indexed_dp["g1_ind"]
+                + g2_hkl[:, np.newaxis] * braggvectors_indexed_dp["g2_ind"]
+            )
             # self.vectors_hkl = vectors_hkl
-            # I think 0,0,0 reflection will always be at index 0 but this will find it 
-            index = np.where((vectors_hkl == 0).all(axis=0))[0][0] # returns tuple and we want the zeroth index
+            # I think 0,0,0 reflection will always be at index 0 but this will find it
+            index = np.where((vectors_hkl == 0).all(axis=0))[0][
+                0
+            ]  # returns tuple and we want the zeroth index
             # print(vectors_hkl.shape)
             # print(direct_beam.shape)
-            vectors_hkl[:,index] = direct_beam
-            
+            vectors_hkl[:, index] = direct_beam
+
             # populate h,k,l fields
             # print(vectors_hkl.shape)
             # bvs_hkl_curr.data['h'] = vectors_hkl[0,:]
             # bvs_hkl_curr.data['k'] = vectors_hkl[1,:]
             # bvs_hkl_curr.data['l'] = vectors_hkl[2,:]
-            bvs_hkl_curr.data['h'],bvs_hkl_curr.data['k'],bvs_hkl_curr.data['l'] = np.vsplit(vectors_hkl,3)
-            
+            (
+                bvs_hkl_curr.data["h"],
+                bvs_hkl_curr.data["k"],
+                bvs_hkl_curr.data["l"],
+            ) = np.vsplit(vectors_hkl, 3)
+
             # add to the PLA
-            bvs_hkl[Rx,Ry] += bvs_hkl_curr
-        
+            bvs_hkl[Rx, Ry] += bvs_hkl_curr
+
         # add the PLA to the Strainmap object
         self.bragg_vectors_indexed_hkl = bvs_hkl
 
@@ -1593,10 +1598,12 @@ def show_bragg_indexing(
         """
         assert isinstance(bragg_directions, PointList)
         # checking if it has h, k or g1_ind, g2_ind
-        assert (
-                all(key in bragg_directions.data.dtype.names for key in ("qx", "qy", "h", "k")) or
-                all(key in bragg_directions.data.dtype.names for key in ("qx", "qy", "g1_ind", "g2_ind"))
-            ), 'pointlist must contain ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields'
+        assert all(
+            key in bragg_directions.data.dtype.names for key in ("qx", "qy", "h", "k")
+        ) or all(
+            key in bragg_directions.data.dtype.names
+            for key in ("qx", "qy", "g1_ind", "g2_ind")
+        ), 'pointlist must contain ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields'
 
         if figax is None:
             fig, ax = show(ar, returnfig=True, **kwargs)
@@ -1615,7 +1622,7 @@ def show_bragg_indexing(
             "pointsize": pointsize,
             "pointcolor": pointcolor,
         }
-        # this can take ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields 
+        # this can take ("qx", "qy", "h", "k") or ("qx", "qy", "g1_ind", "g2_ind") fields
         add_bragg_index_labels(ax, d)
 
         if returnfig:

From 3e2d7feca4fa6beab6be61eba163b8a51ce40059 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Tue, 7 Nov 2023 23:12:29 -0800
Subject: [PATCH 08/18] black

---
 py4DSTEM/process/strain/latticevectors.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/py4DSTEM/process/strain/latticevectors.py b/py4DSTEM/process/strain/latticevectors.py
index 65b5f6c69..587ac0e8e 100644
--- a/py4DSTEM/process/strain/latticevectors.py
+++ b/py4DSTEM/process/strain/latticevectors.py
@@ -47,7 +47,7 @@ def index_bragg_directions(x0, y0, gx, gy, g1, g2):
 
     # Get g1_ind,g2_ind
     g1_ind = M[:, 0]
-    g2_ind  = M[:, 1]
+    g2_ind = M[:, 1]
 
     # Store in a PointList
     coords = [("qx", float), ("qy", float), ("g1_ind", int), ("g2_ind", int)]
@@ -176,7 +176,10 @@ def fit_lattice_vectors(braggpeaks, x0=0, y0=0, minNumPeaks=5):
     """
     assert isinstance(braggpeaks, PointList)
     assert np.all(
-        [name in braggpeaks.dtype.names for name in ("qx", "qy", "intensity", "g1_ind", "g2_ind")]
+        [
+            name in braggpeaks.dtype.names
+            for name in ("qx", "qy", "intensity", "g1_ind", "g2_ind")
+        ]
     )
     braggpeaks = braggpeaks.copy()
 
@@ -246,7 +249,10 @@ def fit_lattice_vectors_all_DPs(braggpeaks, x0=0, y0=0, minNumPeaks=5):
     """
     assert isinstance(braggpeaks, PointListArray)
     assert np.all(
-        [name in braggpeaks.dtype.names for name in ("qx", "qy", "intensity", "g1_ind", "g2_ind")]
+        [
+            name in braggpeaks.dtype.names
+            for name in ("qx", "qy", "intensity", "g1_ind", "g2_ind")
+        ]
     )
 
     # Make RealSlice to contain outputs
@@ -272,7 +278,8 @@ def fit_lattice_vectors_all_DPs(braggpeaks, x0=0, y0=0, minNumPeaks=5):
         # Store data
         if g1x is not None:
             g1g2_map.get_slice("x0").data[Rx, Ry] = qx0
-            g1g2_map.get_slice("y0").data[Rx, Ry] = qy0 # Assume this is a correct change
+            # Assume this is a correct change
+            g1g2_map.get_slice("y0").data[Rx, Ry] = qy0
             g1g2_map.get_slice("g1x").data[Rx, Ry] = g1x
             g1g2_map.get_slice("g1y").data[Rx, Ry] = g1y
             g1g2_map.get_slice("g2x").data[Rx, Ry] = g2x

From 283f8eab0b61eeeaffab2b143e53e99761ba1dd3 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Tue, 7 Nov 2023 23:16:34 -0800
Subject: [PATCH 09/18] changing type hints for python 3.9

---
 py4DSTEM/process/strain/strain.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/py4DSTEM/process/strain/strain.py b/py4DSTEM/process/strain/strain.py
index f622f59ae..89f30a273 100644
--- a/py4DSTEM/process/strain/strain.py
+++ b/py4DSTEM/process/strain/strain.py
@@ -1,7 +1,7 @@
 # Defines the Strain class
 
 import warnings
-from typing import Optional
+from typing import Optional, List, Tuple, Union
 
 import matplotlib.pyplot as plt
 from matplotlib.patches import Circle
@@ -392,8 +392,8 @@ def choose_basis_vectors(
 
     def set_hkl(
         self,
-        g1_hkl: list[int, int, int] | tuple[int, int, int] | np.ndarray[int],
-        g2_hkl: list[int, int, int] | tuple[int, int, int] | np.ndarray[int],
+        g1_hkl: Union[List[int], Tuple[int, int, int], np.ndarray[np.int64]],
+        g2_hkl: Union[List[int], Tuple[int, int, int], np.ndarray[np.int64]],
     ):
         """
         calculate the [h,k,l] reflections from the `g1_ind`,`g2_ind` from known 'g1_hkl` and 'g2_hkl' reflections.

From 3bb6f0a12e7dd090c3d9eb2c4a4eeab05afd6cd0 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Wed, 8 Nov 2023 10:20:28 -0800
Subject: [PATCH 10/18] removing incorrect direct beam stuff

---
 py4DSTEM/process/strain/strain.py | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/py4DSTEM/process/strain/strain.py b/py4DSTEM/process/strain/strain.py
index 89f30a273..6ae27e008 100644
--- a/py4DSTEM/process/strain/strain.py
+++ b/py4DSTEM/process/strain/strain.py
@@ -406,11 +406,6 @@ def set_hkl(
         g1_hkl = np.array(g1_hkl)
         g2_hkl = np.array(g2_hkl)
 
-        # TODO check this is correct
-        # calculate the direct beam
-        direct_beam = np.cross(g1_hkl, g2_hkl)
-        direct_beam = direct_beam / np.gcd.reduce(direct_beam)
-
         # Initialize a PLA
         bvs_hkl = PointListArray(
             shape=self.shape,
@@ -449,13 +444,6 @@ def set_hkl(
                 + g2_hkl[:, np.newaxis] * braggvectors_indexed_dp["g2_ind"]
             )
             # self.vectors_hkl = vectors_hkl
-            # I think 0,0,0 reflection will always be at index 0 but this will find it
-            index = np.where((vectors_hkl == 0).all(axis=0))[0][
-                0
-            ]  # returns tuple and we want the zeroth index
-            # print(vectors_hkl.shape)
-            # print(direct_beam.shape)
-            vectors_hkl[:, index] = direct_beam
 
             # populate h,k,l fields
             # print(vectors_hkl.shape)

From cf256ec982e3322aa7ac3f484d314231055ed323 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Sun, 10 Dec 2023 21:26:18 -0800
Subject: [PATCH 11/18] adding h,k,l property to BVects

---
 py4DSTEM/braggvectors/braggvectors.py | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/py4DSTEM/braggvectors/braggvectors.py b/py4DSTEM/braggvectors/braggvectors.py
index e81eeb62f..0ee705b74 100644
--- a/py4DSTEM/braggvectors/braggvectors.py
+++ b/py4DSTEM/braggvectors/braggvectors.py
@@ -336,7 +336,7 @@ class BVects:
     """
     Enables
 
-        >>> v.qx,v.qy,v.I
+        >>> v.qx,v.qy,v.I, optionally, v.h,v.k,v.l
 
     -like access to a collection of Bragg vector.
     """
@@ -361,11 +361,32 @@ def I(self):
     def data(self):
         return self._data
 
+    @property
+    def h(self):
+        try:
+            return self._data["h"]
+        except ValueError:
+            raise AttributeError("h indicies not set")
+
+    @property
+    def k(self):
+        try:
+            return self._data["k"]
+        except ValueError:
+            raise AttributeError("k indicies not set")
+
+    @property
+    def l(self):
+        try:
+            return self._data["l"]
+        except ValueError:
+            raise AttributeError("l indicies not set")
+
     def __repr__(self):
         space = " " * len(self.__class__.__name__) + "  "
         string = f"{self.__class__.__name__}( "
         string += f"A set of {len(self.data)} bragg vectors."
-        string += " Access data with .qx, .qy, .I, or .data.)"
+        string += " Access data with .qx, .qy, .I, or .data. Optionally .h, .k, and.l, if indexed)"
         return string
 
 

From b5cdd808d15a3b3b882665f84bd1f188a19cc7a6 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Sun, 10 Dec 2023 21:40:00 -0800
Subject: [PATCH 12/18] adding plotting of BVects

---
 py4DSTEM/braggvectors/braggvectors.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/py4DSTEM/braggvectors/braggvectors.py b/py4DSTEM/braggvectors/braggvectors.py
index 0ee705b74..bc8797940 100644
--- a/py4DSTEM/braggvectors/braggvectors.py
+++ b/py4DSTEM/braggvectors/braggvectors.py
@@ -3,6 +3,7 @@
 from py4DSTEM.data import Data
 from emdfile import Custom, PointListArray, PointList, Metadata
 from py4DSTEM.braggvectors.braggvector_methods import BraggVectorMethods
+from py4DSTEM.process.diffraction import plot_diffraction_pattern
 from os.path import basename
 import numpy as np
 from warnings import warn
@@ -389,6 +390,31 @@ def __repr__(self):
         string += " Access data with .qx, .qy, .I, or .data. Optionally .h, .k, and.l, if indexed)"
         return string
 
+    def plot(
+        self,
+        returnfig: bool = False,
+        **kwargs,
+    ):
+        """
+        Plot the diffraction pattern.
+        Calls `py4DSTEM.process.diffraction.plot_diffraction_pattern` and passes kwargs to it.
+
+        Parameters
+        ----------
+        returnfig: bool
+            If True the figure is returned, else its ploted but not returned. Defaults to False
+
+        Returns
+        -------
+        figure : matplotlib object
+            If `returnfig` is True, the figure is returned.
+
+        """
+        if returnfig:
+            return plot_diffraction_pattern(self, **kwargs)
+        else:
+            plot_diffraction_pattern(self, **kwargs)
+
 
 class RawVectorGetter:
     def __init__(

From b3e90e2020499beeb13a2a1da61b5f07a8c57f7a Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Sun, 10 Dec 2023 21:45:30 -0800
Subject: [PATCH 13/18] moving import to avoid circular ref

---
 py4DSTEM/braggvectors/braggvectors.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/py4DSTEM/braggvectors/braggvectors.py b/py4DSTEM/braggvectors/braggvectors.py
index bc8797940..c6e7e3606 100644
--- a/py4DSTEM/braggvectors/braggvectors.py
+++ b/py4DSTEM/braggvectors/braggvectors.py
@@ -3,7 +3,6 @@
 from py4DSTEM.data import Data
 from emdfile import Custom, PointListArray, PointList, Metadata
 from py4DSTEM.braggvectors.braggvector_methods import BraggVectorMethods
-from py4DSTEM.process.diffraction import plot_diffraction_pattern
 from os.path import basename
 import numpy as np
 from warnings import warn
@@ -410,6 +409,8 @@ def plot(
             If `returnfig` is True, the figure is returned.
 
         """
+        from py4DSTEM.process.diffraction import plot_diffraction_pattern
+
         if returnfig:
             return plot_diffraction_pattern(self, **kwargs)
         else:

From 31fb03774a13ca46976cc8edd33f9135080b8818 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Sun, 10 Dec 2023 22:06:39 -0800
Subject: [PATCH 14/18] adding missing returnfig arg

---
 py4DSTEM/braggvectors/braggvectors.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/py4DSTEM/braggvectors/braggvectors.py b/py4DSTEM/braggvectors/braggvectors.py
index c6e7e3606..6ccb46d1f 100644
--- a/py4DSTEM/braggvectors/braggvectors.py
+++ b/py4DSTEM/braggvectors/braggvectors.py
@@ -412,7 +412,7 @@ def plot(
         from py4DSTEM.process.diffraction import plot_diffraction_pattern
 
         if returnfig:
-            return plot_diffraction_pattern(self, **kwargs)
+            return plot_diffraction_pattern(self, returnfig=returnfig, **kwargs)
         else:
             plot_diffraction_pattern(self, **kwargs)
 

From 2684583a458650798f87682bf652e042b3b49bd0 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Sun, 10 Dec 2023 22:07:32 -0800
Subject: [PATCH 15/18] adding plotting to BraggVectors

---
 py4DSTEM/braggvectors/braggvector_methods.py | 45 ++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/py4DSTEM/braggvectors/braggvector_methods.py b/py4DSTEM/braggvectors/braggvector_methods.py
index 70a36dec1..b8ffcbb2e 100644
--- a/py4DSTEM/braggvectors/braggvector_methods.py
+++ b/py4DSTEM/braggvectors/braggvector_methods.py
@@ -813,6 +813,51 @@ def to_strainmap(self, name: str = None):
 
         return StrainMap(self, name) if name else StrainMap(self)
 
+    def plot(
+        self,
+        index: tuple[int,int] | list[int],
+        cal: str = 'cal',
+        returnfig: bool = False,
+        **kwargs
+    ):
+        """
+        Plot Bragg vector, from a specified index.
+        Calls py4DSTEM.process.diffraction.plot_diffraction_pattern(braggvectors.<cal/raw>[index], **kwargs).
+        Optionally can return the figure.
+
+        Parameters
+        ----------
+        index : tuple[int,int] | list[int]
+            index for which Bragg vectors to plot 
+        cal : str, optional
+            Choice to plot calibrated or raw Bragg vectors must be 'raw' or 'cal', by default 'cal'
+        returnfig : bool, optional
+            Boolean to return figure or not, by default False
+
+        Returns
+        -------
+        tuple (figure, axes)
+            matplotlib figure, axes returned if `returnfig` is True
+        """
+        assert cal.lower() in ('cal', 'raw'), f"'cal' must be in ('cal', 'raw') {cal = } passed"
+        from py4DSTEM.process.diffraction import plot_diffraction_pattern
+
+        if cal == 'cal':
+            pl = self.cal[index]
+        else: 
+            pl = self.raw[index]
+        
+        if returnfig:
+            return plot_diffraction_pattern(
+                pl, 
+                returnfig=returnfig,
+                **kwargs,
+                )
+        else: 
+            plot_diffraction_pattern(
+                pl, 
+                **kwargs,
+                )
 
 ######### END BraggVectorMethods CLASS ########
 

From bf4740f55f174fb27670e1202eb9660acd7b0a1b Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Sun, 10 Dec 2023 22:09:02 -0800
Subject: [PATCH 16/18] black

---
 py4DSTEM/braggvectors/braggvector_methods.py | 30 +++++++++++---------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/py4DSTEM/braggvectors/braggvector_methods.py b/py4DSTEM/braggvectors/braggvector_methods.py
index b8ffcbb2e..fc4961015 100644
--- a/py4DSTEM/braggvectors/braggvector_methods.py
+++ b/py4DSTEM/braggvectors/braggvector_methods.py
@@ -815,10 +815,10 @@ def to_strainmap(self, name: str = None):
 
     def plot(
         self,
-        index: tuple[int,int] | list[int],
-        cal: str = 'cal',
+        index: tuple[int, int] | list[int],
+        cal: str = "cal",
         returnfig: bool = False,
-        **kwargs
+        **kwargs,
     ):
         """
         Plot Bragg vector, from a specified index.
@@ -828,7 +828,7 @@ def plot(
         Parameters
         ----------
         index : tuple[int,int] | list[int]
-            index for which Bragg vectors to plot 
+            index for which Bragg vectors to plot
         cal : str, optional
             Choice to plot calibrated or raw Bragg vectors must be 'raw' or 'cal', by default 'cal'
         returnfig : bool, optional
@@ -839,25 +839,29 @@ def plot(
         tuple (figure, axes)
             matplotlib figure, axes returned if `returnfig` is True
         """
-        assert cal.lower() in ('cal', 'raw'), f"'cal' must be in ('cal', 'raw') {cal = } passed"
+        assert cal.lower() in (
+            "cal",
+            "raw",
+        ), f"'cal' must be in ('cal', 'raw') {cal = } passed"
         from py4DSTEM.process.diffraction import plot_diffraction_pattern
 
-        if cal == 'cal':
+        if cal == "cal":
             pl = self.cal[index]
-        else: 
+        else:
             pl = self.raw[index]
-        
+
         if returnfig:
             return plot_diffraction_pattern(
-                pl, 
+                pl,
                 returnfig=returnfig,
                 **kwargs,
-                )
-        else: 
+            )
+        else:
             plot_diffraction_pattern(
-                pl, 
+                pl,
                 **kwargs,
-                )
+            )
+
 
 ######### END BraggVectorMethods CLASS ########
 

From 80fabb126670981541eaa083062356d9963339aa Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Sun, 10 Dec 2023 22:12:12 -0800
Subject: [PATCH 17/18] adding annotation import

---
 py4DSTEM/braggvectors/braggvector_methods.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/py4DSTEM/braggvectors/braggvector_methods.py b/py4DSTEM/braggvectors/braggvector_methods.py
index fc4961015..a4ceb16ca 100644
--- a/py4DSTEM/braggvectors/braggvector_methods.py
+++ b/py4DSTEM/braggvectors/braggvector_methods.py
@@ -1,5 +1,5 @@
 # BraggVectors methods
-
+from __future__ import annotations
 import inspect
 from warnings import warn
 

From cb76e24660b261fb9061772c3563fcbfe4932550 Mon Sep 17 00:00:00 2001
From: bsavitzky <bsavitzky@users.noreply.github.com>
Date: Tue, 23 Jan 2024 20:02:05 -0500
Subject: [PATCH 18/18] minor fixes

---
 py4DSTEM/braggvectors/braggvector_methods.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/py4DSTEM/braggvectors/braggvector_methods.py b/py4DSTEM/braggvectors/braggvector_methods.py
index a4ceb16ca..30ead79a6 100644
--- a/py4DSTEM/braggvectors/braggvector_methods.py
+++ b/py4DSTEM/braggvectors/braggvector_methods.py
@@ -828,7 +828,7 @@ def plot(
         Parameters
         ----------
         index : tuple[int,int] | list[int]
-            index for which Bragg vectors to plot
+            scan position for which Bragg vectors to plot
         cal : str, optional
             Choice to plot calibrated or raw Bragg vectors must be 'raw' or 'cal', by default 'cal'
         returnfig : bool, optional
@@ -839,7 +839,8 @@ def plot(
         tuple (figure, axes)
             matplotlib figure, axes returned if `returnfig` is True
         """
-        assert cal.lower() in (
+        cal = cal.lower()
+        assert cal in (
             "cal",
             "raw",
         ), f"'cal' must be in ('cal', 'raw') {cal = } passed"