Enhance xcube affine_transform_dataset()

CHANGES.md

@@ -20,6 +20,12 @@
 
 ### Other changes
 
+* `xcube.core.resampling.affine_transform_dataset()` has a new 
+  keyword argument `reuse_coords: bool = False`. If set to `True` 
+  the returned dataset will reuse the _same_ spatial coordinates 
+  as the target. This is a workaround for xarray issue 
+  https://github.com/pydata/xarray/issues/6573.
+
 * Deprecated following functions of module `xcube.core.geom`:
   - `is_dataset_y_axis_inverted()` is no longer used;
   - `get_geometry_mask()` is no longer used;

test/core/gridmapping/test_coords.py

@@ -10,14 +10,17 @@
 # noinspection PyProtectedMember
 
 GEO_CRS = pyproj.crs.CRS(4326)
+NOT_A_GEO_CRS = pyproj.crs.CRS(5243)
 
 
 # noinspection PyMethodMayBeStatic
 class Coords1DGridMappingTest(unittest.TestCase):
 
     def test_1d_j_axis_down(self):
-        gm = GridMapping.from_coords(x_coords=xr.DataArray(np.linspace(1.5, 8.5, 8), dims='lon'),
-                                     y_coords=xr.DataArray(np.linspace(4.5, -4.5, 10), dims='lat'),
+        x_coords = xr.DataArray(np.linspace(1.5, 8.5, 8), dims='lon')
+        y_coords = xr.DataArray(np.linspace(4.5, -4.5, 10), dims='lat')
+        gm = GridMapping.from_coords(x_coords=x_coords,
+                                     y_coords=y_coords,
                                      crs=GEO_CRS)
         self.assertEqual((8, 10), gm.size)
         self.assertEqual((8, 10), gm.tile_size)
@@ -27,6 +30,12 @@ def test_1d_j_axis_down(self):
         self.assertEqual(True, gm.is_regular)
         self.assertEqual(False, gm.is_j_axis_up)
         self.assertEqual(False, gm.is_lon_360)
+        self.assertTrue(hasattr(gm, 'x_coords'))
+        self.assertTrue(hasattr(gm, 'y_coords'))
+        # noinspection PyUnresolvedReferences
+        self.assertIs(x_coords, gm.x_coords)
+        # noinspection PyUnresolvedReferences
+        self.assertIs(y_coords, gm.y_coords)
 
     def test_1d_j_axis_up(self):
         gm = GridMapping.from_coords(x_coords=xr.DataArray(np.linspace(1.5, 8.5, 8), dims='lon'),
@@ -121,21 +130,39 @@ def test_1d_xy_coords(self):
         self.assertEqual(('lon', 'lat'), gm.xy_var_names)
         self.assertEqual(('lon', 'lat'), gm.xy_dim_names)
 
+    def test_to_coords(self):
+        gm = GridMapping.regular(size=(10, 6),
+                                 xy_min=(-2600.0, 1200.0),
+                                 xy_res=10.0,
+                                 crs=NOT_A_GEO_CRS)
+        cv = gm.to_coords(reuse_coords=False)
+        self.assertIn("x", cv)
+        self.assertIn("y", cv)
+        self.assertEqual(np.float64, cv["x"].dtype)
+        self.assertEqual(np.float64, cv["y"].dtype)
+
+        gm2 = GridMapping.from_coords(cv['x'].astype(np.float32),
+                                      cv['y'].astype(np.float32),
+                                      gm.crs)
+        cv2 = gm2.to_coords(xy_var_names=("a", "b"),
+                            xy_dim_names=("u", "v"),
+                            reuse_coords=True)
+        self.assertIn("a", cv2)
+        self.assertIn("b", cv2)
+        self.assertEqual(np.float32, cv2["a"].dtype)
+        self.assertEqual(np.float32, cv2["b"].dtype)
+
 
 class Coords2DGridMappingTest(unittest.TestCase):
 
     def test_2d(self):
+        x_coords = xr.DataArray([[10.0, 10.1, 10.2, 10.3], [10.1, 10.2, 10.3, 10.4], [10.2, 10.3, 10.4, 10.5], ],
+                                dims=('lat', 'lon'))
+        y_coords = xr.DataArray([[52.0, 52.2, 52.4, 52.6], [52.2, 52.4, 52.6, 52.8], [52.4, 52.6, 52.8, 53.0], ],
+                                dims=('lat', 'lon'))
         gm = GridMapping.from_coords(
-            x_coords=xr.DataArray([
-                [10.0, 10.1, 10.2, 10.3],
-                [10.1, 10.2, 10.3, 10.4],
-                [10.2, 10.3, 10.4, 10.5],
-            ], dims=('lat', 'lon')),
-            y_coords=xr.DataArray([
-                [52.0, 52.2, 52.4, 52.6],
-                [52.2, 52.4, 52.6, 52.8],
-                [52.4, 52.6, 52.8, 53.0],
-            ], dims=('lat', 'lon')),
+            x_coords=x_coords,
+            y_coords=y_coords,
             crs=GEO_CRS)
         self.assertEqual((4, 3), gm.size)
         self.assertEqual((4, 3), gm.tile_size)
@@ -145,6 +172,12 @@ def test_2d(self):
         self.assertEqual(False, gm.is_regular)
         self.assertEqual(True, gm.is_j_axis_up)
         self.assertEqual(False, gm.is_lon_360)
+        self.assertTrue(hasattr(gm, 'x_coords'))
+        self.assertTrue(hasattr(gm, 'y_coords'))
+        # noinspection PyUnresolvedReferences
+        self.assertIs(x_coords, gm.x_coords)
+        # noinspection PyUnresolvedReferences
+        self.assertIs(y_coords, gm.y_coords)
 
     def test_2d_tile_size_from_chunks(self):
         gm = GridMapping.from_coords(

xcube/core/gridmapping/base.py

@@ -540,7 +540,9 @@ def ij_bboxes_from_xy_bboxes(self,
     def to_coords(self,
                   xy_var_names: Tuple[str, str] = None,
                   xy_dim_names: Tuple[str, str] = None,
-                  exclude_bounds: bool = False) \
+                  exclude_bounds: bool = False,
+                  reuse_coords: bool = False,
+                  ) \
             -> Mapping[str, xr.DataArray]:
         """
         Get CF-compliant axis coordinate variables and cell boundary
@@ -554,14 +556,18 @@ def to_coords(self,
             names (x_dim_name, y_dim_name).
         :param exclude_bounds: If True, do not create bounds
             coordinates. Defaults to False.
+        :param reuse_coords: Whether to either reuse target
+            coordinate arrays from target_gm or to compute
+            new ones.
         :return: dictionary with coordinate variables
         """
         self._assert_regular()
         from .coords import grid_mapping_to_coords
         return grid_mapping_to_coords(self,
                                       xy_var_names=xy_var_names,
                                       xy_dim_names=xy_dim_names,
-                                      exclude_bounds=exclude_bounds)
+                                      exclude_bounds=exclude_bounds,
+                                      reuse_coords=reuse_coords)
 
     def transform(self,
                   crs: Union[str, pyproj.crs.CRS],

xcube/core/gridmapping/coords.py

@@ -59,6 +59,14 @@ def __init__(self,
         self._y_coords = y_coords
         super().__init__(**kwargs)
 
+    @property
+    def x_coords(self):
+        return self._x_coords
+
+    @property
+    def y_coords(self):
+        return self._y_coords
+
 
 class Coords1DGridMapping(CoordsGridMapping):
     """
@@ -288,7 +296,8 @@ def grid_mapping_to_coords(
         grid_mapping: GridMapping,
         xy_var_names: Tuple[str, str] = None,
         xy_dim_names: Tuple[str, str] = None,
-        exclude_bounds: bool = False
+        reuse_coords: bool = False,
+        exclude_bounds: bool = False,
 ) -> Dict[str, xr.DataArray]:
     """
     Get CF-compliant axis coordinate variables and cell
@@ -301,8 +310,12 @@ def grid_mapping_to_coords(
         names (x_var_name, y_var_name).
     :param xy_dim_names: Optional coordinate dimensions
         names (x_dim_name, y_dim_name).
+    :param reuse_coords: Whether to either reuse target
+        coordinate arrays from target_gm or to compute
+        new ones.
     :param exclude_bounds: If True, do not create
-        bounds coordinates. Defaults to False.
+        bounds coordinates. Defaults to False. Ignored if
+        *reuse_coords* is True.
     :return: dictionary with coordinate variables
     """
 
@@ -311,6 +324,27 @@ def grid_mapping_to_coords(
     if xy_dim_names:
         _assert_valid_xy_names(xy_dim_names, name='xy_dim_names')
 
+    if reuse_coords:
+        try:
+            # noinspection PyUnresolvedReferences
+            x, y = grid_mapping.x_coords, grid_mapping.y_coords
+        except AttributeError:
+            x, y = None, None
+        if isinstance(x, xr.DataArray) \
+                and isinstance(y, xr.DataArray) \
+                and x.ndim == 1 \
+                and y.ndim == 1 \
+                and x.size == grid_mapping.width \
+                and y.size == grid_mapping.height:
+            return {
+                name: xr.DataArray(coord.values,
+                                   dims=dim,
+                                   attrs=coord.attrs)
+                for name, dim, coord in zip(xy_var_names,
+                                            xy_dim_names,
+                                            (x, y))
+            }
+
     x_name, y_name = xy_var_names or grid_mapping.xy_var_names
     x_dim_name, y_dim_name = xy_dim_names or grid_mapping.xy_dim_names
     w, h = grid_mapping.size

xcube/core/resampling/affine.py

@@ -40,7 +40,8 @@ def affine_transform_dataset(
         dataset: xr.Dataset,
         source_gm: GridMapping,
         target_gm: GridMapping,
-        var_configs: Mapping[Hashable, Mapping[str, Any]] = None
+        var_configs: Mapping[Hashable, Mapping[str, Any]] = None,
+        reuse_coords: bool = False,
 ) -> xr.Dataset:
     """
     Resample dataset according to an affine transformation.
@@ -52,6 +53,9 @@ def affine_transform_dataset(
         Must be regular. Must have same CRS as *source_gm*.
     :param var_configs: Optional resampling configurations
         for individual variables.
+    :param reuse_coords: Whether to either reuse target
+        coordinate arrays from target_gm or to compute
+        new ones.
     :return: The resampled target dataset.
     """
     if source_gm.crs != target_gm.crs:
@@ -73,7 +77,8 @@ def affine_transform_dataset(
     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
+        exclude_bounds=not has_bounds,
+        reuse_coords=reuse_coords
     )
     return resampled_dataset.assign_coords(new_coords)