Merge remote-tracking branch 'origin/dev' into show

py4DSTEM/datacube/datacube.py

@@ -1126,11 +1126,11 @@ def get_beamstop_mask(
             name = 'gen_params',
             data = {
                 #'gen_func' : 
-                'threshold' : 0.25,
-                'distance_edge' : 4.0,
-                'include_edges' : True,
+                'threshold' : threshold,
+                'distance_edge' : distance_edge,
+                'include_edges' : include_edges,
                 'name' : "mask_beamstop",
-                'returncalc' : True,
+                'returncalc' : returncalc,
             }
         )
 

py4DSTEM/process/calibration/origin.py

@@ -212,10 +212,11 @@ def get_origin_single_dp(dp, r, rscale=1.2):
 
 def get_origin(
     datacube,
-    r=None,
-    rscale=1.2,
-    dp_max=None,
-    mask=None
+    r = None,
+    rscale = 1.2,
+    dp_max = None,
+    mask = None,
+    fast_center = False,
     ):
     """
     Find the origin for all diffraction patterns in a datacube, assuming (a) there is no
@@ -241,6 +242,8 @@ def get_origin(
         mask (ndarray or None): if not None, should be an (R_Nx,R_Ny) shaped
                     boolean array. Origin is found only where mask==True, and masked
                     arrays are returned for qx0,qy0
+        fast_center: (bool)
+            Skip the center of mass refinement step.
 
     Returns:
         (2-tuple of (R_Nx,R_Ny)-shaped ndarrays): the origin, (x,y) at each scan position
@@ -266,10 +269,13 @@ def get_origin(
         ):
             dp = datacube.data[rx, ry, :, :]
             _qx0, _qy0 = np.unravel_index(
-                np.argmax(gaussian_filter(dp, r)), (datacube.Q_Nx, datacube.Q_Ny)
+                np.argmax(gaussian_filter(dp, r, mode='nearest')), (datacube.Q_Nx, datacube.Q_Ny)
             )
-            _mask = np.hypot(qxx - _qx0, qyy - _qy0) < r * rscale
-            qx0[rx, ry], qy0[rx, ry] = get_CoM(dp * _mask)
+            if fast_center:
+                qx0[rx, ry], qy0[rx, ry] = _qx0, _qy0
+            else:
+                _mask = np.hypot(qxx - _qx0, qyy - _qy0) < r * rscale
+                qx0[rx, ry], qy0[rx, ry] = get_CoM(dp * _mask)
 
     else:
         assert mask.shape == (datacube.R_Nx, datacube.R_Ny)
@@ -290,10 +296,13 @@ def get_origin(
             if mask[rx, ry]:
                 dp = datacube.data[rx, ry, :, :]
                 _qx0, _qy0 = np.unravel_index(
-                    np.argmax(gaussian_filter(dp, r)), (datacube.Q_Nx, datacube.Q_Ny)
+                    np.argmax(gaussian_filter(dp, r, mode='nearest')), (datacube.Q_Nx, datacube.Q_Ny)
                 )
-                _mask = np.hypot(qxx - _qx0, qyy - _qy0) < r * rscale
-                qx0.data[rx, ry], qy0.data[rx, ry] = get_CoM(dp * _mask)
+                if fast_center:
+                    qx0[rx, ry], qy0[rx, ry] = _qx0, _qy0
+                else:
+                    _mask = np.hypot(qxx - _qx0, qyy - _qy0) < r * rscale
+                    qx0.data[rx, ry], qy0.data[rx, ry] = get_CoM(dp * _mask)
             else:
                 qx0.mask, qy0.mask = True, True
 

py4DSTEM/process/diffraction/crystal.py

@@ -4,7 +4,6 @@
 import matplotlib.pyplot as plt
 from fractions import Fraction
 from typing import Union, Optional
-from copy import deepcopy
 from scipy.optimize import curve_fit
 import sys
 
@@ -78,6 +77,7 @@ def __init__(
                 1 number: the lattice parameter for a cubic cell
                 3 numbers: the three lattice parameters for an orthorhombic cell
                 6 numbers: the a,b,c lattice parameters and ɑ,β,ɣ angles for any cell
+                3x3 array: row vectors containing the (u,v,w) lattice vectors.
 
         """
         # Initialize Crystal
