From 89b176c2c6d35cdb1d46c2201296ae2c11b97992 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 16:10:02 +0200
Subject: [PATCH 01/18] added maxpooling docs and unit tests

---
 deeptrack/math.py            | 124 +++++++++++++++++++++++++++++++----
 deeptrack/tests/test_math.py |  16 ++++-
 2 files changed, 126 insertions(+), 14 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 13af56cec..1c99100d1 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1249,15 +1249,16 @@ def __init__(
         super().__init__(np.mean, ksize=ksize, **kwargs)
 
 
-#TODO ***AL*** revise MaxPooling - torch, typing, docstring, unit test
 class MaxPooling(Pool):
     """Apply max pooling to images.
 
-    This class reduces the resolution of an image by dividing it into
-    non-overlapping blocks of size `ksize` and applying the max function to
-    each block. The result is a downsampled image where each pixel value
-    represents the maximum value within the corresponding block of the
-    original image.
+    This class inherits from `Pool` to reduce the resolution of an image by
+    dividing it into non-overlapping blocks of size `ksize` and applying the
+    max function to each block. The result is a downsampled image where
+    each pixel value represents the maximum value within the corresponding
+    block of the original image. If the backend is torch, it will return the
+    output of `torch.nn.functional.max_pool2d` instead.
+
     This is useful for reducing the size of an image while retaining the
     most significant features.
 
@@ -1267,7 +1268,7 @@ class MaxPooling(Pool):
         Size of the pooling kernel.
     cval: number
         Value to pad edges with if necessary. Default 0.
-    func_kwargs: dict
+    **kwargs: dict
         Additional parameters sent to the pooling function.
 
     Examples
@@ -1285,10 +1286,12 @@ class MaxPooling(Pool):
 
     Notes
     -----
-    Calling this feature returns a `np.ndarray` by default. If
-    `store_properties` is set to `True`, the returned array will be
-    automatically wrapped in an `Image` object. This behavior is handled
-    internally and does not affect the return type of the `get()` method.
+    Calling this feature returns a pooled image of the input, it will return
+    either numpy or torch depending on the backend. If
+    `store_properties` is set to `True` and the input is a numpy array,
+    the returned array will be automatically wrapped in an `Image` object.
+    This behavior is handled internally and does not affect the return type
+    of the `get()` method.
 
     """
 
@@ -1299,7 +1302,8 @@ def __init__(
     ):
         """Initialize the parameters for max pooling.
 
-        This constructor initializes the parameters for max pooling.
+        This constructor initializes the parameters for max pooling and checks
+        whether to use the numpy or torch implementation, defaults to numpy.
 
         Parameters
         ----------
@@ -1309,9 +1313,103 @@ def __init__(
             Additional keyword arguments.
 
         """
-
         super().__init__(np.max, ksize=ksize, **kwargs)
 
+    def _get_numpy(
+        self,
+        image: NDArray,
+        ksize: int=3,
+        **kwargs,
+    ):
+        """Method to perform average pooling with the numpy backend enabled.
+
+        Returns the result of the image passed to the scikit image block_reduce
+        function with `np.max()` as the pooling function.
+
+        Parameters
+        ----------
+        image: NDArray
+            Input image to be pooled.
+        ksize: int
+            Kernel size of the pooling operation.
+
+        Returns
+        -------
+        NDArray
+            The pooled image as a `NDArray`.
+            
+        """
+        return utils.safe_call(
+            skimage.measure.block_reduce,
+            image=image,
+            func=self.pooling, # This will be np.mean for this class.
+            block_size=ksize,
+            **kwargs,
+        )
+
+    def _get_torch(
+        self,
+        image: torch.Tensor,
+        ksize: int=3,
+        **kwargs,
+    ):
+        """Method to perform max pooling with the torch backend enabled.
+        
+        Returns the result of the image passed to a torch max
+        pooling layer.
+
+        Parameters
+        ----------
+        image: torch.Tensor
+            Input image to be pooled.
+        ksize: int
+            Kernel size of the pooling operation.
+
+        Returns
+        -------
+        torch.Tensor
+            The pooled image as a `torch.Tensor`.
+
+        """
+
+        return torch.nn.functional.max_pool2d(
+            image,
+            kernel_size=ksize,
+        )
+
+    def get(
+        self,
+        image: NDArray | torch.Tensor,
+        ksize: int=3,
+        **kwargs,
+    ):
+        """Method to perform pooling with either torch or numpy backend.
+        
+        Checks the current backend and chooses the appropriate function to pool
+        the input image, either `_get_torch` or `_get_numpy`.
+
+        Parameters
+        ----------
+        image: NDArray | torch.Tensor
+            Input image to be pooled.
+        ksize: int
+            Kernel size of the pooling operation.
+
+        Returns
+        -------
+        NDArray | torch.Tensor
+            The pooled image as `NDArray` or `torch.Tensor` depending on
+            the backend.
+
+        """
+        if self.get_backend() == "numpy":
+            return self._get_numpy(image, ksize, **kwargs,)
+        elif self.get_backend() == "torch":
+            return self._get_torch(image, ksize, **kwargs,)
+        else:
+            raise NotImplementedError(f"Backend {self.backend} not supported")
+
+
 
 #TODO ***AL*** revise MinPooling - torch, typing, docstring, unit test
 class MinPooling(Pool):
diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 197afb404..811140f5b 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -82,12 +82,26 @@ def test_Blur(self):
         #blurred_image = feature.resolve(input_image)
         #self.assertTrue(xp.all(blurred_image == expected_output))
 
+    def test_MaxPooling(self):
+        input_image = np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=float)
+        feature = math.MaxPooling(ksize=2)
+        pooled_image = feature.resolve(input_image)
+        self.assertTrue(np.all(pooled_image == [[3.5, 5.5]]))
+
 
 # Extending the test and setting the backend to torch
 @unittest.skipUnless(TORCH_AVAILABLE, "PyTorch is not installed.")
 class TestMath_Torch(TestMath_Numpy):
     BACKEND = "torch"
