Merge pull request #688 from maclariz/DDF-patch-7

Update digital_dark_field.py

py4DSTEM/io/filereaders/read_mib.py

@@ -163,7 +163,6 @@ def scan_size(path, scan):
     header_path = path[:-3] + "hdr"
     result = {}
     if os.path.exists(header_path):
-
         with open(header_path, encoding="UTF-8") as f:
             for line in f:
                 k, v = line.split("\t", 1)

py4DSTEM/process/diffraction/digital_dark_field.py

@@ -334,7 +334,7 @@ def pointlist_to_array(
         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, Qphi is the angle anticlockwise from horizontal to the right
+        coordinates, Qphi is the angle in degrees anticlockwise from horizontal to the right
 
     Returns
     ----------
@@ -494,47 +494,145 @@ def DDFimage(points_array, aperture_positions, Rshape=None, tol=1):
     return image
 
 
-def DDF_radial_image(points_array, radius, Rshape, tol=1):
+def radial_filtered_array(points_array_w_rphi, radius, tol=1):
+    """
+    Calculates a Filtered points array 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_w_rphi: numpy array
+        as produced by pointlist_to_array with rphi=True and defined in docstring for that function
+    radius: float
+        the radius of diffraction spot you wish to filter by in pixels or calibrated units
+    tol: float
+        the tolerance in pixels or calibrated units for a point of qr in the points_array to be considered to match to the radius
+
+    Returns
+    ----------
+    radial_filtered_points_array: numpy array
+         This will be an 2D numpy array of n points x 7 columns:
+            qx
+            qy
+            I
+            Rx
+            Ry
+            qr
+            qphi
+
+    """
+    radial_filtered_points_array = np.delete(
+        points_array_w_rphi,
+        np.where(np.abs(points_array_w_rphi[:, 5] - radius) > tol),
+        axis=0,
+    )
+    return radial_filtered_points_array
+
+
+def DDF_radial_image(points_array_w_rphi, 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
+    points_array_w_rphi: numpy array
+        as produced by pointlist_to_array with rphi=True and defined in docstring for that function
     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
+        the tolerance in pixels or calibrated units for a point of qr in the points_array to be considered to match to the radius
 
     Returns
     ----------
-    image: numpy array
+    radialimage: 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,
+            np.max(np.max(points_array_w_rphi[:, 3])).astype("int") + 1,
+            np.max(np.max(points_array_w_rphi[:, 4])).astype("int") + 1,
         )
 
-    points_array_edit = np.delete(
-        points_array, np.where(np.abs(points_array[:, 5] - radius) > tol), axis=0
+    radial_filtered_points_array = radial_filtered_array(
+        points_array_w_rphi, radius, tol
     )
+
     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
+                    radial_filtered_points_array[:, 3] == i,
+                    radial_filtered_points_array[:, 4] == j,
                 ),
-                points_array_edit[:, 2],
+                radial_filtered_points_array[:, 2],
                 0,
             ).sum()
 
     return radialimage
