Add multidimensional drift correction and strip-based alignment - EDS, 4DSTEM prior to arXiv

pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
   "colorspacious",
   "dill",
   "numpy>2",
+  "pandas>=2.2",
   "matplotlib",
   "scipy",
   "tqdm",

src/quantem/core/ml/models/kplanes.py

@@ -525,7 +525,7 @@ def param_keys(self) -> list[str]:
     def extract_tau_state(self) -> torch.Tensor:
         """
         Returns the current raw R^9 matrices (detached copy) so they can be
-        used to initialise a phase-2 model via `load_tau_state`.
+        used to initialize a phase-2 model via `load_tau_state`.
 
         Returns
         -------

src/quantem/core/ml/models/so3params.py

@@ -14,12 +14,12 @@
 
 class SO3ParamQuat(nn.Module):
     """
-    Stores T unit quaternions as learnable parameters in R^4 and normalises
+    Stores T unit quaternions as learnable parameters in R^4 and normalizes
     them on every call to `as_matrix()`.
 
     Quaternion convention: [x, y, z, w]  (scalar-last, same as scipy).
 
-    Initialisation
+    Initialization
     --------------
     "random"  – uniform sampling over SO(3) via Shoemake's method.
     "identity" – all rotations start as the identity (good for fine-tuning).
@@ -105,7 +105,7 @@ def extra_repr(self) -> str:
         return f"T={self.quats.shape[0]}"
 
     # ------------------------------------------------------------------
-    # Initialisers
+    # Initializers
     # ------------------------------------------------------------------
 
     @staticmethod

src/quantem/core/ml/optimizer_mixin.py

@@ -52,7 +52,7 @@ class Adam:
         eps : float
             Term added to the denominator for numerical stability. Default: 1e-8.
         weight_decay : float
-            L2 regularisation penalty. Default: 0.
+            L2 regularization penalty. Default: 0.
         """
 
         lr: float = 1e-3
@@ -97,7 +97,7 @@ class AdamW:
         eps : float
             Term added to the denominator for numerical stability. Default: 1e-8.
         weight_decay : float
-            Decoupled L2 regularisation penalty. Default: 0.
+            Decoupled L2 regularization penalty. Default: 0.
         """
 
         lr: float = 1e-3
@@ -138,7 +138,7 @@ class SGD:
         dampening : float
             Dampening for momentum. Default: 0.
         weight_decay : float
-            L2 regularisation penalty. Default: 0.
+            L2 regularization penalty. Default: 0.
         nesterov : bool
             Enables Nesterov momentum. Default: False.
         """
@@ -266,7 +266,7 @@ class Plateau:
         Parameters
         ----------
         mode : {'min', 'max'}
-            Whether the monitored metric should be minimised or maximised. Default: 'min'.
+            Whether the monitored metric should be minimized or maximized. Default: 'min'.
         min_lr_factor : float
             Sets ``min_lr = min_lr_factor * base_lr`` when ``min_lr`` is not provided.
             Default: 1/20.

src/quantem/core/visualization/show_params.py

@@ -54,7 +54,7 @@ class Norm:
         vmax : float or None
             Explicit maximum for ``"manual"`` interval.
         vcenter : float
-            Centre value for ``"centered"`` interval. Default 0.0.
+            Center value for ``"centered"`` interval. Default 0.0.
         half_range : float or None
             Symmetric half-range for ``"centered"`` interval.
         power : float
@@ -241,7 +241,7 @@ class Scalebar:
         pad_px : float
             Padding between bar and plot edge in image pixels. Default 0.5.
         color : str
-            Bar and label colour. Default ``"white"``.
+            Bar and label color. Default ``"white"``.
         loc : ``"lower right"`` | ``"lower left"`` | ``"upper right"`` | ``"upper left"``
             Anchor location. Default ``"lower right"``.
         fontsize : int

src/quantem/imaging/__init__.py

@@ -1,3 +1,15 @@
-from quantem.imaging.drift import DriftCorrection as DriftCorrection
+"""Imaging tools for scientific image analysis."""
+
+from quantem.imaging.drift import (
+    CorrectionResult as CorrectionResult,
+    DriftCorrection as DriftCorrection,
+    StripPass as StripPass,
+    pair_spectrum_image_references as pair_spectrum_image_references,
+)
+from quantem.imaging.drift.io import (
+    read_emd as read_emd,
+    read_emd_eds as read_emd_eds,
+    read_emd_metadata as read_emd_metadata,
+)
 from quantem.imaging.lattice import Lattice as Lattice
 from quantem.imaging.lattice_visualization import PLOT_REGISTRY as PLOT_REGISTRY