Merge remote-tracking branch 'py4dstem/dev' into diffraction

py4DSTEM/process/diffraction/crystal_ACOM.py

@@ -2245,6 +2245,7 @@ def save_ang_file(
     pixel_units="px",
     transpose_xy=True,
     flip_x=False,
+    flip_y=False,
 ):
     """
     This function outputs an ascii text file in the .ang format, containing
@@ -2264,8 +2265,10 @@ def save_ang_file(
         nothing
 
     """
-
-    from orix.io.plugins.ang import file_writer
+    try:
+        from orix.io.plugins.ang import file_writer
+    except ImportError:
+        raise Exception("orix failed to import; try pip installing separately")
 
     xmap = self.orientation_map_to_orix_CrystalMap(
         orientation_map,
@@ -2275,6 +2278,7 @@ def save_ang_file(
         return_color_key=False,
         transpose_xy=transpose_xy,
         flip_x=flip_x,
+        flip_y=flip_y,
     )
 
     file_writer(file_name, xmap)
@@ -2288,6 +2292,7 @@ def orientation_map_to_orix_CrystalMap(
     pixel_units="px",
     transpose_xy=True,
     flip_x=False,
+    flip_y=False,
     return_color_key=False,
 ):
     try:
@@ -2313,12 +2318,20 @@ def orientation_map_to_orix_CrystalMap(
 
     import warnings
 
-    # Get orientation matrices
+    # Get orientation matrices and correlation signal (will be used as iq and ci)
     orientation_matrices = orientation_map.matrix[:, :, ind_orientation].copy()
+    corr_values = orientation_map.corr[:, :, ind_orientation].copy()
+
+    # Check for transpose
     if transpose_xy:
         orientation_matrices = np.transpose(orientation_matrices, (1, 0, 2, 3))
+        corr_values = np.transpose(corr_values, (1, 0))
     if flip_x:
         orientation_matrices = np.flip(orientation_matrices, axis=0)
+        corr_values = np.flip(corr_values, axis=0)
+    if flip_y:
+        orientation_matrices = np.flip(orientation_matrices, axis=1)
+        corr_values = np.flip(corr_values, axis=1)
 
     # Convert the orientation matrices into Euler angles
     # suppress Gimbal lock warnings
@@ -2380,8 +2393,8 @@ def fxn():
         y=coords["y"],
         phase_list=PhaseList(phase),
         prop={
-            "iq": orientation_map.corr[:, :, ind_orientation].ravel(),
-            "ci": orientation_map.corr[:, :, ind_orientation].ravel(),
+            "iq": corr_values.ravel(),
+            "ci": corr_values.ravel(),
         },
         scan_unit=pixel_units,
     )

py4DSTEM/process/phase/parallax.py

@@ -2491,13 +2491,25 @@ def aberration_correct(
                 use_CTF_fit = True
 
         if use_CTF_fit:
-            sin_chi = xp.sin(
-                self._calculate_CTF(im.shape, (sx, sy), *self._aberrations_coefs)
-            )
+            even_radial_orders = (self._aberrations_mn[:, 0] % 2) == 1
+            odd_radial_orders = (self._aberrations_mn[:, 0] % 2) == 0
+
+            odd_coefs = self._aberrations_coefs.copy()
+            odd_coefs[even_radial_orders] = 0
+            chi_odd = self._calculate_CTF(im.shape, (sx, sy), *odd_coefs)
+
+            even_coefs = self._aberrations_coefs.copy()
+            even_coefs[odd_radial_orders] = 0
+            chi_even = self._calculate_CTF(im.shape, (sx, sy), *even_coefs)
+
+            if not chi_even.any():  # check if all zeros
+                chi_even = xp.ones_like(chi_even)
+
         else:
-            sin_chi = xp.sin((xp.pi * self._wavelength * self.aberration_C1) * kra2)
+            chi_even = (xp.pi * self._wavelength * self.aberration_C1) * kra2
+            chi_odd = xp.zeros_like(chi_even)
 
-        CTF_corr = xp.sign(sin_chi)
+        CTF_corr = xp.sign(xp.sin(chi_even)) * xp.exp(-1j * chi_odd)
         CTF_corr[0, 0] = 0
 
         # apply correction to mean reconstructed BF image