Torch-native PDF pipeline + origin finder improvements

src/quantem/core/datastructures/polar4dstem.py

@@ -1,6 +1,7 @@
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Self
 
 import numpy as np
+import torch
 from numpy.typing import NDArray
 from scipy.ndimage import map_coordinates
 
@@ -11,17 +12,19 @@
 
 
 class Polar4dstem(Dataset4d):
-    """4D-STEM dataset in polar coordinates (scan_y, scan_x, phi, r)."""
+    """4D-STEM dataset in polar coordinates (scan_row, scan_col, phi, r_pix)."""
 
     def __init__(
         self,
-        array: NDArray | Any,
-        name: str,
-        origin: NDArray | tuple | list | float | int,
-        sampling: NDArray | tuple | list | float | int,
-        units: list[str] | tuple | list,
+        array: NDArray | None = None,
+        tensor: torch.Tensor | None = None,
+        name: str = "",
+        origin: NDArray | tuple | list | float | int | None = None,
+        sampling: NDArray | tuple | list | float | int | None = None,
+        units: list[str] | tuple | list | None = None,
         signal_units: str = "arb. units",
         metadata: dict | None = None,
+        origin_array: NDArray | None = None,
         _token: object | None = None,
     ):
         if metadata is None:
@@ -32,6 +35,7 @@ def __init__(
             "polar_radial_step",
             "polar_num_annular_bins",
             "polar_two_fold_rotation_symmetry",
+            "polar_ellipticity",
             "polar_origin_row",
             "polar_origin_col",
             "polar_ellipse_params",
@@ -41,6 +45,7 @@ def __init__(
                 metadata[k] = None
         super().__init__(
             array=array,
+            tensor=tensor,
             name=name,
             origin=origin,
             sampling=sampling,
@@ -49,45 +54,30 @@ def __init__(
             metadata=metadata,
             _token=_token,
         )
+        self.origin_array = origin_array
 
     @classmethod
     def from_array(
         cls,
-        array: NDArray | Any,
+        array: NDArray,
         name: str | None = None,
         origin: NDArray | tuple | list | float | int | None = None,
         sampling: NDArray | tuple | list | float | int | None = None,
         units: list[str] | tuple | list | None = None,
         signal_units: str = "arb. units",
         metadata: dict | None = None,
-    ) -> "Polar4dstem":
-        """
-        Create a Polar4dstem from a 4D array shaped (scan_y, scan_x, phi, r).
-
-        Parameters
-        ----------
-        array : NDArray
-            Polar data, shape (scan_y, scan_x, n_phi, n_r).
-        name : str, optional
-            Dataset name, by default "Polar 4D-STEM dataset".
-        origin : NDArray | tuple | list | float | int, optional
-            Origin of each axis in calibrated units, by default zeros.
-        sampling : NDArray | tuple | list | float | int, optional
-            Sampling of each axis, by default ones.
-        units : list[str] | tuple | list, optional
-            Units of each axis, by default ["pixels", "pixels", "deg", "pixels"].
-        signal_units : str, optional
-            Units of the array values, by default "arb. units".
-        metadata : dict, optional
-            Metadata. Missing "polar_*" keys are set to None.
-
-        Returns
-        -------
-        Polar4dstem
-        """
+    ) -> Self:
+        if isinstance(array, torch.Tensor):
+            raise TypeError(
+                f"Polar4dstem.from_array requires a numpy array, got {type(array).__name__}. "
+                "Use Polar4dstem.from_tensor for torch tensors."
+            )
         array = np.asarray(array)
         if array.ndim != 4:
-            raise ValueError("Polar4dstem.from_array expects a 4D array.")
+            raise ValueError(
+                f"Found array with shape: {array.shape}. "
+                "Polar4dstem.from_array expects a 4D array."
+            )
         if origin is None:
             origin = np.zeros(4, dtype=float)
         if sampling is None:
@@ -107,14 +97,52 @@ def from_array(
             _token=cls._token,
         )
 
+    @classmethod
+    def from_tensor(
+        cls,
+        tensor: torch.Tensor,
+        name: str | None = None,
+        origin: NDArray | tuple | list | float | int | None = None,
+        sampling: NDArray | tuple | list | float | int | None = None,
+        units: list[str] | tuple | list | None = None,
+        signal_units: str = "arb. units",
+        metadata: dict | None = None,
+    ) -> Self:
+        """Create a Polar4dstem from a torch tensor."""
+        if not isinstance(tensor, torch.Tensor):
+            raise TypeError(
+                f"Polar4dstem.from_tensor requires torch.Tensor, got {type(tensor).__name__}."
+            )
+        if tensor.ndim != 4:
+            raise ValueError(
+                f"Found tensor with shape: {tuple(tensor.shape)}. "
+                "Polar4dstem.from_tensor expects a 4D tensor."
+            )
+        if origin is None:
+            origin = np.zeros(4, dtype=float)
+        if sampling is None:
+            sampling = np.ones(4, dtype=float)
+        if units is None:
+            units = ["pixels", "pixels", "deg", "pixels"]
+        if metadata is None:
+            metadata = {}
+        return cls(
+            tensor=tensor,
+            name=name if name is not None else "Polar 4D-STEM dataset",
+            origin=origin,
+            sampling=sampling,
+            units=units,
+            signal_units=signal_units,
+            metadata=metadata,
+            _token=cls._token,
+        )
+
     @property
     def n_phi(self) -> int:
-        """Number of azimuthal bins (axis 2)."""
         return int(self.shape[2])
 
     @property
     def n_r(self) -> int:
-        """Number of radial bins (axis 3)."""
         return int(self.shape[3])
 
 

src/quantem/core/utils/filter.py

@@ -135,6 +135,17 @@ def gaussian_filter_2d(
     return img.squeeze_(0).squeeze_(0)  # Make 2D again
 
 
+def gaussian_filter_1d(
+    arr: torch.Tensor, kernel_1d: torch.Tensor
+) -> torch.Tensor:  # Replicate-padded torch alternative to ``scipy.ndimage.gaussian_filter1d``
+    padding = len(kernel_1d) // 2  # Ensure that signal size does not change
+    arr = arr.unsqueeze(0).unsqueeze_(0)  # Make copy, make 3D for ``conv1d()``
+    # Replicate edge values so the output has no zero-padded ringing at boundaries
+    arr = torch.nn.functional.pad(arr, (padding, padding), mode="replicate")
+    arr = torch.nn.functional.conv1d(arr, weight=kernel_1d.view(1, 1, -1))
+    return arr.squeeze_(0).squeeze_(0)  # Make 1D again
+
+
 def gaussian_filter_2d_stack(stack: torch.Tensor, kernel_1d: torch.Tensor) -> torch.Tensor:
     """
     Apply 2D Gaussian blur to each slice stack[:, i, :] in a vectorized way.

src/quantem/diffraction/__init__.py

@@ -1,3 +1,4 @@
+from quantem.diffraction.polar import PairDistributionFunction as PairDistributionFunction
 from quantem.diffraction.bragg_vectors import BraggVectors as BraggVectors
 from quantem.diffraction.crystal import Crystal as Crystal
 from quantem.diffraction.crystal_map import CrystalMap as CrystalMap