@@ -92,42 +92,59 @@ def __init__(
         else:
             raise Exception("Number of positions and atomic numbers do not match")
 
-        # unit cell, as either [a a a 90 90 90], [a b c 90 90 90], or [a b c alpha beta gamma]
+        # unit cell, as one of:
+        # [a a a 90 90 90]
+        # [a b c 90 90 90]
+        # [a b c alpha beta gamma]
         cell = np.asarray(cell, dtype="float_")
         if np.size(cell) == 1:
             self.cell = np.hstack([cell, cell, cell, 90, 90, 90])
         elif np.size(cell) == 3:
             self.cell = np.hstack([cell, 90, 90, 90])
         elif np.size(cell) == 6:
             self.cell = cell
+        elif np.shape(cell)[0] == 3 and np.shape(cell)[1] == 3:
+            self.lat_real = np.array(cell)
+            a = np.linalg.norm(self.lat_real[0,:])
+            b = np.linalg.norm(self.lat_real[1,:])
+            c = np.linalg.norm(self.lat_real[2,:])
+            alpha = np.rad2deg(np.arccos(np.clip(np.sum(
+                self.lat_real[1,:]*self.lat_real[2,:])/b/c,-1,1)))
+            beta  = np.rad2deg(np.arccos(np.clip(np.sum(
+                self.lat_real[0,:]*self.lat_real[2,:])/a/c,-1,1)))
+            gamma = np.rad2deg(np.arccos(np.clip(np.sum(
+                self.lat_real[0,:]*self.lat_real[1,:])/a/b,-1,1)))
+            self.cell = (a,b,c,alpha,beta,gamma)
         else:
-            raise Exception("Cell cannot contain " + np.size(cell) + " elements")
+            raise Exception("Cell cannot contain " + np.size(cell) + " entries")
 
         # pymatgen flag
         self.pymatgen_available = False
 
         # Calculate lattice parameters
         self.calculate_lattice()
-        
+
     def calculate_lattice(self):
-        # calculate unit cell lattice vectors
-        a = self.cell[0]
-        b = self.cell[1]
-        c = self.cell[2]
-        alpha = np.deg2rad(self.cell[3])
-        beta = np.deg2rad(self.cell[4])
-        gamma = np.deg2rad(self.cell[5])
-        f = np.cos(beta) * np.cos(gamma) - np.cos(alpha)
-        vol = a*b*c*np.sqrt(1 \
-            + 2*np.cos(alpha)*np.cos(beta)*np.cos(gamma) \
-            - np.cos(alpha)**2 - np.cos(beta)**2 - np.cos(gamma)**2)
-        self.lat_real = np.array(
-            [
-                [a,               0,                 0],
-                [b*np.cos(gamma), b*np.sin(gamma),   0],
-                [c*np.cos(beta), -c*f/np.sin(gamma), vol/(a*b*np.sin(gamma))],
-            ]
-        )
+
+        if not hasattr(self, 'lat_real'):
+            # calculate unit cell lattice vectors
+            a = self.cell[0]
+            b = self.cell[1]
+            c = self.cell[2]
+            alpha = np.deg2rad(self.cell[3])
+            beta = np.deg2rad(self.cell[4])
+            gamma = np.deg2rad(self.cell[5])
+            f = np.cos(beta) * np.cos(gamma) - np.cos(alpha)
+            vol = a*b*c*np.sqrt(1 \
+                + 2*np.cos(alpha)*np.cos(beta)*np.cos(gamma) \
+                - np.cos(alpha)**2 - np.cos(beta)**2 - np.cos(gamma)**2)
+            self.lat_real = np.array(
+                [
+                    [a,               0,                 0],
+                    [b*np.cos(gamma), b*np.sin(gamma),   0],
+                    [c*np.cos(beta), -c*f/np.sin(gamma), vol/(a*b*np.sin(gamma))],
+                ]
+            )
 
         # Inverse lattice, metric tensors
         self.metric_real = self.lat_real @ self.lat_real.T
@@ -139,6 +156,49 @@ def calculate_lattice(self):
             self.pymatgen_available = True            
         else:
             self.pymatgen_available = False
+
+    def get_strained_crystal(
+        self,
+        exx = 0.0,
+        eyy = 0.0,
+        ezz = 0.0,
+        exy = 0.0,
+        exz = 0.0,
+        eyz = 0.0,
+        deformation_matrix = None,
+        return_deformation_matrix = False,
+        ):
+        """
+        This method returns new Crystal class with strain applied. The directions of (x,y,z)
+        are with respect to the default Crystal orientation, which can be checked with
+        print(Crystal.lat_real) applied to the original Crystal.
+
+        Strains are given in fractional values, so exx = 0.01 is 1% strain along the x direction.
+        """
+
+        # deformation matrix
+        if deformation_matrix is None:
+            deformation_matrix = np.array([
+                [1.0+exx,   1.0*exy,    1.0*exz],
+                [1.0*exy,   1.0+eyy,    1.0*eyz],
+                [1.0*exz,   1.0*eyz,    1.0+ezz],
+            ])
+
+        # new unit cell
+        lat_new = self.lat_real @ deformation_matrix
+
+        # make new crystal class
+        from py4DSTEM.process.diffraction import Crystal
+        crystal_strained = Crystal(
+            positions = self.positions.copy(),
+            numbers = self.numbers.copy(),
+            cell = lat_new,
+        )
+
+        if return_deformation_matrix:
+            return crystal_strained, deformation_matrix
+        else:
+            return crystal_strained
 
 
     def from_CIF(CIF, conventional_standard_structure=True):
@@ -386,13 +446,28 @@ def calculate_structure_factors(
         k_max: float = 2.0,
         tol_structure_factor: float = 1e-4,
         return_intensities: bool = False,
-    ):
+        ):
+
+
         """
         Calculate structure factors for all hkl indices up to max scattering vector k_max
 
-        Args:
-            k_max (numpy float):                max scattering vector to include (1/Angstroms)
-            tol_structure_factor (numpy float): tolerance for removing low-valued structure factors
+        Parameters
+        --------
+        
+        k_max: float                
+            max scattering vector to include (1/Angstroms)
+        tol_structure_factor: float
+            tolerance for removing low-valued structure factors
+        return_intensities: bool
+            return the intensities and positions of all structure factor peaks.
+        
+        Returns
+        --------
+        (q_SF, I_SF)
+            Tuple of the q vectors and intensities of each structure factor.
+        
+
         """
 
         # Store k_max
@@ -425,7 +500,7 @@ def calculate_structure_factors(
         hkl = np.vstack([xa.ravel(), ya.ravel(), za.ravel()])
         # g_vec_all = self.lat_inv @ hkl
         g_vec_all =  (hkl.T @ self.lat_inv).T
-
+                
         # Delete lattice vectors outside of k_max
         keep = np.linalg.norm(g_vec_all, axis=0) <= self.k_max
         self.hkl = hkl[:, keep]