Enable Rectification and Reprojection for Arrays of Any Dimensionality

CHANGES.md

@@ -12,7 +12,10 @@
   `xcube_resampling.rectify._compute_source_tile_indexing`.
 - Added support for rectifying datasets with decreasing x-coordinate in 
   `xcube_resampling.rectify_dataset`.
-
+- Support data variables with arbitrary dimensionality (≥ 3D) in
+  `xcube_resampling.rectify_dataset` and `xcube_resampling.reproject_dataset`.
+  Spatial dimensions are automatically reordered to the trailing `(y, x)`
+  positions when required.
 
 ## Changes in 0.2.4
 

tests/sampledata.py

@@ -57,6 +57,29 @@ def create_2x2x2_dataset_with_irregular_coords():
     )
 
 
+def create_2x2x2x2_dataset_with_irregular_coords():
+    lon = np.array([[1.0, 6.0], [0.0, 2.0]])
+    lat = np.array([[56.0, 53.0], [52.0, 50.0]])
+    time = pd.date_range("2025-08-01", periods=2)
+    member = np.arange(2)
+    rad = np.array([[1.0, 2.0], [3.0, 4.0]])
+    rad = np.repeat(rad[np.newaxis, :, :], 2, axis=0)
+    rad = np.repeat(rad[np.newaxis, :, :, :], 2, axis=0)
+    rad = np.transpose(rad, (2, 0, 3, 1))
+    return xr.Dataset(
+        dict(
+            rad=xr.DataArray(rad, dims=("y", "member", "x", "time")),
+            time_series=xr.DataArray(np.array([1, 2]), dims="time"),
+        ),
+        coords=dict(
+            lon=xr.DataArray(lon, dims=("y", "x")),
+            lat=xr.DataArray(lat, dims=("y", "x")),
+            time=time,
+            member=member,
+        ),
+    )
+
+
 def create_8x6_dataset_with_regular_coords():
     res = 0.1
     return xr.Dataset(
@@ -142,6 +165,30 @@ def create_2x5x5_dataset_regular_utm():
     return ds
 
 
+def create_2x2x5x5_dataset_regular_utm():
+    x = np.arange(565300.0, 565800.0, 100.0)
+    y = np.arange(5934300.0, 5933800.0, -100.0)
+    time = pd.date_range("2025-08-01", periods=2)
+    members = np.arange(2)
+    spatial_ref = np.array(0)
+    band_1 = np.arange(25).reshape((5, 5))
+    band_1 = np.repeat(band_1[np.newaxis, :, :], 2, axis=0)
+    band_1 = np.repeat(band_1[np.newaxis, :, :, :], 2, axis=0)
+    band_1 = np.transpose(band_1, (2, 0, 3, 1))
+    ds = xr.Dataset(
+        dict(
+            band_1=xr.DataArray(
+                band_1,
+                dims=("y", "members", "x", "time"),
+                attrs=dict(grid_mapping="spatial_ref"),
+            )
+        ),
+        coords=dict(time=time, x=x, y=y, members=members, spatial_ref=spatial_ref),
+    )
+    ds.spatial_ref.attrs = pyproj.CRS.from_epsg("32632").to_cf()
+    return ds
+
+
 def create_large_dataset_for_reproject():
     nt, nx, ny = 10, 100, 100
     chunks = {"time": 2, "x": 25, "y": 25}

tests/test_rectify.py

@@ -31,6 +31,7 @@
     create_2x2_dataset_with_irregular_coords,
     create_2x2_dataset_with_irregular_coords_antimeridian,
     create_2x2x2_dataset_with_irregular_coords,
+    create_2x2x2x2_dataset_with_irregular_coords,
     create_4x4_dataset_with_irregular_coords,
 )
 
@@ -164,6 +165,54 @@ def test_rectify_2x2x2_to_default(self):
             ),
         )
 