-    pass
+
+    def test_MaxPooling(self):
+        input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
+                                        [5.0, 6.0, 7.0, 8.0] ]]])
+        feature = math.MaxPooling(ksize=2)
+        pooled_image = feature(input_image, ksize=2)
+        expected = torch.tensor([[[[3.5, 5.5]]]])
+        self.assertEqual(pooled_image.shape, expected.shape)
+        self.assertTrue(torch.allclose(pooled_image, expected))
 
 
 class TestMath(unittest.TestCase):

From 2eefbdf3c8045d78292807618efd69cfe0c62c6d Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 30 Jul 2025 16:13:52 +0200
Subject: [PATCH 02/18] Update test_math.py

---
 deeptrack/tests/test_math.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 811140f5b..f04532f41 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -86,7 +86,7 @@ def test_MaxPooling(self):
         input_image = np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=float)
         feature = math.MaxPooling(ksize=2)
         pooled_image = feature.resolve(input_image)
-        self.assertTrue(np.all(pooled_image == [[3.5, 5.5]]))
+        self.assertTrue(np.all(pooled_image == [[6.0, 8.0]]))
 
 
 # Extending the test and setting the backend to torch
@@ -99,7 +99,7 @@ def test_MaxPooling(self):
                                         [5.0, 6.0, 7.0, 8.0] ]]])
         feature = math.MaxPooling(ksize=2)
         pooled_image = feature(input_image, ksize=2)
-        expected = torch.tensor([[[[3.5, 5.5]]]])
+        expected = torch.tensor([[[[6.0, 8.0]]]])
         self.assertEqual(pooled_image.shape, expected.shape)
         self.assertTrue(torch.allclose(pooled_image, expected))
 

From d9d27cfb050d7ea4cd740eb4f880a1c07c07355f Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 6 Aug 2025 18:48:31 +0200
Subject: [PATCH 03/18] Implemented feedback from Mirja

---
 deeptrack/math.py | 39 ++++++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 1c99100d1..bf2dd0a80 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1254,20 +1254,24 @@ class MaxPooling(Pool):
 
     This class inherits from `Pool` to reduce the resolution of an image by
     dividing it into non-overlapping blocks of size `ksize` and applying the
-    max function to each block. The result is a downsampled image where
-    each pixel value represents the maximum value within the corresponding
-    block of the original image. If the backend is torch, it will return the
-    output of `torch.nn.functional.max_pool2d` instead.
+    `max` function to each block. The result is a downsampled image where each
+    pixel value represents the maximum value within the corresponding block of
+    the original image. This is useful for reducing the size of an image while
+    retaining the most significant features.
+
+    If the backend is numpy, the downsampling is performed using
+    `skimage.measure.block_reduce`.
+    If the backend is torch, the downsampling
+    is performed using `torch.nn.functional.max_pool2d`.
+
+ 
+
 
-    This is useful for reducing the size of an image while retaining the
-    most significant features.
 
     Parameters
     ----------
     ksize: int
         Size of the pooling kernel.
-    cval: number
-        Value to pad edges with if necessary. Default 0.
     **kwargs: dict
         Additional parameters sent to the pooling function.
 
