Added "nearest", "wrap", and "truncate" modes to `ndfilters.generic_f…

…ilter()`.

ndfilters/_generic.py

@@ -14,7 +14,7 @@ def generic_filter(
     size: int | tuple[int, ...],
     axis: None | int | tuple[int, ...] = None,
     where: bool | np.ndarray = True,
-    mode: Literal["mirror"] = "mirror",
+    mode: Literal["mirror", "nearest", "wrap", "truncate"] = "mirror",
     args: tuple = (),
 ) -> np.ndarray:
     """
@@ -42,7 +42,8 @@ def generic_filter(
     mode
         The method used to extend the input array beyond its boundaries.
         See :func:`scipy.ndimage.generic_filter` for the definitions.
-        Currently, only "reflect" mode is supported.
+        Currently, only "mirror", "nearest", "wrap", and "truncate" modes are
+        supported.
     args
         Extra arguments to pass to function.
 
@@ -98,9 +99,6 @@ def function(a: np.ndarray, args: tuple) -> float:
                 f"{size=} should have the same number of elements as {axis=}."
             )
 
-    if mode != "mirror":  # pragma: nocover
-        raise ValueError(f"Only mode='reflected' is supported, got {mode=}")
-
     axis_numba = ~np.arange(len(axis))[::-1]
 
     shape = array.shape
@@ -138,6 +136,30 @@ def function(a: np.ndarray, args: tuple) -> float:
     return result
 
 
+@numba.njit
+def _rectify_index_lower(index: int, size: int, mode: str) -> int:
+    if mode == "mirror":
+        return -index
+    elif mode == "nearest":
+        return 0
+    elif mode == "wrap":
+        return index % size
+    else:
+        raise ValueError
+
+
+@numba.njit
+def _rectify_index_upper(index: int, size: int, mode: str) -> int:
+    if mode == "mirror":
+        return ~(index % size + 1)
+    elif mode == "nearest":
+        return size - 1
+    elif mode == "wrap":
+        return index % size
+    else:
+        raise ValueError
+
+
 @numba.njit(parallel=True)
 def _generic_filter_1d(
     array: np.ndarray,
@@ -157,18 +179,22 @@ def _generic_filter_1d(
 
         for ix in numba.prange(array_shape_x):
 
-            values = np.empty(shape=size)
-            mask = np.empty(shape=size, dtype=np.bool_)
+            values = np.zeros(shape=size)
+            mask = np.zeros(shape=size, dtype=np.bool_)
 
             for kx in range(kernel_shape_x):
 
                 px = kx - kernel_shape_x // 2
                 jx = ix + px
 
                 if jx < 0:
-                    jx = -jx
+                    if mode == "truncate":
+                        continue
+                    jx = _rectify_index_lower(jx, array_shape_x, mode)
                 elif jx >= array_shape_x:
-                    jx = ~(jx % array_shape_x + 1)
+                    if mode == "truncate":
+                        continue
+                    jx = _rectify_index_upper(jx, array_shape_x, mode)
 
                 values[kx] = array[it, jx]
                 mask[kx] = where[it, jx]
@@ -198,28 +224,36 @@ def _generic_filter_2d(
         for ix in numba.prange(array_shape_x):
             for iy in numba.prange(array_shape_y):
 
-                values = np.empty(shape=size)
-                mask = np.empty(shape=size, dtype=np.bool_)
+                values = np.zeros(shape=size)
+                mask = np.zeros(shape=size, dtype=np.bool_)
 
                 for kx in range(kernel_shape_x):
 
                     px = kx - kernel_shape_x // 2
                     jx = ix + px
 
                     if jx < 0:
-                        jx = -jx
+                        if mode == "truncate":
+                            continue
+                        jx = _rectify_index_lower(jx, array_shape_x, mode)
                     elif jx >= array_shape_x:
-                        jx = ~(jx % array_shape_x + 1)
+                        if mode == "truncate":
+                            continue
+                        jx = _rectify_index_upper(jx, array_shape_x, mode)
 
                     for ky in range(kernel_shape_y):
 
                         py = ky - kernel_shape_y // 2
                         jy = iy + py
 
                         if jy < 0:
-                            jy = -jy
+                            if mode == "truncate":
+                                continue
+                            jy = _rectify_index_lower(jy, array_shape_y, mode)
                         elif jy >= array_shape_y:
-                            jy = ~(jy % array_shape_y + 1)
+                            if mode == "truncate":
+                                continue
+                            jy = _rectify_index_upper(jy, array_shape_y, mode)
 
                         values[kx, ky] = array[it, jx, jy]
                         mask[kx, ky] = where[it, jx, jy]
@@ -253,38 +287,50 @@ def _generic_filter_3d(
             for iy in numba.prange(array_shape_y):
                 for iz in numba.prange(array_shape_z):
 
-                    values = np.empty(shape=size)
-                    mask = np.empty(shape=size, dtype=np.bool_)
+                    values = np.zeros(shape=size)
+                    mask = np.zeros(shape=size, dtype=np.bool_)
 
                     for kx in range(kernel_shape_x):
 
                         px = kx - kernel_shape_x // 2
                         jx = ix + px
 
                         if jx < 0:
-                            jx = -jx
+                            if mode == "truncate":
+                                continue
+                            jx = _rectify_index_lower(jx, array_shape_x, mode)
                         elif jx >= array_shape_x:
-                            jx = ~(jx % array_shape_x + 1)
+                            if mode == "truncate":
+                                continue
+                            jx = _rectify_index_upper(jx, array_shape_x, mode)
 
                         for ky in range(kernel_shape_y):
 
                             py = ky - kernel_shape_y // 2
                             jy = iy + py
 
                             if jy < 0:
-                                jy = -jy
+                                if mode == "truncate":
+                                    continue
+                                jy = _rectify_index_lower(jy, array_shape_y, mode)
                             elif jy >= array_shape_y:
-                                jy = ~(jy % array_shape_y + 1)
+                                if mode == "truncate":
+                                    continue
+                                jy = _rectify_index_upper(jy, array_shape_y, mode)
 
                             for kz in range(kernel_shape_z):
 
                                 pz = kz - kernel_shape_z // 2
                                 jz = iz + pz
 
                                 if jz < 0:
-                                    jz = -jz
+                                    if mode == "truncate":
+                                        continue
+                                    jz = _rectify_index_lower(jz, array_shape_z, mode)
                                 elif jz >= array_shape_z:
-                                    jz = ~(jz % array_shape_z + 1)
+                                    if mode == "truncate":
+                                        continue
+                                    jz = _rectify_index_upper(jz, array_shape_z, mode)
 
                                 values[kx, ky, kz] = array[it, jx, jy, jz]
                                 mask[kx, ky, kz] = where[it, jx, jy, jz]

ndfilters/_trimmed_mean.py

@@ -14,7 +14,7 @@ def trimmed_mean_filter(
     size: int | tuple[int, ...],
     axis: None | int | tuple[int, ...] = None,
     where: bool | np.ndarray = True,
-    mode: Literal["mirror"] = "mirror",
+    mode: Literal["mirror", "nearest", "wrap", "truncate"] = "mirror",
     proportion: float = 0.25,
 ) -> np.ndarray:
     """
@@ -36,7 +36,8 @@ def trimmed_mean_filter(
     mode
         The method used to extend the input array beyond its boundaries.
         See :func:`scipy.ndimage.generic_filter` for the definitions.
-        Currently, only "reflect" mode is supported.
+        Currently, only "mirror", "nearest", "wrap", and "truncate" modes are
+        supported.
     proportion
         The proportion to cut from the top and bottom of the distribution.
 
@@ -83,6 +84,8 @@ def _trimmed_mean(
     (proportion,) = args
 
     nobs = array.size
+    if nobs == 0:
+        return np.nan
     lowercut = int(proportion * nobs)
     uppercut = nobs - lowercut
     if lowercut > uppercut:  # pragma: nocover