+
+
+def DDFradialazimuthimage(points_array_w_rphi, radius, phi0, phi1, 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, and only within a defined azimuthal range
+
+    Parameters
+    ----------
+    points_array_w_rphi: numpy array
+        as produced by pointlist_to_array with rphi=True and defined in docstring for that function
+    radius: float
+        the radius of diffraction spot you wish to image in pixels or calibrated units
+    phi0: float
+        Angle in degrees anticlockwise from horizontal-right for setting minimum qphi for inclusion in the image calculation
+    phi1: float
+        Angle in degrees anticlockwise from horizontal-right for setting maximum qphi for inclusion in the image calculation
+    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 of qr in the points_array to be considered to match to the radius
+
+    Returns
+    ----------
+    image: numpy array
+        2D numpy array with dimensions determined by Rshape
+
+    """
+    if Rshape is None:
+        Rshape = (
+            np.max(np.max(points_array_w_rphi[:, 3])).astype("int") + 1,
+            np.max(np.max(points_array_w_rphi[:, 4])).astype("int") + 1,
+        )
+
+    radial_filtered_points_array = radial_filtered_array(
+        points_array_w_rphi, radius, tol
+    )
+
+    rphi_filtered_points_array = np.delete(
+        radial_filtered_points_array,
+        np.where(
+            np.logical_or(
+                radial_filtered_points_array[:, 6] < phi0,
+                radial_filtered_points_array[:, 6] >= phi1,
+            )
+        ),
+        axis=0,
+    )
+    radiusazimuthimage = np.zeros(shape=Rshape)
+
+    for i in range(Rshape[0]):
+        for j in range(Rshape[1]):
+            radiusazimuthimage[i, j] = np.where(
+                np.logical_and(
+                    rphi_filtered_points_array[:, 3] == i,
+                    rphi_filtered_points_array[:, 4] == j,
+                ),
+                rphi_filtered_points_array[:, 2],
+                0,
+            ).sum()
+    return radiusazimuthimage

py4DSTEM/process/phase/parallax.py

@@ -884,15 +884,17 @@ def guess_common_aberrations(
         sampling = 1 / (
             np.array(self._reciprocal_sampling) * self._region_of_interest_shape
         )
-        aberrations_basis, aberrations_basis_du, aberrations_basis_dv = (
-            calculate_aberration_gradient_basis(
-                aberrations_mn,
-                sampling,
-                self._region_of_interest_shape,
-                self._wavelength,
-                rotation_angle=np.deg2rad(rotation_angle_deg),
-                xp=xp,
-            )
+        (
+            aberrations_basis,
+            aberrations_basis_du,
+            aberrations_basis_dv,
+        ) = calculate_aberration_gradient_basis(
+            aberrations_mn,
+            sampling,
+            self._region_of_interest_shape,
+            self._wavelength,
+            rotation_angle=np.deg2rad(rotation_angle_deg),
+            xp=xp,
         )
 
         # shifts
@@ -2432,7 +2434,6 @@ def score_CTF(coefs):
 
             # Plot the measured/fitted shifts comparison
             if plot_BF_shifts_comparison:
-
                 fitted_shifts = (
                     xp.tensordot(gradients, xp.array(self._aberrations_coefs), axes=1)
                     .reshape((2, -1))
@@ -3055,7 +3056,6 @@ def show_shifts(
         shifts = shifts_px * scale_arrows * xp.array(self._reciprocal_sampling)
 
         if plot_rotated_shifts and hasattr(self, "rotation_Q_to_R_rads"):
-
             if figax is None:
                 figsize = kwargs.pop("figsize", (8, 4))
                 fig, ax = plt.subplots(1, 2, figsize=figsize)

py4DSTEM/process/phase/xray_magnetic_ptychography.py

@@ -892,7 +892,6 @@ def _gradient_descent_adjoint(
 
         match (self._recon_mode, self._active_measurement_index):
             case (0, 0) | (1, 0):  # reverse
-
                 magnetic_conj = xp.exp(1.0j * xp.conj(object_patches[1]))
 
                 probe_magnetic_abs = xp.abs(shifted_probes * magnetic_conj)
@@ -930,7 +929,6 @@ def _gradient_descent_adjoint(
                 )
 
                 if not fix_probe:
-
                     electrostatic_magnetic_abs = xp.abs(
                         electrostatic_conj * magnetic_conj
                     )
@@ -962,7 +960,6 @@ def _gradient_descent_adjoint(
                     )
 
             case (0, 1) | (1, 2) | (2, 1):  # forward
-
                 magnetic_conj = xp.exp(-1.0j * xp.conj(object_patches[1]))
 
                 probe_magnetic_abs = xp.abs(shifted_probes * magnetic_conj)
@@ -992,7 +989,6 @@ def _gradient_descent_adjoint(
                 )
 
                 if not fix_probe:
-
                     electrostatic_magnetic_abs = xp.abs(
                         electrostatic_conj * magnetic_conj
                     )
@@ -1024,7 +1020,6 @@ def _gradient_descent_adjoint(
                     )
 
             case (1, 1) | (2, 0):  # neutral
-
                 probe_abs = xp.abs(shifted_probes)
                 probe_normalization = self._sum_overlapping_patches_bincounts(
                     probe_abs**2,
@@ -1047,7 +1042,6 @@ def _gradient_descent_adjoint(
                 )
 
                 if not fix_probe:
-
                     electrostatic_abs = xp.abs(electrostatic_conj)
                     electrostatic_normalization = xp.sum(
                         electrostatic_abs**2,