Make optics/fluorescence compatible with torch

deeptrack/image.py

@@ -59,11 +59,11 @@ class is central to DeepTrack2, acting as a container for numerical data
 - `pad_image_to_fft(image, axes)`
 
     pad_image_to_fft(
-        image: Image | np.ndarray,
+        image: Image | np.ndarray | torch.tensor,
         axes: Iterable[int] = (0, 1),
-    ) -> Image | np.ndarray
+    ) -> Image | np.ndarray | torch.tensor
 
-    Pads an image to optimize Fast Fourier Transform (FFT) performance.
+    Pad an image to optimize Fast Fourier Transform (FFT) performance.
 
 Examples
 --------
@@ -96,11 +96,17 @@ class is central to DeepTrack2, acting as a container for numerical data
 import operator as ops
 from typing import Any, Callable, Iterable
 
+import array_api_compat as apc
 import numpy as np
+from numpy.typing import NDArray
 
+from deeptrack.backend import config, TORCH_AVAILABLE, xp
 from deeptrack.properties import Property
 from deeptrack.types import NumberLike
 
+if TORCH_AVAILABLE:
+    import torch
+
 
 #TODO ***??*** revise _binary_method - typing, docstring, unit test
 def _binary_method(
@@ -1694,28 +1700,27 @@ def coerce(
 _FASTEST_SIZES = np.sort(_FASTEST_SIZES)
 
 
-#TODO ***??*** revise pad_image_to_fft - typing, docstring, unit test
 def pad_image_to_fft(
-    image: Image | np.ndarray | np.ndarray,
+    image: Image | NDArray | torch.Tensor,
     axes: Iterable[int] = (0, 1),
-) -> Image | np.ndarray:
-    """Pads an image to optimize Fast Fourier Transform (FFT) performance.
+) -> Image | NDArray | torch.Tensor:
+    """Pad an image to optimize Fast Fourier Transform (FFT) performance.
 
     This function pads an image by adding zeros to the end of specified axes
     so that their lengths match the nearest larger size in `_FASTEST_SIZES`.
     These sizes are selected to optimize FFT computations.
 
     Parameters
     ----------
-    image: Image | np.ndarray
+    image: Image | np.ndarray | torch.tensor
         The input image to pad. It should be an instance of the `Image` class
         or any array-like structure compatible with FFT operations.
     axes: Iterable[int], optional
         The axes along which to apply padding. Defaults to `(0, 1)`.
 
     Returns
     -------
-    Image | np.ndarray
+    Image | np.ndarray | torch.tensor
         The padded image with dimensions optimized for FFT performance.
 
     Raises
@@ -1725,30 +1730,37 @@ def pad_image_to_fft(
 
     Examples
     --------
-    >>> import numpy as np
     >>> from deeptrack.image import Image, pad_image_to_fft
 
     Pad an Image object:
-
-    >>> img = Image(np.zeros((7, 13)))
+    >>> import numpy as np
+    >>>
+    >>> img = Image(np.ones((7, 13)))
     >>> padded_img = pad_image_to_fft(img)
     >>> print(padded_img.shape)
     (8, 16)
 
     Pad a NumPy array:
-
-    >>> img = np.zeros((5, 11)))
+    >>> img = np.ones((5, 11))
     >>> padded_img = pad_image_to_fft(img)
     >>> print(padded_img.shape)
     (6, 12)
 
+    Pad a PyTorch tensor:
+    >>> import torch
+    >>>
+    >>> img = torch.ones(7, 11)
+    >>> padded_img = pad_image_to_fft(img)
+    >>> print(padded_img.shape)
+    (8, 12)
+
     """
 
     def _closest(
         dim: int,
     ) -> int:
 
-        # Returns the smallest value frin _FASTEST_SIZES larger than dim.
+        # Return the smallest value from _FASTEST_SIZES that is >= dim.
         for size in _FASTEST_SIZES:
             if size >= dim:
                 return size
@@ -1763,7 +1775,18 @@ def _closest(
         new_shape[axis] = _closest(new_shape[axis])
 
     # Calculate the padding for each axis.
-    pad_width = [(0, increase) for increase in np.array(new_shape) - image.shape]
+    pad_width = [
+        (0, increase)
+        for increase in np.array(new_shape) - np.array(image.shape)
+    ]
+
+    # Apply zero-padding with torch.nn.functional.pad if the input is a
+    # PyTorch tensor
+    if apc.is_torch_array(image):
+        pad = []
+        for before, after in reversed(pad_width):
+            pad.extend([before, after])
+        return torch.nn.functional.pad(image, pad, mode="constant", value=0)
 
-    # Pad the image using constant mode (add zeros).
+    # Apply zero-padding with np.pad if the input is a NumPy array or an Image
     return np.pad(image, pad_width, mode="constant")