+    def test_rectify_2x2x2x2_to_default(self):
+        source_ds = create_2x2x2x2_dataset_with_irregular_coords()
+
+        target_gm = GridMapping.regular(
+            size=(4, 4), xy_min=(0, 50), xy_res=2, crs=CRS_WGS84
+        )
+        target_ds = rectify_dataset(source_ds, target_gm=target_gm, interp_methods=0)
+        self.assertEqual(
+            set(source_ds.variables).union(["spatial_ref"]), set(target_ds.variables)
+        )
+        self.assertEqual(("member", "time", "lat", "lon"), target_ds.rad.dims)
+        np.testing.assert_almost_equal(
+            target_ds.rad.values,
+            np.array(
+                [
+                    [
+                        [
+                            [nan, nan, nan, nan],
+                            [nan, 1.0, 2.0, nan],
+                            [3.0, 3.0, 2.0, nan],
+                            [nan, 4.0, nan, nan],
+                        ],
+                        [
+                            [nan, nan, nan, nan],
+                            [nan, 1.0, 2.0, nan],
+                            [3.0, 3.0, 2.0, nan],
+                            [nan, 4.0, nan, nan],
+                        ],
+                    ],
+                    [
+                        [
+                            [nan, nan, nan, nan],
+                            [nan, 1.0, 2.0, nan],
+                            [3.0, 3.0, 2.0, nan],
+                            [nan, 4.0, nan, nan],
+                        ],
+                        [
+                            [nan, nan, nan, nan],
+                            [nan, 1.0, 2.0, nan],
+                            [3.0, 3.0, 2.0, nan],
+                            [nan, 4.0, nan, nan],
+                        ],
+                    ],
+                ],
+                dtype=target_ds.rad.dtype,
+            ),
+        )
+
     def test_rectify_2x2_to_7x7(self):
         source_ds = create_2x2_dataset_with_irregular_coords()
         # Add offset to "rad" so its values do not lie on a plane

tests/test_reproject.py

@@ -10,6 +10,7 @@
 from xcube_resampling.reproject import reproject_dataset
 
 from .sampledata import (
+    create_2x2x5x5_dataset_regular_utm,
     create_2x5x5_dataset_regular_utm,
     create_5x5_dataset_regular_utm,
     create_large_dataset_for_reproject,
@@ -75,6 +76,58 @@ def test_reproject_target_gm_3d(self):
             ),
         )
 
