pass kwargs down to the reducer

Author: d-v-b

src/xarray_multiscale/multiscale.py

@@ -36,6 +36,7 @@ def multiscale(
     chunks: ChunkOption | Sequence[int] | dict[Hashable, int] = "preserve",
     chained: bool = True,
     namer: Callable[[int], str] = _default_namer,
+    **kwargs: Any,
 ) -> list[xarray.DataArray]:
     """
     Generate a coordinate-aware multiscale representation of an array.
@@ -90,6 +91,9 @@ def multiscale(
         index and return a string. The default function simply prepends the string
         representation of the integer with the character "s".
 
+    **kwargs: Any
+        Additional keyword arguments that will be passed to the reduction function.
+
     Returns
     -------
     result : list[xarray.DataArray]
@@ -127,7 +131,7 @@ def multiscale(
         else:
             scale = tuple(s**level for s in scale_factors)
             source = result[0]
-        downscaled = downscale(source, reduction, scale, preserve_dtype)
+        downscaled = downscale(source, reduction, scale, preserve_dtype, **kwargs)
         downscaled.name = namer(level)
         result.append(downscaled)
 
@@ -201,7 +205,7 @@ def downscale(
     if to_downscale.chunks is not None:
         downscaled_data = downscale_dask(to_downscale.data, reduction, scale_factors, **kwargs)
     else:
-        downscaled_data = reduction(to_downscale.data, scale_factors)
+        downscaled_data = reduction(to_downscale.data, scale_factors, **kwargs)
     if preserve_dtype:
         downscaled_data = downscaled_data.astype(array.dtype)
     downscaled_coords = downscale_coords(to_downscale, scale_factors)

tests/test_multiscale.py

@@ -1,6 +1,7 @@
 import dask.array as da
 import numpy as np
 import pytest
+from src.xarray_multiscale.reducers import windowed_rank
 from xarray import DataArray
 from xarray.testing import assert_equal
 
@@ -127,6 +128,14 @@ def test_multiscale(ndim: int, chained: bool):
     assert np.array_equal(pyr[0].data, base_array)
 
 
+@pytest.mark.parametrize("rank", (-1, 0, 1))
+def test_multiscale_rank_kwargs(rank: int):
+    data = np.arange(16)
+    window_size = (4,)
+    pyr = multiscale(data, windowed_rank, window_size, rank=rank)
+    assert np.array_equal(pyr[1].data, windowed_rank(data, window_size=window_size, rank=rank))
+
+
 def test_chunking():
     ndim = 3
     shape = (16,) * ndim

tests/test_reducers.py

@@ -69,7 +69,7 @@ def test_windowed_rank():
     window_size = (2, 2, 2)
 
     # 2nd brightest voxel
-    rank = np.product(window_size) - 2
+    rank = np.prod(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)