fixing abtem bug, adding conversion to polar property

py4dstem · Dec 21, 2023 · 3eda4ab · 3eda4ab
1 parent c4654db
commit 3eda4ab
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 4 deletions.
diff --git a/py4DSTEM/process/phase/iterative_parallax.py b/py4DSTEM/process/phase/iterative_parallax.py
@@ -160,12 +160,12 @@ def to_h5(self, group):
                 data=self._asnumpy(self._recon_BF_subpixel_aligned),
             )
 
-        if hasattr(self, "aberration_dict"):
+        if hasattr(self, "aberration_dict_cartesian"):
             self.metadata = Metadata(
                 name="aberrations_metadata",
                 data={
                     v["aberration name"]: v["value [Ang]"]
-                    for k, v in self.aberration_dict.items()
+                    for k, v in self.aberration_dict_cartesian.items()
                 },
             )
 
@@ -2457,7 +2457,7 @@ def score_CTF(coefs):
 
             fig.tight_layout()
 
-        self.aberration_dict = {
+        self.aberration_dict_cartesian = {
             tuple(self._aberrations_mn[a0]): {
                 "aberration name": _aberration_names.get(
                     tuple(self._aberrations_mn[a0, :2]), "-"
@@ -3160,3 +3160,26 @@ def object_cropped(self):
                 )
             else:
                 return self._crop_padded_object(self._recon_BF)
+
+    @property
+    def aberration_dict_polar(self):
+        """converts cartesian aberration dictionary to the polar convention used in ptycho"""
+        polar_dict = {}
+        unique_aberrations = np.unique(self._aberrations_mn[:, :2], axis=0)
+        aberrations_dict = self.aberration_dict_cartesian
+
+        for aberration_order in unique_aberrations:
+            m, n = aberration_order
+            modulus_name = "C" + str(m) + str(n)
+
+            if n != 0:
+                value_a = aberrations_dict[(m, n, 0)]["value [Ang]"]
+                value_b = aberrations_dict[(m, n, 1)]["value [Ang]"]
+                polar_dict[modulus_name] = np.sqrt(value_a**2 + value_b**2)
+
+                argument_name = "phi" + str(m) + str(n)
+                polar_dict[argument_name] = np.arctan2(value_b, value_a) / n
+            else:
+                polar_dict[modulus_name] = aberrations_dict[(m, n, 0)]["value [Ang]"]
+
+        return polar_dict
diff --git a/py4DSTEM/process/phase/utils.py b/py4DSTEM/process/phase/utils.py
@@ -413,7 +413,8 @@ def polar_coordinates(self, x, y):
         """Calculate a polar grid for a given Cartesian grid."""
         xp = self._xp
         alpha = xp.sqrt(x.reshape((-1, 1)) ** 2 + y.reshape((1, -1)) ** 2)
-        phi = xp.arctan2(x.reshape((-1, 1)), y.reshape((1, -1)))
+        # phi = xp.arctan2(x.reshape((-1, 1)), y.reshape((1, -1))) # bug in abtem-legacy and py4DSTEM<=0.14.9
+        phi = xp.arctan2(y.reshape((1, -1)), x.reshape((-1, 1)))
         return alpha, phi
 
     def build(self):