+    def test_reproject_target_gm_4d(self):
+        source_ds = create_2x2x5x5_dataset_regular_utm()
+
+        # test projected CRS, similar resolution
+        target_gm = GridMapping.regular(
+            size=(5, 5), xy_min=(4320120, 3382520), xy_res=80, crs="epsg:3035"
+        )
+        target_ds = reproject_dataset(source_ds, target_gm)
+        self.assertEqual(set(source_ds.variables), set(target_ds.variables))
+        self.assertEqual(("members", "time", "y", "x"), target_ds.band_1.dims)
+
+        np.testing.assert_almost_equal(
+            target_ds.band_1.values,
+            np.array(
+                [
+                    [
+                        [
+                            [1, 1, 2, 3, 4],
+                            [6, 6, 7, 8, 9],
+                            [11, 12, 12, 13, 14],
+                            [16, 17, 17, 18, 19],
+                            [21, 17, 17, 18, 19],
+                        ],
+                        [
+                            [1, 1, 2, 3, 4],
+                            [6, 6, 7, 8, 9],
+                            [11, 12, 12, 13, 14],
+                            [16, 17, 17, 18, 19],
+                            [21, 17, 17, 18, 19],
+                        ],
+                    ],
+                    [
+                        [
+                            [1, 1, 2, 3, 4],
+                            [6, 6, 7, 8, 9],
+                            [11, 12, 12, 13, 14],
+                            [16, 17, 17, 18, 19],
+                            [21, 17, 17, 18, 19],
+                        ],
+                        [
+                            [1, 1, 2, 3, 4],
+                            [6, 6, 7, 8, 9],
+                            [11, 12, 12, 13, 14],
+                            [16, 17, 17, 18, 19],
+                            [21, 17, 17, 18, 19],
+                        ],
+                    ],
+                ],
+                dtype=target_ds.band_1.dtype,
+            ),
+        )
+
     def test_reproject_target_gm_j_axis_up(self):
         source_ds = create_5x5_dataset_regular_utm()
         target_gm = GridMapping.regular(

tests/test_utils.py

@@ -23,6 +23,7 @@
     _get_spatial_interp_method,
     _prep_spatial_interp_methods_downscale,
     _select_variables,
+    _transpose_dims,
     bbox_overlap,
     clip_dataset_by_bbox,
     get_spatial_coords,
@@ -496,3 +497,24 @@ def test_create_empty_dataset_3d(self):
         np.testing.assert_array_equal(
             np.isnan(target_ds.rad), np.ones_like(target_ds.rad, dtype=bool)
         )
+
+    def test_transpose_dims(self):
+        gm = GridMapping.regular_from_bbox([1000, 1000, 2000, 2000], 100, "EPSG:3035")
+
+        data = xr.Dataset(
+            {
+                "var0": xr.DataArray(np.random.rand(2, 4, 3), dims=("x", "time", "y")),
+                "var1": xr.DataArray(np.random.rand(2, 3), dims=("x", "y")),
+                "var2": xr.DataArray(np.random.rand(4, 3, 2), dims=("time", "y", "x")),
+                "var3": xr.DataArray(np.random.rand(3, 2), dims=("y", "x")),
+                "var4": xr.DataArray(np.random.rand(4, 2), dims=("time", "x")),
+                "var5": xr.DataArray(np.random.rand(2), dims="x"),
+            },
+        )
+        result = _transpose_dims(data, gm)
+        self.assertEqual(("time", "y", "x"), result["var0"].dims)
+        self.assertEqual(("y", "x"), result["var1"].dims)
+        self.assertEqual(("time", "y", "x"), result["var2"].dims)
+        self.assertEqual(("y", "x"), result["var3"].dims)
+        self.assertEqual(("time", "x"), result["var4"].dims)
+        self.assertEqual(("x",), result["var5"].dims)

xcube_resampling/rectify.py

@@ -52,6 +52,7 @@
     _reorganize_tiled_array,
     _sample_array_at_indices,
     _select_variables,
+    _transpose_dims,
     bbox_overlap,
     normalize_grid_mapping,
 )
