Merge pull request #1003 from xcube-dev/forman-1001-resample_in_space_issues

Address issues with `resample_in_space()`

Author: forman

CHANGES.md

@@ -53,9 +53,19 @@
   A new keyword argument `max_depth` defines the maximum subdirectory depths 
   used to search for datasets in case `roots` is given. It defaults to `1`.
 
+* The behaviour of function `resample_in_space()` of module 
+  `xcube.core.resampling` changed in this version. (#1001)
+  1. If the resampling boils down to a plain affine transformation,
+     the returned target dataset will now have the _same_ spatial coordinates
+     as the provided target grid mapping passed by `target_gm`.
+  2. In the case of up-sampling, we no longer recover `NaN` values by default
+     as it may require considerable CPU overhead.
+     To enforce the old behaviour, provide the `var_configs` keyword-argument
+     and set `recover_nan` to `True` for desired variables.
+
 * The class `MaskSet()` of module `xcube.core.maskset` now correctly recognises
   the variable attributes `flag_values`, `flag_masks`, `flag_meanings` when
-  their values are lists (ESA CCI LC data encodes them as JSON arrays).
+  their values are lists (ESA CCI LC data encodes them as JSON arrays). (#1002)
 
 
 ### Incompatible API changes

test/core/resampling/test_affine.py

@@ -46,7 +46,7 @@ class AffineTransformDatasetTest(unittest.TestCase):
     def test_subset(self):
         target_gm = GridMapping.regular((3, 3), (50.0, 10.0), res, source_gm.crs)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertIsInstance(target_ds, xr.Dataset)
         self.assertEqual(
@@ -66,7 +66,7 @@ def test_subset(self):
 
         target_gm = GridMapping.regular((3, 3), (50.1, 10.1), res, source_gm.crs)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertIsInstance(target_ds, xr.Dataset)
         self.assertEqual(
@@ -86,7 +86,7 @@ def test_subset(self):
 
         target_gm = GridMapping.regular((3, 3), (50.05, 10.05), res, source_gm.crs)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertIsInstance(target_ds, xr.Dataset)
         self.assertEqual(
@@ -102,13 +102,15 @@ def test_different_geographic_crses(self):
         expected = np.array([[1.25, 1.5, 0.75], [1.0, 1.25, 1.5], [1.75, 1.0, 1.25]])
 
         target_gm = GridMapping.regular((3, 3), (50.05, 10.05), res, CRS_WGS84)
-        target_ds = affine_transform_dataset(source_ds, source_gm, target_gm)
+        target_ds = affine_transform_dataset(
+            source_ds, source_gm=source_gm, target_gm=target_gm
+        )
         self.assertEqual((3, 3), target_ds.refl.shape)
         np.testing.assert_almost_equal(target_ds.refl.values, expected)
 
         target_gm = GridMapping.regular((3, 3), (50.05, 10.05), res, CRS_CRS84)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertEqual((3, 3), target_ds.refl.shape)
         np.testing.assert_almost_equal(target_ds.refl.values, expected)
@@ -122,12 +124,14 @@ def test_different_geographic_crses(self):
             ' must be equal, was "WGS 84" and'
             ' "ETRS89-extended / LAEA Europe"',
         ):
-            affine_transform_dataset(source_ds, source_gm, target_gm)
+            affine_transform_dataset(
+                source_ds, source_gm=source_gm, target_gm=target_gm
+            )
 
     def test_downscale_x2(self):
         target_gm = GridMapping.regular((8, 6), (50, 10), 2 * res, source_gm.crs)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertIsInstance(target_ds, xr.Dataset)
         self.assertEqual(
@@ -152,7 +156,7 @@ def test_downscale_x2(self):
     def test_downscale_x2_and_shift(self):
         target_gm = GridMapping.regular((8, 6), (49.8, 9.8), 2 * res, source_gm.crs)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertIsInstance(target_ds, xr.Dataset)
         self.assertEqual(
@@ -177,7 +181,7 @@ def test_downscale_x2_and_shift(self):
     def test_upscale_x2(self):
         target_gm = GridMapping.regular((8, 6), (50, 10), res / 2, source_gm.crs)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertIsInstance(target_ds, xr.Dataset)
         self.assertEqual(
@@ -202,7 +206,7 @@ def test_upscale_x2(self):
     def test_upscale_x2_and_shift(self):
         target_gm = GridMapping.regular((8, 6), (49.9, 9.95), res / 2, source_gm.crs)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertIsInstance(target_ds, xr.Dataset)
         self.assertEqual(
@@ -227,7 +231,7 @@ def test_upscale_x2_and_shift(self):
     def test_shift(self):
         target_gm = GridMapping.regular((8, 6), (50.2, 10.1), res, source_gm.crs)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertIsInstance(target_ds, xr.Dataset)
         self.assertEqual(
@@ -251,7 +255,7 @@ def test_shift(self):
 
         target_gm = GridMapping.regular((8, 6), (49.8, 9.9), res, source_gm.crs)
         target_ds = affine_transform_dataset(
-            source_ds, source_gm, target_gm, gm_name="crs"
+            source_ds, source_gm=source_gm, target_gm=target_gm, gm_name="crs"
         )
         self.assertIsInstance(target_ds, xr.Dataset)
         self.assertEqual(

test/core/resampling/test_spatial.py

@@ -25,7 +25,11 @@ def test_affine_transform_dataset(self):
         )
 
         target_ds = resample_in_space(
-            source_ds, source_gm, target_gm, encode_cf=True, gm_name="crs"
+            source_ds,
+            source_gm=source_gm,
+            target_gm=target_gm,
+            encode_cf=True,
+            gm_name="crs",
         )
 
         self.assertIn("crs", target_ds)

xcube/core/resampling/affine.py

@@ -3,7 +3,7 @@
 # https://opensource.org/licenses/MIT.
 
 import math
-from typing import Union, Callable, Optional, Tuple, Any
+from typing import Union, Callable, Optional, Any
 from collections.abc import Sequence, Mapping, Hashable
 
 import numpy as np
@@ -14,17 +14,19 @@
 from xcube.core.gridmapping import GridMapping
 from xcube.core.gridmapping.helpers import AffineTransformMatrix
 from xcube.util.assertions import assert_true
-from .cf import maybe_encode_grid_mapping
+from .cf import complete_resampled_dataset
 
 NDImage = Union[np.ndarray, da.Array]
 Aggregator = Callable[[NDImage], NDImage]
 
 
 def affine_transform_dataset(
-    dataset: xr.Dataset,
-    source_gm: GridMapping,
-    target_gm: GridMapping,
-    var_configs: Mapping[Hashable, Mapping[str, Any]] = None,
+    source_ds: xr.Dataset,
+    *,
+    target_ds: Optional[xr.Dataset] = None,
+    source_gm: Optional[GridMapping] = None,
+    target_gm: Optional[GridMapping] = None,
+    var_configs: Optional[Mapping[Hashable, Mapping[str, Any]]] = None,
     encode_cf: bool = True,
     gm_name: Optional[str] = None,
     reuse_coords: bool = False,
@@ -35,12 +37,22 @@ def affine_transform_dataset(
     *source_gm* and the CRS of *target_gm* are both geographic or equal.
     Otherwise, a ``ValueError`` will be raised.
 
+    New in 1.6: If *target_ds* is given, its coordinate
+    variables are copied by reference into the returned
+    dataset.
+
     Args:
-        dataset: The source dataset
-        source_gm: Source grid mapping of *dataset*. Must be regular.
-            Must have same CRS as *target_gm*.
-        target_gm: Target grid mapping. Must be regular. Must have same
-            CRS as *source_gm*.
+        source_ds: The source dataset
+        target_ds: An optional dataset that provides the
+            target grid mapping, if *target_gm* is not provided.
+            The coordinate variables of *target_dataset* are copied
+            by reference into the returned dataset.
+        source_gm: Optional source grid mapping of *dataset*.
+            If not provided, computed from *source_ds*.
+            Must be regular and must have same CRS as *target_gm*.
+        target_gm: Optional target grid mapping. If not provided,
+            computed from *target_ds* or source grid mapping.
+            Must be regular and must have same CRS as source grid mapping.
         var_configs: Optional resampling configurations for individual
             variables.
         encode_cf: Whether to encode the target grid mapping into the
@@ -54,6 +66,18 @@ def affine_transform_dataset(
     Returns:
         The resampled target dataset.
     """
+    if source_gm is None:
+        # No source grid mapping given, so do derive it from dataset
+        source_gm = GridMapping.from_dataset(source_ds)
+
+    if target_gm is None:
+        # No target grid mapping given, so do derive it
+        # from target dataset or source grid mapping
+        if target_ds is not None:
+            target_gm = GridMapping.from_dataset(target_ds)
+        else:
+            target_gm = source_gm.to_regular()
+
     # Are source and target both geographic grid mappings?
     both_geographic = source_gm.crs.is_geographic and target_gm.crs.is_geographic
     if not (both_geographic or source_gm.crs == target_gm.crs):
@@ -65,24 +89,28 @@ def affine_transform_dataset(
     GridMapping.assert_regular(source_gm, name="source_gm")
     GridMapping.assert_regular(target_gm, name="target_gm")
     resampled_dataset = resample_dataset(
-        dataset=dataset,
+        dataset=source_ds,
         matrix=target_gm.ij_transform_to(source_gm),
         size=target_gm.size,
         tile_size=target_gm.tile_size,
         xy_dim_names=source_gm.xy_dim_names,
         var_configs=var_configs,
     )
     has_bounds = any(
-        dataset[var_name].attrs.get("bounds") for var_name in source_gm.xy_var_names
+        source_ds[var_name].attrs.get("bounds") for var_name in source_gm.xy_var_names
     )
     new_coords = target_gm.to_coords(
         xy_var_names=source_gm.xy_var_names,
         xy_dim_names=source_gm.xy_dim_names,
         exclude_bounds=not has_bounds,
         reuse_coords=reuse_coords,
     )
-    return maybe_encode_grid_mapping(
-        encode_cf, resampled_dataset.assign_coords(new_coords), target_gm, gm_name
+    return complete_resampled_dataset(
+        encode_cf,
+        resampled_dataset.assign_coords(new_coords),
+        target_gm,
+        gm_name,
+        target_ds.coords if target_ds else None,
     )
 
 
@@ -127,7 +155,8 @@ def resample_dataset(
             else:
                 spline_order = 1
                 aggregator = np.nanmean
-                recover_nan = True
+                # forman: changed default from True to False (v1.6, 2024-06-05)
+                recover_nan = False
             var_data = resample_ndimage(
                 var.data,
                 scale=(j_scale, i_scale),
@@ -260,7 +289,7 @@ def _transform_array(
             # Yes, then
             # 1. replace NaN by zero
             filled_im = da.where(mask, 0.0, image)
-            # 2. transform the zeo-filled image
+            # 2. transform the zero-filled image
             scaled_im = ndinterp.affine_transform(
                 filled_im, matrix, **at_kwargs, cval=0.0
             )
@@ -306,17 +335,19 @@ def _normalize_image(im: NDImage) -> da.Array:
     return da.asarray(im)
 
 
-def _normalize_offset(offset: Optional[Sequence[float]], ndim: int) -> tuple[int, ...]:
+def _normalize_offset(
+    offset: Optional[Sequence[float]], ndim: int
+) -> tuple[float, ...]:
     return _normalize_pair(offset, 0.0, ndim, "offset")
 
 
-def _normalize_scale(scale: Optional[Sequence[float]], ndim: int) -> tuple[int, ...]:
+def _normalize_scale(scale: Optional[Sequence[float]], ndim: int) -> tuple[float, ...]:
     return _normalize_pair(scale, 1.0, ndim, "scale")
 
 
 def _normalize_pair(
     pair: Optional[Sequence[float]], default: float, ndim: int, name: str
-) -> tuple[int, ...]:
+) -> tuple[float, ...]:
     if pair is None:
         pair = [default, default]
     elif isinstance(pair, (int, float)):

xcube/core/resampling/cf.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2018-2024 by xcube team and contributors
 # Permissions are hereby granted under the terms of the MIT License:
 # https://opensource.org/licenses/MIT.
-
+from collections.abc import Hashable, Mapping
 from typing import Optional
 
 import xarray as xr
@@ -82,12 +82,29 @@ def encode_grid_mapping(
     return ds_copy
 
 
-def maybe_encode_grid_mapping(
-    encode_cf: bool, ds: xr.Dataset, gm: GridMapping, gm_name: Optional[str]
+def complete_resampled_dataset(
+    encode_cf: bool,
+    ds: xr.Dataset,
+    gm: GridMapping,
+    gm_name: Optional[str],
+    coords: Optional[Mapping[Hashable, xr.DataArray]],
 ) -> xr.Dataset:
     """Internal helper."""
     if encode_cf:
-        return encode_grid_mapping(
+        ds = encode_grid_mapping(
             ds, gm, gm_name=gm_name, force=True if gm_name else None
         )
+    if coords:
+        compatible_coords = {
+            k: v for k, v in coords.items() if is_var_compatible(v, ds)
+        }
+        ds = ds.assign_coords(compatible_coords)
     return ds
+
+
+def is_var_compatible(var: xr.DataArray, ds: xr.Dataset):
+    """Internal helper."""
+    for d in var.dims:
+        if d in ds.sizes and var.sizes[d] != ds.sizes[d]:
+            return False
+    return True