Updating ACOM, new phase mapping module code, some fixes

py4DSTEM/process/diffraction/crystal.py

@@ -710,41 +710,69 @@ def generate_diffraction_pattern(
         zone_axis_cartesian: Optional[np.ndarray] = None,
         proj_x_cartesian: Optional[np.ndarray] = None,
         foil_normal_cartesian: Optional[Union[list, tuple, np.ndarray]] = None,
-        sigma_excitation_error: float = 0.02,
+        sigma_excitation_error: float = 0.01,
         tol_excitation_error_mult: float = 3,
-        tol_intensity: float = 1e-4,
+        tol_intensity: float = 1e-5,
         k_max: Optional[float] = None,
+        precession_angle_degrees=None,
         keep_qz=False,
         return_orientation_matrix=False,
     ):
         """
-        Generate a single diffraction pattern, return all peaks as a pointlist.
+        Generate a single diffraction pattern, return all peaks as a pointlist. This function performs a
+        kinematical calculation, with optional precession of the beam.
 
-        Args:
-            orientation (Orientation):       an Orientation class object
-            ind_orientation                  If input is an Orientation class object with multiple orientations,
-                                             this input can be used to select a specific orientation.
-
-            orientation_matrix (array):      (3,3) orientation matrix, where columns represent projection directions.
-            zone_axis_lattice (array):        (3,) projection direction in lattice indices
-            proj_x_lattice (array):           (3,) x-axis direction in lattice indices
-            zone_axis_cartesian (array):     (3,) cartesian projection direction
-            proj_x_cartesian (array):        (3,) cartesian projection direction
-
-            foil_normal:                     3 element foil normal - set to None to use zone_axis
-            proj_x_axis (np float vector):   3 element vector defining image x axis (vertical)
-            accel_voltage (float):           Accelerating voltage in Volts. If not specified,
-                                             we check to see if crystal already has voltage specified.
-            sigma_excitation_error (float):  sigma value for envelope applied to s_g (excitation errors) in units of inverse Angstroms
-            tol_excitation_error_mult (float): tolerance in units of sigma for s_g inclusion
-            tol_intensity (np float):        tolerance in intensity units for inclusion of diffraction spots
-            k_max (float):                   Maximum scattering vector
-            keep_qz (bool):                  Flag to return out-of-plane diffraction vectors
-            return_orientation_matrix (bool): Return the orientation matrix
+        TODO - switch from numerical precession to analytic (requires geometry projection).
+        TODO - verify projection geometry for 2D material diffraction.
+
+        Parameters
+        ----------
+
+            orientation (Orientation)
+                an Orientation class object
+            ind_orientation
+                If input is an Orientation class object with multiple orientations,
+                this input can be used to select a specific orientation.
+
+            orientation_matrix (3,3) numpy.array
+                orientation matrix, where columns represent projection directions.
+            zone_axis_lattice (3,) numpy.array
+                projection direction in lattice indices
+            proj_x_lattice (3,) numpy.array
+                x-axis direction in lattice indices
+            zone_axis_cartesian (3,) numpy.array
+                cartesian projection direction
+            proj_x_cartesian (3,) numpy.array
+                cartesian projection direction
+            foil_normal
+                3 element foil normal - set to None to use zone_axis
+            proj_x_axis (3,) numpy.array
+                3 element vector defining image x axis (vertical)
+            accel_voltage (float)
+                Accelerating voltage in Volts. If not specified,
+                we check to see if crystal already has voltage specified.
+            sigma_excitation_error (float)
+                sigma value for envelope applied to s_g (excitation errors) in units of inverse Angstroms
+            tol_excitation_error_mult (float)
+                tolerance in units of sigma for s_g inclusion
+            tol_intensity (numpy float)
+                tolerance in intensity units for inclusion of diffraction spots
+            k_max (float)
+                Maximum scattering vector
+            precession_angle_degrees (float)
+                Precession angle for library calculation.  Set to None for no precession.
+            keep_qz (bool)
+                Flag to return out-of-plane diffraction vectors
+            return_orientation_matrix (bool)
+                Return the orientation matrix
+
+        Returns
+        ----------
+        bragg_peaks (PointList)
+            list of all Bragg peaks with fields [qx, qy, intensity, h, k, l]
+        orientation_matrix (array, optional)
+            3x3 orientation matrix
 
-        Returns:
-            bragg_peaks (PointList):         list of all Bragg peaks with fields [qx, qy, intensity, h, k, l]
-            orientation_matrix (array):      3x3 orientation matrix (optional)
         """
 
         if not (hasattr(self, "wavelength") and hasattr(self, "accel_voltage")):