@@ -134,6 +135,7 @@ def rectify_dataset(
 
     source_gm = source_gm or GridMapping.from_dataset(source_ds)
     source_ds = normalize_grid_mapping(source_ds, source_gm)
+    source_ds = _transpose_dims(source_ds, source_gm)
 
     target_gm = target_gm or source_gm.to_regular(tile_size=tile_size)
 
@@ -214,31 +216,19 @@ def _assemble_rectified_dataset(
     yx_dims = (source_gm.xy_dim_names[1], source_gm.xy_dim_names[0])
     for var_name, da_src in source_ds.data_vars.items():
         if da_src.dims[-2:] == yx_dims:
-            assert len(da_src.dims) in (
-                2,
-                3,
-            ), f"Data variable {var_name} has {len(da_src.dims)} dimensions."
             fill_value = _get_fill_value(fill_values, var_name, da_src)
             interp_method = _get_spatial_interp_method_str(
                 interp_methods, var_name, da_src
             )
-
-            rectified = _rectify_data_array(
-                da_src.data,
+            target_ds[var_name] = _rectify_data_array(
+                da_src,
                 indexing,
                 pixel_target_ij,
                 target_gm,
                 interp_method,
                 fill_value,
             )
-            dims = da_src.dims[:-2] + (
-                target_gm.xy_dim_names[1],
-                target_gm.xy_dim_names[0],
-            )
-            target_ds[var_name] = xr.DataArray(rectified, dims=dims, attrs=da_src.attrs)
-
         elif yx_dims[0] not in da_src.dims and yx_dims[1] not in da_src.dims:
-            # Non-spatial variable → copy as-is
             target_ds[var_name] = da_src
 
     # optional source pixel index outputs
@@ -287,7 +277,7 @@ def _downscale_source_dataset(
     interp_methods: SpatialInterpMethods | None,
     agg_methods: SpatialAggMethods | None,
     prevent_nan_propagations: PreventNaNPropagations,
-) -> (xr.Dataset, GridMapping):
+) -> tuple[xr.Dataset, GridMapping]:
     if interp_methods in [0, "nearest"]:
         return source_ds, source_gm
 
@@ -304,9 +294,9 @@ def _downscale_source_dataset(
         (source_gm.xy_dim_names[1], source_gm.xy_dim_names[0]),
         downscaled_size,
         source_gm.tile_size,
-        _prep_spatial_interp_methods_downscale(interp_methods),
-        agg_methods,
-        prevent_nan_propagations,
+        interp_methods=_prep_spatial_interp_methods_downscale(interp_methods),
+        agg_methods=agg_methods,
+        prevent_nan_propagations=prevent_nan_propagations or True,
     )
     source_gm = GridMapping.from_dataset(source_ds)
 
@@ -661,51 +651,56 @@ def _compute_source_tile_indexing(
 
 
 def _rectify_data_array(
-    data: da.Array,
+    data_array: xr.DataArray,
     indexing: SourceTileIndexing,
     pixel_target_ij: da.Array,
     target_gm: GridMapping,
     interp_method: SpatialInterpMethod,
     fill_value: FloatInt,
-) -> da.Array:
-
-    expanded = data.ndim == 2
-    if expanded:
-        data = data[None, ...]
+) -> xr.DataArray:
 
-    is_numpy_array = False
-    if isinstance(data, np.ndarray):
+    if isinstance(data_array.data, np.ndarray):
         is_numpy_array = True
-        data = da.asarray(data)
+        array = da.asarray(data_array.data)
+    else:
+        is_numpy_array = False
+        array = data_array.data
+
+    # collapse non-spatial dims
+    non_spatial_dims = data_array.dims[:-2]
+    orig_shape = array.shape
+    n_leading = int(np.prod(orig_shape[:-2])) if len(orig_shape) > 2 else 1
+    array_flat = array.reshape((n_leading, orig_shape[-2], orig_shape[-1]))
 
-    tiled = _reorganize_tiled_array(data, indexing, fill_value)
+    # reorganize data array slice to align with the chunks of pixel_target_ij
+    tiled = _reorganize_tiled_array(array_flat, indexing, fill_value)
 
     rectified_chunks = []
     offset = 0
-
-    for chunk in data.chunks[0]:
+    for chunk in tiled.chunks[0]:
         rectified_chunks.append(
             da.map_blocks(
                 _rectify_block,
                 pixel_target_ij,
                 tiled[offset : offset + chunk],
                 interp_method=interp_method,
                 fill_value=fill_value,
-                dtype=data.dtype,
+                dtype=tiled.dtype,
                 chunks=(chunk, *pixel_target_ij.chunks[1:]),
             )
         )
         offset += chunk
-
     result = da.concatenate(rectified_chunks, axis=0)
 
-    if expanded:
-        result = result[0]
+    # restore original non-spatial shape
+    new_shape = orig_shape[:-2] + (target_gm.height, target_gm.width)
+    result = result.reshape(new_shape)
+    dims = non_spatial_dims + (target_gm.xy_dim_names[1], target_gm.xy_dim_names[0])
 
     if is_numpy_array and not target_gm.is_tiled:
         result = result.compute()
 
-    return result
+    return xr.DataArray(data=result, dims=dims, attrs=data_array.attrs)
 
 
 def _rectify_block(