@@ -1287,11 +1291,10 @@ class MaxPooling(Pool):
     Notes
     -----
     Calling this feature returns a pooled image of the input, it will return
-    either numpy or torch depending on the backend. If
-    `store_properties` is set to `True` and the input is a numpy array,
-    the returned array will be automatically wrapped in an `Image` object.
-    This behavior is handled internally and does not affect the return type
-    of the `get()` method.
+    either numpy or torch depending on the backend. If `store_properties` is
+    set to `True` and the input is a numpy array, the returned array will be
+    automatically wrapped in an `Image` object. This behavior is handled
+    internally and does not affect the return type of the `get()` method.
 
     """
 
@@ -1342,7 +1345,7 @@ def _get_numpy(
         return utils.safe_call(
             skimage.measure.block_reduce,
             image=image,
-            func=self.pooling, # This will be np.mean for this class.
+            func=self.pooling, # This will be np.max for this class.
             block_size=ksize,
             **kwargs,
         )
@@ -1354,9 +1357,8 @@ def _get_torch(
         **kwargs,
     ):
         """Method to perform max pooling with the torch backend enabled.
-        
-        Returns the result of the image passed to a torch max
-        pooling layer.
+
+        Returns the result of the image passed to a torch max pooling layer.
 
         Parameters
         ----------
@@ -1384,7 +1386,7 @@ def get(
         **kwargs,
     ):
         """Method to perform pooling with either torch or numpy backend.
-        
+
         Checks the current backend and chooses the appropriate function to pool
         the input image, either `_get_torch` or `_get_numpy`.
 
@@ -1410,7 +1412,6 @@ def get(
             raise NotImplementedError(f"Backend {self.backend} not supported")
 
 
-
 #TODO ***AL*** revise MinPooling - torch, typing, docstring, unit test
 class MinPooling(Pool):
     """Apply min pooling to images.

From 460551d2af7a88c31f65148e61cfd84a0bb083a7 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 6 Aug 2025 11:12:47 -0700
Subject: [PATCH 04/18] Added type and shape check

---
 deeptrack/tests/test_math.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index f04532f41..012481fe6 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -78,7 +78,7 @@ def test_Blur(self):
         #input_image = xp.asarray(np.array([[1, 2], [3, 4]], dtype=float))
         #expected_output = xp.asarray(np.array([[1, 1.5], [2, 2.5]]))
 
-        #eature = math.Blur(filter_function=uniform_filter, size=2)
+        #feature = math.Blur(filter_function=uniform_filter, size=2)
         #blurred_image = feature.resolve(input_image)
         #self.assertTrue(xp.all(blurred_image == expected_output))
 
@@ -87,6 +87,7 @@ def test_MaxPooling(self):
         feature = math.MaxPooling(ksize=2)
         pooled_image = feature.resolve(input_image)
         self.assertTrue(np.all(pooled_image == [[6.0, 8.0]]))
+        self.assertEqual(pooled_image.shape, (1, 2))
 
 
 # Extending the test and setting the backend to torch
@@ -102,6 +103,7 @@ def test_MaxPooling(self):
         expected = torch.tensor([[[[6.0, 8.0]]]])
         self.assertEqual(pooled_image.shape, expected.shape)
         self.assertTrue(torch.allclose(pooled_image, expected))
+        self.assertTrue(isinstance(pooled_image, torch.tensor))
 
 
 class TestMath(unittest.TestCase):

From 20a32dde3335fa413a25e68b92b960b6f9e181d8 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 6 Aug 2025 11:15:43 -0700
Subject: [PATCH 05/18] =?UTF-8?q?=C3=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 012481fe6..0293c80d7 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -103,7 +103,7 @@ def test_MaxPooling(self):
         expected = torch.tensor([[[[6.0, 8.0]]]])
         self.assertEqual(pooled_image.shape, expected.shape)
         self.assertTrue(torch.allclose(pooled_image, expected))
-        self.assertTrue(isinstance(pooled_image, torch.tensor))
+        self.assertTrue(isinstance(pooled_image, torch.Tensor))
 
 
 class TestMath(unittest.TestCase):

From 57c0213a15b2024b5b5f7cdb2a3dfca8c4a43472 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 6 Aug 2025 13:49:52 -0700
Subject: [PATCH 06/18] Added shape handling for len(dim) = 2

---
 deeptrack/math.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index bf2dd0a80..05783551d 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1264,10 +1264,6 @@ class MaxPooling(Pool):
     If the backend is torch, the downsampling
     is performed using `torch.nn.functional.max_pool2d`.
 
