From 34a8494ca36dfe3e56d4cc5cbf270700386a8383 Mon Sep 17 00:00:00 2001
From: maclariz <ian.maclaren@glasgow.ac.uk>
Date: Wed, 4 Sep 2024 19:54:35 +0100
Subject: [PATCH] Update digital_dark_field.py

Basic radial functions added
---
 .../process/diffraction/digital_dark_field.py | 122 +++++++++++++++---
 1 file changed, 101 insertions(+), 21 deletions(-)

diff --git a/py4DSTEM/process/diffraction/digital_dark_field.py b/py4DSTEM/process/diffraction/digital_dark_field.py
index ee0ded110..c971dbf5d 100644
--- a/py4DSTEM/process/diffraction/digital_dark_field.py
+++ b/py4DSTEM/process/diffraction/digital_dark_field.py
@@ -314,6 +314,7 @@ def pointlist_to_array(
     ellipse=None,
     pixel=None,
     rotate=None,
+    rphi=True,
 ):
     """
     This function turns the py4DSTEM BraggVectors to a simple numpy array that is more
@@ -331,16 +332,21 @@ def pointlist_to_array(
         if True, applies pixel calibration to bragg_peaks
     rotate: bool
         if True, applies rotational calibration to bragg_peaks
+        rphi: bool
+                if True, generates two extra columns of Qr and Qphi for addressing in polar
+                coordinates
 
     Returns
     ----------
     points_array: numpy array
-         This will be an 2D numpy array of n points x 5 columns:
+         This will be an 2D numpy array of n points x 5 (or 7) columns:
             qx
             qy
             I
             Rx
             Ry
+            (qr)
+            (qphi)
     """
     if center is None:
         center = bragg_peaks.calstate["center"]
@@ -364,27 +370,55 @@ def pointlist_to_array(
             rotate=rotate,
         )
 
-        if i == j == 0:
-            points_array = np.array(
-                [
-                    vectors.qx,
-                    vectors.qy,
-                    vectors.I,
-                    vectors.qx.shape[0] * [i],
-                    vectors.qx.shape[0] * [j],
-                ]
-            ).T
+        if rphi == True:
+            if i == j == 0:
+                points_array = np.array(
+                    [
+                        vectors.qx,
+                        vectors.qy,
+                        vectors.I,
+                        vectors.qx.shape[0] * [i],
+                        vectors.qx.shape[0] * [j],
+                        (vectors.qx**2 + vectors.qy**2) ** 0.5,
+                        np.degrees(np.arctan2(vectors.qx, vectors.qy)),
+                    ]
+                ).T
+            else:
+                nps = np.array(
+                    [
+                        vectors.qx,
+                        vectors.qy,
+                        vectors.I,
+                        vectors.qx.shape[0] * [i],
+                        vectors.qx.shape[0] * [j],
+                        (vectors.qx**2 + vectors.qy**2) ** 0.5,
+                        np.degrees(np.arctan2(vectors.qx, vectors.qy)),
+                    ]
+                ).T
+                points_array = np.vstack((points_array, nps))
+
         else:
-            nps = np.array(
-                [
-                    vectors.qx,
-                    vectors.qy,
-                    vectors.I,
-                    vectors.qx.shape[0] * [i],
-                    vectors.qx.shape[0] * [j],
-                ]
-            ).T
-            points_array = np.vstack((points_array, nps))
+            if i == j == 0:
+                points_array = np.array(
+                    [
+                        vectors.qx,
+                        vectors.qy,
+                        vectors.I,
+                        vectors.qx.shape[0] * [i],
+                        vectors.qx.shape[0] * [j],
+                    ]
+                ).T
+            else:
+                nps = np.array(
+                    [
+                        vectors.qx,
+                        vectors.qy,
+                        vectors.I,
+                        vectors.qx.shape[0] * [i],
+                        vectors.qx.shape[0] * [j],
+                    ]
+                ).T
+                points_array = np.vstack((points_array, nps))
 
     return points_array
 
@@ -458,3 +492,49 @@ def DDFimage(points_array, aperture_positions, Rshape=None, tol=1):
                 intensities2[row, 1].astype(int), intensities2[row, 2].astype(int)
             ] += intensities2[row, 0]
     return image
+
+
+def DDF_radial_image(points_array, radius, Rshape, tol=1):
+    """
+    Calculates a Digital Dark Field image from a list of detected diffraction peak positions in a points_array matching a specific qr radius, within a defined matching tolerance
+
+    Parameters
+    ----------
+    points_array: numpy array
+        as produced by pointlist_to_array and defined in docstring for that function, must be the version with r and phi included
+    radius: float
+        the radius of diffraction spot you wish to image in pixels or calibrated units
+    Rshape: tuple, list, array
+        a 2 element vector giving the real space dimensions.  If not specified, this is determined from the max along points_array
+    tol: float
+        the tolerance in pixels or calibrated units for a point in the points_array to be considered to match to an aperture position in the aperture_positions array
+
+    Returns
+    ----------
+    image: numpy array
+        2D numpy array with dimensions determined by Rshape
+
+    """
+
+    if Rshape is None:
+        Rshape = (
+            np.max(np.max(points_array[:, 3])).astype("int") + 1,
+            np.max(np.max(points_array[:, 4])).astype("int") + 1,
+        )
+
+    points_array_edit = np.delete(
+        points_array, np.where(np.abs(points_array[:, 5] - radius) > tol), axis=0
+    )
+    radialimage = np.zeros(shape=Rshape)
+
+    for i in range(Rshape[0]):
+        for j in range(Rshape[1]):
+            radialimage[i, j] = np.where(
+                np.logical_and(
+                    points_array_edit[:, 3] == i, points_array_edit[:, 4] == j
+                ),
+                points_array_edit[:, 2],
+                0,
+            ).sum()
+
+    return radialimage