@@ -779,17 +807,27 @@ def generate_diffraction_pattern(
 
         # Calculate excitation errors
         if foil_normal is None:
-            sg = self.excitation_errors(g)
+            sg = self.excitation_errors(
+                g,
+                precession_angle_degrees=precession_angle_degrees,
+            )
         else:
             foil_normal = (
                 orientation_matrix.T
                 @ (-1 * foil_normal[:, None] / np.linalg.norm(foil_normal))
             ).ravel()
-            sg = self.excitation_errors(g, foil_normal)
+            sg = self.excitation_errors(
+                g,
+                foil_normal=foil_normal,
+                precession_angle_degrees=precession_angle_degrees,
+            )
 
         # Threshold for inclusion in diffraction pattern
         sg_max = sigma_excitation_error * tol_excitation_error_mult
-        keep = np.abs(sg) <= sg_max
+        if precession_angle_degrees is None:
+            keep = np.abs(sg) <= sg_max
+        else:
+            keep = np.min(np.abs(sg), axis=1) <= sg_max
 
         # Maximum scattering angle cutoff
         if k_max is not None:
@@ -799,9 +837,15 @@ def generate_diffraction_pattern(
         g_diff = g[:, keep]
 
         # Diffracted peak intensities and labels
-        g_int = self.struct_factors_int[keep] * np.exp(
-            (sg[keep] ** 2) / (-2 * sigma_excitation_error**2)
-        )
+        if precession_angle_degrees is None:
+            g_int = self.struct_factors_int[keep] * np.exp(
+                (sg[keep] ** 2) / (-2 * sigma_excitation_error**2)
+            )
+        else:
+            g_int = self.struct_factors_int[keep] * np.mean(
+                np.exp((sg[keep] ** 2) / (-2 * sigma_excitation_error**2)),
+                axis=1,
+            )
         hkl = self.hkl[:, keep]
 
         # Intensity tolerance
@@ -1323,20 +1367,55 @@ def excitation_errors(
         self,
         g,
         foil_normal=None,
+        precession_angle_degrees=None,
+        precession_steps=72,
     ):
         """
         Calculate the excitation errors, assuming k0 = [0, 0, -1/lambda].
         If foil normal is not specified, we assume it is [0,0,-1].
+
+        Precession is currently implemented using numerical integration.
         """
-        if foil_normal is None:
-            return (2 * g[2, :] - self.wavelength * np.sum(g * g, axis=0)) / (
-                2 - 2 * self.wavelength * g[2, :]
-            )
+
+        if precession_angle_degrees is None:
+            if foil_normal is None:
+                return (2 * g[2, :] - self.wavelength * np.sum(g * g, axis=0)) / (
+                    2 - 2 * self.wavelength * g[2, :]
+                )
+            else:
+                return (2 * g[2, :] - self.wavelength * np.sum(g * g, axis=0)) / (
+                    2 * self.wavelength * np.sum(g * foil_normal[:, None], axis=0)
+                    - 2 * foil_normal[2]
+                )
+
         else:
-            return (2 * g[2, :] - self.wavelength * np.sum(g * g, axis=0)) / (
-                2 * self.wavelength * np.sum(g * foil_normal[:, None], axis=0)
-                - 2 * foil_normal[2]
+            t = np.deg2rad(precession_angle_degrees)
+            p = np.linspace(
+                0,
+                2.0 * np.pi,
+                precession_steps,
+                endpoint=False,
             )
+            if foil_normal is None:
+                foil_normal = np.array((0.0, 0.0, -1.0))
+
+            k = np.reshape(
+                (-1 / self.wavelength)
+                * np.vstack(
+                    (
+                        np.sin(t) * np.cos(p),
+                        np.sin(t) * np.sin(p),
+                        np.cos(t) * np.ones(p.size),
+                    )
+                ),
+                (3, 1, p.size),
+            )
+
+            term1 = np.sum((g[:, :, None] + k) * foil_normal[:, None, None], axis=0)
+            term2 = np.sum((g[:, :, None] + 2 * k) * g[:, :, None], axis=0)
+            sg = np.sqrt(term1**2 - term2) - term1
+
+            return sg
 
     def calculate_bragg_peak_histogram(
         self,