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] 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 ########