Update of noises (#405)

* update of background noise

* update of background noises unittest

* update noises background

* unit testing now iwth BackendTestBase, for both numpy and torch.

* Added docs and torch implementation for Poisson

* modified  get function for Poisson

* docs polishing

* Implemented Gaussian and Complex Gaussian

* Docs polishing

* Torch tests for Gaussian & Complex Gaussian

* syntax

* Implemented Feedback from Mirja

* docs

* update background

* Implemented second round feedback from Mirja

* update noises

* Update test_noises.py

* imports

* Update test_noises.py

* u

* u

* u

* xp.asarray

* rest of the functions

* removed Image wrapper for Poisson

* u

* complex

* u

* Update test_noises.py

* Update test_noises.py

* whitespace

* added checks for array types

* update test noises

---------

Co-authored-by: Alex <95913221+Pwhsky@users.noreply.github.com>

Author: mirjagranfors

deeptrack/noises.py

@@ -1,5 +1,4 @@
-"""
-Features for introducing noise to images.
+"""Features for introducing noise to images.
 
 This module provides classes to add various types of noise to images, 
 including constant offsets, Gaussian noise, and Poisson-distributed noise.
@@ -63,21 +62,59 @@ class Noise(Feature):
     """Base abstract noise class."""
 
 
-#TODO ***MG*** revise Background - torch, typing, docstring, unit test
 class Background(Noise):
-    """Adds a constant value to an image
+    """Add a constant value to an image.
 
     Parameters
     ----------
-    offset : float
-        The value to add to the image
+    offset: float
+        The value to add to the image.
+    **kwargs: Any
+        Additional keyword arguments passed to the parent `Noise` class.
+
+    Methods
+    -------
+    get(
+        image: np.ndarray, torch.Tensor, or Image,
+        offset: float,
+        **kwargs,
+    ) -> np.ndarray, torch.Tensor, or Image
+        Adds the constant offset to the input image.
+
+    Examples
+    --------
+    >>> import deeptrack as dt
+
+    Create an input image with zeros:
+    >>> import numpy as np
+    >>>
+    >>> input_image = np.zeros((2,2))
+
+    Define the Background noise feature with offset 0.5:
+    >>> noise = dt.Background(offset=0.5)
+
+    Apply the noise to the input image and print the resulting image:
+    >>> output_image = noise.resolve(input_image)
+    >>> print(output_image)
+    [[0.5 0.5]
+    [0.5 0.5]]
+
     """
 
     def __init__(
         self: Background,
         offset: PropertyLike[float],
         **kwargs: Any,
     ):
+        """Initialize the Background noise feature.
+
+        Parameters
+        ----------
+        offset: PropertyLike[float]
+            The constant value to be added to the image.
+        **kwargs: Any
+            Additional arguments passed to the parent `Noise` class.
+        """
         super().__init__(offset=offset, **kwargs)
 
     def get(
@@ -86,23 +123,76 @@ def get(
         offset: float,
         **kwargs: Any,
     ) -> NDArray[Any] | torch.Tensor | Image:
+        """Add the given offset to the image.
 
-        return image + offset
+        Parameters
+        ----------
+        image: np.ndarray, torch.Tensor, or Image
+            The input image.
+        offset: float
+            The value to add to the image.
 
+        Returns
+        -------
+        np.ndarray, torch.Tensor, or Image
+            The image with offset added.
+        """
+
+        return image + offset
 
 Offset = Background
 
 
-#TODO ***JH*** revise Gaussian - torch, typing, docstring, unit test
 class Gaussian(Noise):
-    """Adds IID Gaussian noise to an image.
+    """Add IID Gaussian noise to an image.
+
+    Gaussian noise is sampled from a Gaussian distribution and added pixel-wise
+    to the input image.
 
     Parameters
     ----------
-    mu : float
+    mu: float
         The mean of the Gaussian distribution.
-    sigma : float
+    sigma: float
         The standard deviation of the Gaussian distribution.
+
+    Notes
+    -----
+    If the backend is NumPy, the calculations use NumPy-compatible functions,
+    and the output will be a np.array. If the backend is PyTorch, the
+    calculations use PyTorch-compatible functions, and the output will be a
+    torch.Tensor.
+
+    Methods
+    -------
+    get(
+        image: np.ndarray, torch.Tensor, or Image,
+        snr: float,
+        background: float,
+        max_val: float, optional,
+        **kwargs,
+    ) -> np.ndarray, torch.Tensor, or Image
+        Returns an image with Gaussian noise added.
+
+    Examples
+    --------
+    Add Gaussian noise to an image.
+    >>> import deeptrack as dt
+
+    Create an input image with constant values:
+    >>> import numpy as np
+    >>>
+    >>> input_image = np.ones((2,2)) * 3
+    
+    Define the Gaussian noise feature with mean 1 and standard deviation 0.1:
+    >>> noise = dt.Gaussian(mu=1, sigma=0.1)
+
+    Apply the noise to the input image and print the resulting image:
+    >>> output_image = noise.resolve(input_image)
+    >>> print(output_image)
+    [[4.01965863 4.20688642]
+    [4.02184982 3.87875873]]
+
     """
 
     def __init__(
@@ -111,7 +201,6 @@ def __init__(
         sigma: PropertyLike[float] = 1,
         **kwargs: Any,
     ):
-
         super().__init__(mu=mu, sigma=sigma, **kwargs)
 
     def get(
@@ -122,21 +211,69 @@ def get(
         **kwargs: Any,
     ) -> NDArray[Any] | torch.Tensor | Image:
 
-        noisy_image = mu + image + np.random.randn(*image.shape) * sigma
+        # For a Numpy backend.
+        if self.get_backend() == "numpy":
+            noisy_image = mu + image + np.random.randn(*image.shape) * sigma
+
+        # For a Torch backend.
+        elif self.get_backend() == "torch":
+            noisy_image = mu + image + torch.randn(*image.shape) * sigma
 
         return noisy_image
 
 
-#TODO ***JH*** revise ComplexGaussian - torch, typing, docstring, unit test
 class ComplexGaussian(Noise):
-    """Adds complex-valued IID Gaussian noise to an image.
+    """Add complex-valued IID Gaussian noise to an image.
+
+    Complex Gaussian noise is sampled by combining two independent Gaussian
+    distributions for real and imaginary values and is then added pixel-wise
+    to the input image.
 
     Parameters
     ----------
-    mu : float
+    mu: float
         The mean of the Gaussian distribution.
-    sigma : float
+    sigma: float
         The standard deviation of the Gaussian distribution.
+
+    Notes
+    -----
+    If the backend is NumPy, the calculations use NumPy-compatible functions,
+    and the output will be a np.array. If the backend is PyTorch, the
+    calculations use PyTorch-compatible functions, and the output will be a
+    torch.Tensor.
+
+    Methods
+    -------
+    get(
+        image: np.ndarray, torch.Tensor, or Image,
+        snr: float,
+        background: float,
+        max_val: float, optional,
+        **kwargs,
+    ) -> np.ndarray, torch.Tensor, or Image
+        Returns an image with complex Gaussian noise added.
+
+    Examples
+    --------
+    Add complex Gaussian noise to an image.
+
+    >>> import deeptrack as dt
+
+    Create an input image with constant values:
+    >>> import numpy as np
+    >>>
+    >>> input_image = np.ones((2,2)) * 3
+    
+    Define the Gaussian noise feature with mean 1 and standard deviation 0.1:
+    >>> noise = dt.ComplexGaussian(mu=1, sigma=0.1)
+
+    Apply the noise to the input image and print the resulting image:
+    >>> output_image = noise.resolve(input_image)
+    >>> print(output_image)
+    [[3.79975648-0.06967551j 4.09943404+0.06499738j]
+    [3.99886747-0.23549974j 4.15725117-0.07847024j]]
+
     """
 
     def __init__(
@@ -145,7 +282,6 @@ def __init__(
         sigma: PropertyLike[float] = 1,
         **kwargs: Any,
     ):
-
         super().__init__(mu=mu, sigma=sigma, **kwargs)
 
     def get(
@@ -156,28 +292,78 @@ def get(
         **kwargs: Any,
     ) -> NDArray[Any] | torch.Tensor | Image:
 
-        real_noise = np.random.randn(*image.shape)
-        imag_noise = np.random.randn(*image.shape) * 1j
-        noisy_image = mu + image + (real_noise + imag_noise) * sigma
- 
+        # For a Numpy backend.
+        if self.get_backend() == "numpy":
+            real_noise = np.random.randn(*image.shape)
+            imag_noise = np.random.randn(*image.shape) * 1j
+            noisy_image = mu + image + (real_noise + imag_noise) * sigma
+
+        # For a Torch backend.
+        elif self.get_backend() == "torch":
+            real_noise = torch.randn(*image.shape)
+            imag_noise = torch.randn(*image.shape) * 1j
+            noisy_image = mu + image + (real_noise + imag_noise) * sigma
+
         return noisy_image
 
 
-#TODO ***AL*** revise Poisson - torch, typing, docstring, unit test
 class Poisson(Noise):
-    """Adds Poisson-distributed noise to an image.
+    """Add Poisson-distributed noise to an image.
+
+    Poisson noise is sampled and added pixel-wise depending on the
+    intensity of the pixel in the original image to achieve a desired
+    signal-to-noise ratio `snr`. 
 
     Parameters
     ----------
-    snr : float
+    snr: float
         Signal-to-noise ratio of the final image. The signal is determined
         by the peak value of the image.
-    background : float
+    background: float
         Value to be be used as the background. This is used to calculate the
         signal of the image.
-    max_val : float, optional
+    max_val: float, optional
         Maximum allowable value to prevent overflow in noise computation.
         Default is 1e8.
+
+    Notes
+    -----
+    If the backend is NumPy, the calculations use NumPy-compatible functions,
+    and the output will be a np.array. If the backend is PyTorch, the
+    calculations use PyTorch-compatible functions, and the output will be a
+    torch.Tensor.
+
+    Methods
+    -------
+    get(
+        image: np.ndarray, torch.Tensor, or Image,
+        snr: float,
+        background: float,
+        max_val: float, optional,
+        **kwargs,
+    ) -> np.ndarray, torch.Tensor, or Image
+        Returns an image with Poisson noise added.
+
+    Examples
+    --------
+    Add Poisson noise to an image.
+
+    >>> import deeptrack as dt
+
+    Create an input image with ones:
+    >>> import numpy as np
+    >>>
+    >>> input_image = np.ones((2,2))
+    
+    Define the Poisson noise feature with a low SNR:
+    >>> noise = dt.Poisson(snr=1)
+
+    Apply the noise to the input image and print the resulting image:
+    >>> output_image = noise.resolve(input_image)
+    >>> print(output_image)
+    [[2. 1.]
+    [0. 4.]]
+
     """
 
     def __init__(
@@ -188,7 +374,6 @@ def __init__(
         max_val: PropertyLike[float] = 1e8,
         **kwargs,
     ):
-
         super().__init__(
             *args,
             snr=snr,
@@ -206,18 +391,43 @@ def get(
         **kwargs: Any,
     ) -> NDArray[Any] | torch.Tensor | Image:
 
-        image[image < 0] = 0
-        image_max = np.max(image)
-        peak = np.abs(image_max - background)
-
-        rescale = snr ** 2 / peak ** 2
-        rescale = np.clip(rescale, 1e-10, max_val / np.abs(image_max))
-        try:
-            noisy_image = Image(np.random.poisson(image * rescale) / rescale)
-            noisy_image.merge_properties_from(image)  # TODO Should only be done if input is Image!
-            return noisy_image
-        except ValueError:
-            raise ValueError(
-                "NumPy poisson function errored due to too large value. "
-                "Set max_val in dt.Poisson to a lower value to fix."
+        # For a numpy backend.
+        if self.get_backend() == "numpy":
+            image[image < 0] = 0
+            image_max = np.max(image)
+            peak = np.abs(image_max - background)
+
+            rescale = snr ** 2 / peak ** 2
+            rescale = np.clip(
+                rescale, 1e-10, max_val / np.abs(image_max)
+            )
+            try:
+                noisy_image = Image(
+                    np.random.poisson(image * rescale) / rescale
+                )
+                noisy_image.merge_properties_from(image)
+                return noisy_image
+            except ValueError:
+                raise ValueError(
+                    "NumPy poisson function errored due to too large value. "
+                    "Set max_val in dt.Poisson to a lower value to fix."
+                )
+
+        # For a Torch backend.
+        elif self.get_backend() == "torch":
+            image = torch.clamp(image, min=0)
+            image_max = torch.max(image)
+            peak = torch.abs(image_max - background)
+
+            rescale = snr ** 2 / peak ** 2
+            rescale = torch.clamp(
+                rescale, min=1e-10, max=max_val / torch.abs(image_max)
             )
+            try:
+                noisy_image = torch.poisson(image * rescale) / rescale
+                return noisy_image
+            except ValueError:
+                raise ValueError(
+                    "Torch Poisson function errored due to too large value. "
+                    "Set max_val in dt.Poisson to a lower value to fix."
+                )