Merge branch 'main' of https://github.com/JaneliaSciComp/xarray-multiscale into use-hatch

Author: d-v-b

src/xarray_multiscale/__init__.py

@@ -5,7 +5,7 @@
     windowed_max,
     windowed_mean,
     windowed_min,
-    windowed_mode,
+    windowed_rank,
 )
 
 __all__ = [
@@ -15,4 +15,5 @@
     "windowed_mode",
     "windowed_max",
     "windowed_min",
+    "windowed_rank",
 ]

src/xarray_multiscale/reducers.py

@@ -378,3 +378,60 @@ def windowed_mode_countless(
     final_result: npt.NDArray[Any] = _lor(reduce(_lor, results), sections[-1]) - 1  # type: ignore[arg-type]
 
     return final_result
+
+
+def windowed_rank(
+    array: npt.NDArray[Any], window_size: tuple[int, ...], rank: int = -1
+) -> npt.NDArray[Any]:
+    """
+    Compute the windowed rank order filter of an array.
+    Input will be coerced to a numpy array.
+
+    Parameters
+    ----------
+    array: Array-like, e.g. Numpy array, Dask array
+        The array to be downscaled. The array must have a ``reshape``
+        method.
+
+    window_size: tuple[int, ...]
+        The window to use for aggregation. The array is partitioned into
+        non-overlapping regions with size equal to ``window_size``, and the
+        values in each window are sorted to generate the result.
+
+    rank: int, default=-1
+        The index to take from the sorted values in each window. If non-negative, then
+        rank must be between 0 and the product of the elements of ``window_size`` minus one,
+        (inclusive).
+        Rank may be negative, in which case it denotes an index relative to the end of the sorted
+        values following normal python indexing rules.
+        E.g., when rank is -1 (the default), this takes the maxmum value of each window.
+
+    Returns
+    -------
+    Numpy array
+        The result of the windowed rank filter. The length of each axis of this array
+        will be a fraction of the input.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from xarray_multiscale.reducers import windowed_rank
+    >>> data = np.arange(16).reshape(4, 4)
+    >>> windowed_rank(data, (2, 2), -2)
+    array([[ 4  6]
+           [12 14]])
+    """
+    max_rank = np.prod(window_size) - 1
+    if rank > max_rank or rank < -max_rank - 1:
+        msg = (
+            f"Invalid rank: {rank} for window_size: {window_size} ",
+            f"If rank is negative then between either -1 and {-max_rank-1}, inclusive",
+            f"If rank is non-negtaive, then it must be between 0 and {max_rank}, inclusive.",
+        )
+        raise ValueError(msg)
+    reshaped = reshape_windowed(array, window_size)
+    transposed_shape = tuple(range(0, reshaped.ndim, 2)) + tuple(range(1, reshaped.ndim, 2))
+    transposed = reshaped.transpose(transposed_shape)
+    collapsed = transposed.reshape(tuple(reshaped.shape[slice(0, None, 2)]) + (-1,))
+    result: npt.NDArray[Any] = np.take(np.sort(collapsed, axis=-1), rank, axis=-1)
+    return result

tests/test_reducers.py

@@ -10,6 +10,7 @@
     windowed_mean,
     windowed_min,
     windowed_mode,
+    windowed_rank,
 )
 
 
@@ -62,6 +63,29 @@ def test_windowed_mode():
     assert np.array_equal(results, answer)
 
 
+def test_windowed_rank():
+    initial_array = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+    larger_array = np.tile(initial_array, (2, 2, 2))
+    window_size = (2, 2, 2)
+
+    # 2nd brightest voxel
+    rank = np.product(window_size) - 2
+    answer = np.array([[[7, 7], [7, 7]], [[7, 7], [7, 7]]])
+    results = windowed_rank(larger_array, window_size, rank)
+    assert np.array_equal(results, answer)
+
+    # Test negative rank
+    rank = -8
+    answer = np.array([[[1, 1], [1, 1]], [[1, 1], [1, 1]]])
+    results = windowed_rank(larger_array, window_size, rank)
+    assert np.array_equal(results, answer)
+
+    # Test out-of-bounds rank
+    rank = 100
+    with pytest.raises(ValueError):
+        windowed_rank(larger_array, window_size, rank)
+
+
 @pytest.mark.parametrize("windows_per_dim", (1, 2, 3, 4, 5))
 @pytest.mark.parametrize(
     "window_size", ((1,), (2,), (1, 2), (2, 2), (2, 2, 2), (1, 2, 3), (3, 3, 3, 3))