From d669df366d86bf47a882d7b666e136fb470b6ea1 Mon Sep 17 00:00:00 2001
From: Stephanie Ribet <stephanie.ribet@gmail.com>
Date: Fri, 1 Dec 2023 15:56:38 -0800
Subject: [PATCH] rotation option

---
 py4DSTEM/datacube/virtualimage.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/py4DSTEM/datacube/virtualimage.py b/py4DSTEM/datacube/virtualimage.py
index b026bfd62..c63911d2c 100644
--- a/py4DSTEM/datacube/virtualimage.py
+++ b/py4DSTEM/datacube/virtualimage.py
@@ -682,6 +682,7 @@ def make_bragg_mask(
         max_q,
         return_sum=True,
         include_origin=True,
+        rotation_deg=0,
         **kwargs,
     ):
         """
@@ -698,6 +699,7 @@ def make_bragg_mask(
                 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
@@ -705,6 +707,14 @@ def make_bragg_mask(
         nas = np.asarray
         g1, g2, origin = nas(g1), nas(g2), nas(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))
         H = int(max_q / L1) + 1