Torch-native PDF pipeline + origin finder improvements

src/quantem/core/datastructures/__init__.py

@@ -2,6 +2,7 @@
 from quantem.core.datastructures.vector import Vector as Vector
 
 from quantem.core.datastructures.dataset4dstem import Dataset4dstem as Dataset4dstem
+from quantem.core.datastructures.polar4dstem import Polar4dstem as Polar4dstem
 from quantem.core.datastructures.dataset4d import Dataset4d as Dataset4d
 from quantem.core.datastructures.dataset3d import Dataset3d as Dataset3d
 from quantem.core.datastructures.dataset2d import Dataset2d as Dataset2d

src/quantem/core/datastructures/polar4dstem.py

@@ -0,0 +1,139 @@
+from typing import Self
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+
+from quantem.core.datastructures.dataset4d import Dataset4d
+
+
+class Polar4dstem(Dataset4d):
+    """4D-STEM dataset in polar coordinates (scan_row, scan_col, phi, r_pix)."""
+
+    def __init__(
+        self,
+        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:
+            metadata = {}
+        mdata_keys_polar = [
+            "polar_radial_min",
+            "polar_radial_max",
+            "polar_radial_step",
+            "polar_num_annular_bins",
+            "polar_two_fold_rotation_symmetry",
+            "polar_ellipticity",
+        ]
+        for k in mdata_keys_polar:
+            if k not in metadata:
+                metadata[k] = None
+        super().__init__(
+            array=array,
+            tensor=tensor,
+            name=name,
+            origin=origin,
+            sampling=sampling,
+            units=units,
+            signal_units=signal_units,
+            metadata=metadata,
+            _token=_token,
+        )
+        self.origin_array = origin_array
+
+    @classmethod
+    def from_array(
+        cls,
+        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,
+    ) -> 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(
+                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:
+            sampling = np.ones(4, dtype=float)
+        if units is None:
+            units = ["pixels", "pixels", "deg", "pixels"]
+        if metadata is None:
+            metadata = {}
+        return cls(
+            array=array,
+            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,
+        )
+
+    @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:
+        return int(self.shape[2])
+
+    @property
+    def n_r(self) -> int:
+        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

@@ -0,0 +1 @@
+from quantem.diffraction.polar import PairDistributionFunction as PairDistributionFunction