add recursive mode

Author: d-v-b

src/xarray_multiscale/multiscale.py

@@ -64,21 +64,26 @@ def multiscale(
         tuple(s ** l for s in scale_factors) for l in levels
     )
     result = [_ingest_array(array, scales[0])]
-    base_attrs = result[0].attrs
-    base_coords = result[0].coords
 
-    for scale in scales[1:]:
-        downscaled = downscale(
-            result[0], reduction, scale, pad_mode=pad_mode, preserve_dtype=preserve_dtype
-        )
+    for level in levels[1:]:
+        if recursive:
+            scale = scale_factors
+            downscaled = downscale(result[-1], reduction, scale, pad_mode=pad_mode) 
+        else:
+            scale = scales[level]
+            downscaled = downscale(result[0], reduction, scale, pad_mode=pad_mode)
         result.append(downscaled)
 
+    if preserve_dtype:
+        result = [r.astype(array.dtype) for r in result]
+
     if chunks is not None:
         if isinstance(chunks, Sequence):
             _chunks = {k: v for k, v in zip(result[0].dims, chunks)}
         else:
             _chunks = chunks
         result = [r.chunk(_chunks) for r in result]
+
     return result
 
 
@@ -202,7 +207,6 @@ def downscale(
     reduction: Callable,
     scale_factors: Sequence[int],
     pad_mode: Optional[str] = None,
-    preserve_dtype: bool = True,
     **kwargs,
 ) -> DataArray:
     """
@@ -239,8 +243,6 @@ def downscale(
         **kwargs,
     )
 
-    if preserve_dtype:
-        coarsened = coarsened.astype(array.dtype)
 
     if isinstance(array, xarray.DataArray):
         base_coords = array.coords

tests/test_multiscale.py

@@ -54,33 +54,20 @@ def test_downscale_2d():
     arr_xarray = DataArray(arr_dask)
 
     downscaled_numpy_float = downscale(
-        arr_numpy, np.mean, scale, preserve_dtype=False
-    ).compute()
+        arr_numpy, np.mean, scale).compute()
 
     downscaled_dask_float = downscale(
-        arr_dask, np.mean, scale, preserve_dtype=False
-    ).compute()
+        arr_dask, np.mean, scale).compute()
 
     downscaled_xarray_float = downscale(
-        arr_xarray, np.mean, scale, preserve_dtype=False
-    ).compute()
+        arr_xarray, np.mean, scale).compute()
 
     answer_float = np.array([[0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.5, 0.5]])
+    
     assert np.array_equal(downscaled_numpy_float, answer_float)
     assert np.array_equal(downscaled_dask_float, answer_float)
     assert np.array_equal(downscaled_xarray_float, answer_float)
 
-    downscaled_numpy_int = downscale(
-        arr_numpy, np.mean, scale, dtype=arr_numpy.dtype
-    ).compute()
-    downscaled_dask_int = downscale(
-        arr_dask, np.mean, scale, dtype=arr_numpy.dtype
-    ).compute()
-
-    answer_int = answer_float.astype("int")
-    assert np.array_equal(downscaled_numpy_int, answer_int)
-    assert np.array_equal(downscaled_dask_int, answer_int)
-
 
 def test_multiscale():
     ndim = 3
@@ -94,7 +81,7 @@ def test_multiscale():
 
     pyr_trimmed = multiscale(array, np.mean, 2, pad_mode=None)
     pyr_padded = multiscale(array, np.mean, 2, pad_mode="reflect")
-
+    pyr_trimmed_recursive = multiscale(array, np.mean, 2, pad_mode=None, recursive=True)
     assert [p.shape for p in pyr_padded] == [
         shape,
         (5, 5, 5),
@@ -111,4 +98,7 @@ def test_multiscale():
     assert np.array_equal(
         pyr_trimmed[-2].data.mean().compute(), pyr_trimmed[-1].data.compute().mean()
     )
+    assert np.array_equal(
+        pyr_trimmed_recursive[-2].data.mean().compute(), pyr_trimmed_recursive[-1].data.compute().mean()
+    )
     assert np.allclose(pyr_padded[0].data.mean().compute(), 0.17146776406035666)