- 
-
-
-
     Parameters
     ----------
     ksize: int
@@ -1286,7 +1282,7 @@ class MaxPooling(Pool):
     >>> max_pooling = dt.MaxPooling(ksize=8)
     >>> output_image = max_pooling(input_image)
     >>> print(output_image.shape)
-    (8, 8)
+    (4, 4)
 
     Notes
     -----
@@ -1373,6 +1369,15 @@ def _get_torch(
             The pooled image as a `torch.Tensor`.
 
         """
+        # If needed, expand tensor shape
+        if len(image.shape) == 2:
+            expanded_image = image.unsqueeze(0)
+
+            pooled_image = torch.nn.functional.max_pool2d(
+                expanded_image, kernel_size=ksize,
+            )
+            # Remove the expanded dim.
+            return pooled_image.squeeze(0)
 
         return torch.nn.functional.max_pool2d(
             image,

From 199589c105b77a51ab4cc936db1f16d2105049ac Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 6 Aug 2025 13:57:03 -0700
Subject: [PATCH 07/18] Update test_math with len(dim) = 2

---
 deeptrack/tests/test_math.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 0293c80d7..d35936330 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -96,6 +96,7 @@ class TestMath_Torch(TestMath_Numpy):
     BACKEND = "torch"
 
     def test_MaxPooling(self):
+        # (1, 1, 2, 4)
         input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
                                         [5.0, 6.0, 7.0, 8.0] ]]])
         feature = math.MaxPooling(ksize=2)
@@ -104,6 +105,16 @@ def test_MaxPooling(self):
         self.assertEqual(pooled_image.shape, expected.shape)
         self.assertTrue(torch.allclose(pooled_image, expected))
         self.assertTrue(isinstance(pooled_image, torch.Tensor))
+        
+        # (2, 4)
+        input_image = torch.tensor([ [1.0, 2.0, 3.0, 4.0],
+                                        [5.0, 6.0, 7.0, 8.0] ])
+        feature = math.MaxPooling(ksize=2)
+        pooled_image = feature(input_image, ksize=2)
+        expected = torch.tensor([[6.0, 8.0]])
+        self.assertEqual(pooled_image.shape, expected.shape)
+        self.assertTrue(torch.allclose(pooled_image, expected))
+        self.assertTrue(isinstance(pooled_image, torch.Tensor))
 
 
 class TestMath(unittest.TestCase):

From d7d0629a4be83e0ad1627e4fac9efb1389bb8bfd Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Mon, 1 Sep 2025 16:08:13 +0200
Subject: [PATCH 08/18] type hints

---
 deeptrack/math.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 05783551d..5c3dafce5 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1268,7 +1268,7 @@ class MaxPooling(Pool):
     ----------
     ksize: int
         Size of the pooling kernel.
-    **kwargs: dict
+    **kwargs: Any
         Additional parameters sent to the pooling function.
 
     Examples
@@ -1315,11 +1315,11 @@ def __init__(
         super().__init__(np.max, ksize=ksize, **kwargs)
 
     def _get_numpy(
-        self,
-        image: NDArray,
+        self: MaxPooling,
+        image: NDArray[Any],
         ksize: int=3,
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> NDArray[Any]:
         """Method to perform average pooling with the numpy backend enabled.
 
         Returns the result of the image passed to the scikit image block_reduce
@@ -1347,11 +1347,11 @@ def _get_numpy(
         )
 
     def _get_torch(
-        self,
+        self: MaxPooling,
         image: torch.Tensor,
         ksize: int=3,
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> torch.Tensor:
         """Method to perform max pooling with the torch backend enabled.
 
         Returns the result of the image passed to a torch max pooling layer.
@@ -1385,11 +1385,11 @@ def _get_torch(
         )
 
     def get(
-        self,
-        image: NDArray | torch.Tensor,
+        self: MaxPooling,
+        image: NDArray[Any] | torch.Tensor,
         ksize: int=3,
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> NDArray[Any] | torch.Tensor:
         """Method to perform pooling with either torch or numpy backend.
 
         Checks the current backend and chooses the appropriate function to pool

From ef45e072e0207cf8be95e01d7203846faf170997 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Wed, 3 Sep 2025 13:05:11 +0200
Subject: [PATCH 09/18] implemented xp for tests

---
 deeptrack/tests/test_math.py | 29 ++++++-----------------------
 1 file changed, 6 insertions(+), 23 deletions(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index d35936330..6e96746a2 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -83,10 +83,13 @@ def test_Blur(self):
         #self.assertTrue(xp.all(blurred_image == expected_output))
 
     def test_MaxPooling(self):
-        input_image = np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=float)
+        input_image = xp.asarray([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=float)
         feature = math.MaxPooling(ksize=2)
         pooled_image = feature.resolve(input_image)
-        self.assertTrue(np.all(pooled_image == [[6.0, 8.0]]))
+
+        expected = xp.asarray([[6.0, 8.0]], dtype=float)
+
+        self.assertTrue(xp.all(pooled_image == expected))
         self.assertEqual(pooled_image.shape, (1, 2))
 
 
@@ -94,27 +97,7 @@ def test_MaxPooling(self):
 @unittest.skipUnless(TORCH_AVAILABLE, "PyTorch is not installed.")
 class TestMath_Torch(TestMath_Numpy):
     BACKEND = "torch"
-
-    def test_MaxPooling(self):
-        # (1, 1, 2, 4)
-        input_image = torch.tensor([[[ [1.0, 2.0, 3.0, 4.0],
-                                        [5.0, 6.0, 7.0, 8.0] ]]])
-        feature = math.MaxPooling(ksize=2)
-        pooled_image = feature(input_image, ksize=2)
-        expected = torch.tensor([[[[6.0, 8.0]]]])
-        self.assertEqual(pooled_image.shape, expected.shape)
-        self.assertTrue(torch.allclose(pooled_image, expected))
-        self.assertTrue(isinstance(pooled_image, torch.Tensor))
-        
-        # (2, 4)
-        input_image = torch.tensor([ [1.0, 2.0, 3.0, 4.0],
-                                        [5.0, 6.0, 7.0, 8.0] ])
-        feature = math.MaxPooling(ksize=2)
-        pooled_image = feature(input_image, ksize=2)
-        expected = torch.tensor([[6.0, 8.0]])
-        self.assertEqual(pooled_image.shape, expected.shape)
-        self.assertTrue(torch.allclose(pooled_image, expected))
-        self.assertTrue(isinstance(pooled_image, torch.Tensor))
+    pass
 
 
 class TestMath(unittest.TestCase):

From b512d840eaa357bc413a31895a026bbb1872e03d Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Fri, 5 Sep 2025 14:30:02 +0200
Subject: [PATCH 10/18] Update math.py

---
 deeptrack/math.py | 116 +++++++++++++++++++++++-----------------------
 1 file changed, 57 insertions(+), 59 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 5c3dafce5..cf99f548c 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1252,16 +1252,16 @@ def __init__(
 class MaxPooling(Pool):
     """Apply max pooling to images.
 
-    This class inherits from `Pool` to reduce the resolution of an image by
-    dividing it into non-overlapping blocks of size `ksize` and applying the
-    `max` function to each block. The result is a downsampled image where each
-    pixel value represents the maximum value within the corresponding block of
-    the original image. This is useful for reducing the size of an image while
+    `MaxPooling` reduces the resolution of an image by dividing it into
+    non-overlapping blocks of size `ksize` and applying the `max` function
+    to each block. The result is a downsampled image where each pixel value
+    represents the maximum value within the corresponding block of the
+    original image. This is useful for reducing the size of an image while
     retaining the most significant features.
 
-    If the backend is numpy, the downsampling is performed using
+    If the backend is NumPy, the downsampling is performed using
     `skimage.measure.block_reduce`.
-    If the backend is torch, the downsampling
+    If the backend is PyTorch, the downsampling
     is performed using `torch.nn.functional.max_pool2d`.
 
     Parameters
@@ -1274,24 +1274,17 @@ class MaxPooling(Pool):
     Examples
     --------
     >>> import deeptrack as dt
-    >>> import numpy as np
     Create an input image:
+    >>> import numpy as np
+    >>>
     >>> input_image = np.random.rand(32, 32)
 
-    Define a max pooling feature:
+    Define and use a max pooling feature:
     >>> max_pooling = dt.MaxPooling(ksize=8)
     >>> output_image = max_pooling(input_image)
-    >>> print(output_image.shape)
+    >>> output_image.shape
     (4, 4)
 
-    Notes
-    -----
-    Calling this feature returns a pooled image of the input, it will return
-    either numpy or torch depending on the backend. If `store_properties` is
-    set to `True` and the input is a numpy array, the returned array will be
-    automatically wrapped in an `Image` object. This behavior is handled
-    internally and does not affect the return type of the `get()` method.
-
     """
 
     def __init__(
@@ -1302,7 +1295,7 @@ def __init__(
         """Initialize the parameters for max pooling.
 
         This constructor initializes the parameters for max pooling and checks
-        whether to use the numpy or torch implementation, defaults to numpy.
+        whether to use the NumPy or PyTorch implementation, defaults to NumPy.
 
         Parameters
         ----------
@@ -1314,16 +1307,50 @@ def __init__(
         """
         super().__init__(np.max, ksize=ksize, **kwargs)
 
+    
+    def get(
+        self: MaxPooling,
+        image: NDArray[Any] | torch.Tensor,
+        ksize: int=3,
+        **kwargs: Any,
+    ) -> NDArray[Any] | torch.Tensor:
+        """Max pooling of input.
+
+        Checks the current backend and chooses the appropriate function to pool
+        the input image, either `._get_torch()` or `._get_numpy()`.
+
+        Parameters
+        ----------
+        image: array or tensor
+            Input array or tensor be pooled.
+        ksize: int
+            Kernel size of the pooling operation.
+
+        Returns
+        -------
+        array or tensor
+            The pooled image as `NDArray` or `torch.Tensor` depending on
+            the backend.
+
+        """
+        if self.get_backend() == "numpy":
+            return self._get_numpy(image, ksize, **kwargs)
+        elif self.get_backend() == "torch":
+            return self._get_torch(image, ksize, **kwargs)
+        else:
+            raise NotImplementedError(f"Backend {self.backend} not supported")
+
+
     def _get_numpy(
         self: MaxPooling,
         image: NDArray[Any],
         ksize: int=3,
         **kwargs: Any,
     ) -> NDArray[Any]:
-        """Method to perform average pooling with the numpy backend enabled.
+        """Max pooling pooling with the NumPy backend enabled.
 
-        Returns the result of the image passed to the scikit image block_reduce
-        function with `np.max()` as the pooling function.
+        Returns the result of the input array passed to the scikit
+        image `block_reduce()` function with `np.max()` as the pooling function.
 
         Parameters
         ----------
@@ -1341,7 +1368,7 @@ def _get_numpy(
         return utils.safe_call(
             skimage.measure.block_reduce,
             image=image,
-            func=self.pooling, # This will be np.max for this class.
+            func=np.min
             block_size=ksize,
             **kwargs,
         )
@@ -1352,14 +1379,15 @@ def _get_torch(
         ksize: int=3,
         **kwargs: Any,
     ) -> torch.Tensor:
-        """Method to perform max pooling with the torch backend enabled.
+        """Perform max pooling with the PyTorch backend enabled.
+
 
-        Returns the result of the image passed to a torch max pooling layer.
+        Returns the result of the tensor passed to a PyTorch max pooling layer.
 
         Parameters
         ----------
         image: torch.Tensor
-            Input image to be pooled.
+            Input tensor to be pooled.
         ksize: int
             Kernel size of the pooling operation.
 
@@ -1369,8 +1397,10 @@ def _get_torch(
             The pooled image as a `torch.Tensor`.
 
         """
-        # If needed, expand tensor shape
+
+        # If input tensor is 2D
         if len(image.shape) == 2:
+            # Add batch dimension for max pooling.
             expanded_image = image.unsqueeze(0)
 
             pooled_image = torch.nn.functional.max_pool2d(
@@ -1384,38 +1414,6 @@ def _get_torch(
             kernel_size=ksize,
         )
 
-    def get(
-        self: MaxPooling,
-        image: NDArray[Any] | torch.Tensor,
-        ksize: int=3,
-        **kwargs: Any,
-    ) -> NDArray[Any] | torch.Tensor:
-        """Method to perform pooling with either torch or numpy backend.
-
-        Checks the current backend and chooses the appropriate function to pool
-        the input image, either `_get_torch` or `_get_numpy`.
-
-        Parameters
-        ----------
-        image: NDArray | torch.Tensor
-            Input image to be pooled.
-        ksize: int
-            Kernel size of the pooling operation.
-
-        Returns
-        -------
-        NDArray | torch.Tensor
-            The pooled image as `NDArray` or `torch.Tensor` depending on
-            the backend.
-
-        """
-        if self.get_backend() == "numpy":
-            return self._get_numpy(image, ksize, **kwargs,)
-        elif self.get_backend() == "torch":
-            return self._get_torch(image, ksize, **kwargs,)
-        else:
-            raise NotImplementedError(f"Backend {self.backend} not supported")
-
 
 #TODO ***AL*** revise MinPooling - torch, typing, docstring, unit test
 class MinPooling(Pool):

From e450da45b5b863b2d6557c35a2f14f91adbe004c Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Fri, 5 Sep 2025 14:34:36 +0200
Subject: [PATCH 11/18] Update test_math.py


From d5233b1d4d5d85f29c38e0844e710d156219e23b Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Fri, 5 Sep 2025 14:40:04 +0200
Subject: [PATCH 12/18] Update math.py

---
 deeptrack/math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index cf99f548c..6bfded59f 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1368,7 +1368,7 @@ def _get_numpy(
         return utils.safe_call(
             skimage.measure.block_reduce,
             image=image,
-            func=np.min
+            func=np.min,
             block_size=ksize,
             **kwargs,
         )

From 5ccaded0cea6d6f83c9ba7b039d9d363b320bd01 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Fri, 5 Sep 2025 14:47:20 +0200
Subject: [PATCH 13/18] Update test_math.py

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index 6e96746a2..e4b2dbc5d 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -123,7 +123,7 @@ def test_MaxPooling(self):
         input_image = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
         feature = math.MaxPooling(ksize=2)
         pooled_image = feature.resolve(input_image)
-        self.assertTrue(np.all(pooled_image == [[5, 6], [8, 9]]))
+        self.assertTrue(np.all(pooled_image == np.array([[5, 6], [8, 9]])))
 
     def test_MinPooling(self):
         input_image = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])

From 0ab8a2748794f65320abfca689e268d69a74f748 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Fri, 5 Sep 2025 14:57:20 +0200
Subject: [PATCH 14/18] Update test_math.py

---
 deeptrack/tests/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/tests/test_math.py b/deeptrack/tests/test_math.py
index ab9b406c0..09799e292 100644
--- a/deeptrack/tests/test_math.py
+++ b/deeptrack/tests/test_math.py
@@ -134,7 +134,7 @@ def test_MaxPooling(self):
         input_image = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
         feature = math.MaxPooling(ksize=2)
         pooled_image = feature.resolve(input_image)
-        self.assertTrue(np.all(pooled_image == np.array([[5, 6], [8, 9]])))
+        self.assertTrue(xp.all(pooled_image == xp.asarray([[5, 6], [8, 9]]) ) )
 
     def test_MinPooling(self):
         input_image = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])

From f4141ddc74a24bef9e4dc7b5facfd52bd4353f6e Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Fri, 5 Sep 2025 15:14:00 +0200
Subject: [PATCH 15/18] Update math.py

---
 deeptrack/math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 87b5700b7..7fcd51fcb 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1369,7 +1369,7 @@ def _get_numpy(
         return utils.safe_call(
             skimage.measure.block_reduce,
             image=image,
-            func=np.min,
+            func=np.max,
             block_size=ksize,
             **kwargs,
         )

From 5360352ca4d2ef6e3907ed67148e3a3b7cdfc446 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Fri, 5 Sep 2025 15:30:59 +0200
Subject: [PATCH 16/18] Update math.py

---
 deeptrack/math.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 7fcd51fcb..46ce23f65 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1295,7 +1295,6 @@ def __init__(
     ):
         """Initialize the parameters for max-pooling.
 
-
         This constructor initializes the parameters for max-pooling.
 
         Parameters
@@ -1306,6 +1305,7 @@ def __init__(
             Additional keyword arguments.
 
         """
+        
         super().__init__(np.max, ksize=ksize, **kwargs)
 
     
@@ -1330,7 +1330,7 @@ def get(
         Returns
         -------
         array or tensor
-            The pooled image as `NDArray` or `torch.Tensor` depending on
+            The pooled input as `NDArray` or `torch.Tensor` depending on
             the backend.
 
         """
@@ -1356,7 +1356,7 @@ def _get_numpy(
         Parameters
         ----------
         image: NDArray
-            Input image to be pooled.
+            Input array to be pooled.
         ksize: int
             Kernel size of the pooling operation.
 
@@ -1383,7 +1383,8 @@ def _get_torch(
         """Perform max pooling with the PyTorch backend enabled.
 
 
-        Returns the result of the tensor passed to a PyTorch max pooling layer.
+        Returns the result of the tensor passed to a PyTorch max
+        pooling layer.
 
         Parameters
         ----------

From 00bfa84ec09c082c9b3b37a88cc50673b5f0dae4 Mon Sep 17 00:00:00 2001
From: Alex <95913221+Pwhsky@users.noreply.github.com>
Date: Fri, 5 Sep 2025 15:37:15 +0200
Subject: [PATCH 17/18] u

---
 deeptrack/math.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index 46ce23f65..e2db42f05 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1305,10 +1305,9 @@ def __init__(
             Additional keyword arguments.
 
         """
-        
+
         super().__init__(np.max, ksize=ksize, **kwargs)
 
-    
     def get(
         self: MaxPooling,
         image: NDArray[Any] | torch.Tensor,
@@ -1341,7 +1340,6 @@ def get(
         else:
             raise NotImplementedError(f"Backend {self.backend} not supported")
 
-
     def _get_numpy(
         self: MaxPooling,
         image: NDArray[Any],

From 8ff8c73e63d49d521085fa7273e1756ff4ba9123 Mon Sep 17 00:00:00 2001
From: Giovanni Volpe <giovanni.volpe@physics.gu.se>
Date: Fri, 5 Sep 2025 16:26:53 +0200
Subject: [PATCH 18/18] Update math.py

---
 deeptrack/math.py | 37 +++++++++++++++++++++----------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/deeptrack/math.py b/deeptrack/math.py
index e2db42f05..eb9be8040 100644
--- a/deeptrack/math.py
+++ b/deeptrack/math.py
@@ -1261,8 +1261,9 @@ class MaxPooling(Pool):
 
     If the backend is NumPy, the downsampling is performed using
     `skimage.measure.block_reduce`.
-    If the backend is PyTorch, the downsampling
-    is performed using `torch.nn.functional.max_pool2d`.
+
+    If the backend is PyTorch, the downsampling is performed using
+    `torch.nn.functional.max_pool2d`.
 
     Parameters
     ----------
@@ -1274,12 +1275,13 @@ class MaxPooling(Pool):
     Examples
     --------
     >>> import deeptrack as dt
+
     Create an input image:
     >>> import numpy as np
     >>>
     >>> input_image = np.random.rand(32, 32)
 
-    Define and use a max pooling feature:
+    Define and use a max-pooling feature:
 
     >>> max_pooling = dt.MaxPooling(ksize=8)
     >>> output_image = max_pooling(input_image)
@@ -1314,7 +1316,7 @@ def get(
         ksize: int=3,
         **kwargs: Any,
     ) -> NDArray[Any] | torch.Tensor:
-        """Max pooling of input.
+        """Max-pooling of input.
 
         Checks the current backend and chooses the appropriate function to pool
         the input image, either `._get_torch()` or `._get_numpy()`.
@@ -1333,12 +1335,14 @@ def get(
             the backend.
 
         """
+
         if self.get_backend() == "numpy":
             return self._get_numpy(image, ksize, **kwargs)
-        elif self.get_backend() == "torch":
+
+        if self.get_backend() == "torch":
             return self._get_torch(image, ksize, **kwargs)
-        else:
-            raise NotImplementedError(f"Backend {self.backend} not supported")
+
+        raise NotImplementedError(f"Backend {self.backend} not supported")
 
     def _get_numpy(
         self: MaxPooling,
@@ -1346,24 +1350,25 @@ def _get_numpy(
         ksize: int=3,
         **kwargs: Any,
     ) -> NDArray[Any]:
-        """Max pooling pooling with the NumPy backend enabled.
+        """Max-pooling pooling with the NumPy backend enabled.
 
-        Returns the result of the input array passed to the scikit
-        image `block_reduce()` function with `np.max()` as the pooling function.
+        Returns the result of the input array passed to the scikit image
+        `block_reduce()` function with `np.max()` as the pooling function.
 
         Parameters
         ----------
-        image: NDArray
+        image: array
             Input array to be pooled.
         ksize: int
             Kernel size of the pooling operation.
 
         Returns
         -------
-        NDArray
-            The pooled image as a `NDArray`.
+        array
+            The pooled image as a NumPy array.
             
         """
+
         return utils.safe_call(
             skimage.measure.block_reduce,
             image=image,
@@ -1378,7 +1383,7 @@ def _get_torch(
         ksize: int=3,
         **kwargs: Any,
     ) -> torch.Tensor:
-        """Perform max pooling with the PyTorch backend enabled.
+        """Max-pooling with the PyTorch backend enabled.
 
 
         Returns the result of the tensor passed to a PyTorch max
@@ -1400,13 +1405,13 @@ def _get_torch(
 
         # If input tensor is 2D
         if len(image.shape) == 2:
-            # Add batch dimension for max pooling.
+            # Add batch dimension for max-pooling
             expanded_image = image.unsqueeze(0)
 
             pooled_image = torch.nn.functional.max_pool2d(
                 expanded_image, kernel_size=ksize,
             )
-            # Remove the expanded dim.
+            # Remove the expanded dim
             return pooled_image.squeeze(0)
 
         return torch.nn.functional.max_pool2d(