Merge branch 'py4dstem:dev' into maclariz-DDF-patch-1

README.md

@@ -1,7 +1,7 @@
 
 > :warning: **py4DSTEM version 0.14 update** :warning: Warning: this is a major update and we expect some workflows to break.  You can still install previous versions of py4DSTEM [as discussed here](#legacyinstall)
 
-> :warning: **Phase retrieval refactor version 0.14.9** :warning: Warning: The phase-retrieval modules in py4DSTEM (DPC, parallax, and ptychography) underwent a major refactor in version 0.14.9 and as such older tutorial notebooks will not work as expected. Notably, class names have been pruned to remove the trailing "Reconstruction" (`DPCReconstruction` -> `DPC` etc.), and regularization functions have dropped the `_iter` suffix (and are instead specified as boolean flags). We are working on updating the tutorial notebooks to reflect these changes. In the meantime, there's some more information in the relevant pull request [here](https://github.com/py4dstem/py4DSTEM/pull/597#issuecomment-1890325568).
+> :warning: **Phase retrieval refactor version 0.14.9** :warning: Warning: The phase-retrieval modules in py4DSTEM (DPC, parallax, and ptychography) underwent a major refactor in version 0.14.9 and as such older tutorial notebooks will not work as expected. Notably, class names have been pruned to remove the trailing "Reconstruction" (`DPCReconstruction` -> `DPC` etc.), and regularization functions have dropped the `_iter` suffix (and are instead specified as boolean flags). See the [updated tutorials](https://github.com/py4dstem/py4DSTEM_tutorials) for more information.
 
 ![py4DSTEM logo](/images/py4DSTEM_logo.png)
 

py4DSTEM/__init__.py

@@ -1,6 +1,9 @@
 from py4DSTEM.version import __version__
 from emdfile import tqdmnd
 
+from importlib.metadata import packages_distributions
+
+is_package_lite = "py4DSTEM-lite" in packages_distributions()["py4DSTEM"]
 
 ### io
 
@@ -52,8 +55,11 @@
     BraggVectorMap,
 )
 
-from py4DSTEM.process import classification
-
+try:
+    from py4DSTEM.process import classification
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
 
 # diffraction
 from py4DSTEM.process.diffraction import Crystal, Orientation
@@ -70,7 +76,11 @@
 # strain
 from py4DSTEM.process.strain.strain import StrainMap
 
-from py4DSTEM.process import wholepatternfit
+try:
+    from py4DSTEM.process import wholepatternfit
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
 
 
 ### more submodules

py4DSTEM/braggvectors/diskdetection.py

@@ -5,12 +5,18 @@
 from scipy.ndimage import gaussian_filter
 
 from emdfile import tqdmnd
+from py4DSTEM import is_package_lite
 from py4DSTEM.braggvectors.braggvectors import BraggVectors
 from py4DSTEM.data import QPoints
 from py4DSTEM.datacube import DataCube
 from py4DSTEM.preprocess.utils import get_maxima_2D
 from py4DSTEM.process.utils.cross_correlate import get_cross_correlation_FT
-from py4DSTEM.braggvectors.diskdetection_aiml import find_Bragg_disks_aiml
+
+try:
+    from py4DSTEM.braggvectors.diskdetection_aiml import find_Bragg_disks_aiml
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
 
 
 def find_Bragg_disks(

py4DSTEM/datacube/virtualimage.py

@@ -5,7 +5,6 @@
 # for bragg virtual imaging methods, goto diskdetection.virtualimage.py
 
 import numpy as np
-import dask.array as da
 from typing import Optional
 import inspect
 
@@ -220,7 +219,9 @@ def get_virtual_image(
                     virtual_image[rx, ry] = np.sum(self.data[rx, ry] * mask)
 
             # dask
-            if dask is True:
+            if dask:
+                import dask.array as da
+
                 # set up a generalized universal function for dask distribution
                 def _apply_mask_dask(self, mask):
                     virtual_image = np.sum(

py4DSTEM/io/filereaders/__init__.py

@@ -1,6 +1,12 @@
+from py4DSTEM import is_package_lite
+from py4DSTEM.io.filereaders.empad import read_empad
 from py4DSTEM.io.filereaders.read_dm import read_dm
 from py4DSTEM.io.filereaders.read_K2 import read_gatan_K2_bin
-from py4DSTEM.io.filereaders.empad import read_empad
 from py4DSTEM.io.filereaders.read_mib import load_mib
-from py4DSTEM.io.filereaders.read_arina import read_arina
+
+try:
+    from py4DSTEM.io.filereaders.read_arina import read_arina
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
 from py4DSTEM.io.filereaders.read_abTEM import read_abTEM

py4DSTEM/io/importfile.py

@@ -7,7 +7,6 @@
 from py4DSTEM.io.filereaders import (
     load_mib,
     read_abTEM,
-    read_arina,
     read_dm,
     read_empad,
     read_gatan_K2_bin,
@@ -90,6 +89,8 @@ def import_file(
     elif filetype == "mib":
         data = load_mib(filepath, mem=mem, binfactor=binfactor, **kwargs)
     elif filetype == "arina":
+        from py4DSTEM.io.filereaders import read_arina
+
         data = read_arina(filepath, mem=mem, binfactor=binfactor, **kwargs)
     elif filetype == "abTEM":
         data = read_abTEM(filepath, mem=mem, binfactor=binfactor, **kwargs)

py4DSTEM/process/__init__.py

@@ -1,9 +1,21 @@
+from py4DSTEM import is_package_lite
 from py4DSTEM.process.polar import PolarDatacube
 from py4DSTEM.process.strain.strain import StrainMap
 
 from py4DSTEM.process import phase
 from py4DSTEM.process import calibration
 from py4DSTEM.process import utils
-from py4DSTEM.process import classification
+
+try:
+    from py4DSTEM.process import classification
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
+
 from py4DSTEM.process import diffraction
-from py4DSTEM.process import wholepatternfit
+
+try:
+    from py4DSTEM.process import wholepatternfit
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc

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
@@ -1309,20 +1353,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,