New function "add_spatial_ref"

Author: forman

test/core/gridmapping/test_cfconv.py

@@ -2,13 +2,20 @@
 import unittest
 from typing import Set
 
+import fsspec
 import numpy as np
 import pyproj
 import xarray as xr
 
+from test.s3test import MOTO_SERVER_ENDPOINT_URL
+from test.s3test import S3Test
+from xcube.core.gridmapping import GridMapping
 from xcube.core.gridmapping.cfconv import GridCoords
 from xcube.core.gridmapping.cfconv import GridMappingProxy
+from xcube.core.gridmapping.cfconv import add_spatial_ref
 from xcube.core.gridmapping.cfconv import get_dataset_grid_mapping_proxies
+from xcube.core.new import new_cube
+from xcube.core.store.fs.registry import new_fs_data_store
 
 CRS_WGS84 = pyproj.crs.CRS(4326)
 CRS_CRS84 = pyproj.crs.CRS.from_string("urn:ogc:def:crs:OGC:1.3:CRS84")
@@ -192,3 +199,76 @@ def _gen_2d(self):
         self.assertEqual(('y', 'x'), lon.dims)
         self.assertEqual(('y', 'x'), lat.dims)
         return noise, crs, lon, lat
+
+
+class AddSpatialRefTest(S3Test):
+
+    def test_add_spatial_ref(self):
+        original_dataset = new_cube(x_name='x',
+                                    y_name='y',
+                                    x_start=0,
+                                    y_start=0,
+                                    x_res=10,
+                                    y_res=10,
+                                    x_units='metres',
+                                    y_units='metres',
+                                    drop_bounds=True,
+                                    width=100,
+                                    height=100,
+                                    variables=dict(A=1.3, B=8.3))
+
+        root = 'eurodatacube-test/xcube-eea'
+        data_id = 'test.zarr'
+        crs_var_name = 'spatial_ref'
+        crs = pyproj.CRS.from_string("EPSG:3035")
+
+        storage_options = dict(
+            anon=False,
+            client_kwargs=dict(
+                endpoint_url=MOTO_SERVER_ENDPOINT_URL,
+            )
+        )
+
+        fs: fsspec.AbstractFileSystem = fsspec.filesystem('s3',
+                                                          **storage_options)
+        if fs.isdir(root):
+            fs.rm(root, recursive=True)
+        fs.mkdirs(root, exist_ok=True)
+
+        data_store = new_fs_data_store('s3',
+                                       root=root,
+                                       storage_options=storage_options)
+
+        data_store.write_data(original_dataset, data_id=data_id)
+        opened_dataset = data_store.open_data(data_id)
+        self.assertEqual({'A', 'B', 'x', 'y', 'time'},
+                         set(opened_dataset.variables))
+
+        with self.assertRaises(ValueError) as cm:
+            GridMapping.from_dataset(opened_dataset)
+        self.assertEqual(
+            ('cannot find any grid mapping in dataset',),
+            cm.exception.args
+        )
+
+        path = f"{root}/{data_id}"
+        group_store = fs.get_mapper(path, create=True)
+
+        add_spatial_ref(group_store,
+                        crs,
+                        crs_var_name=crs_var_name)
+
+        self.assertTrue(fs.exists(f"{path}/{crs_var_name}"))
+        self.assertTrue(fs.exists(f"{path}/{crs_var_name}/.zarray"))
+        self.assertTrue(fs.exists(f"{path}/{crs_var_name}/.zattrs"))
+
+        opened_dataset = data_store.open_data(data_id)
+        self.assertEqual({'A', 'B', 'x', 'y', 'time', 'spatial_ref'},
+                         set(opened_dataset.variables))
+        self.assertEqual(crs_var_name,
+                         opened_dataset.A.attrs.get('grid_mapping'))
+        self.assertEqual(crs_var_name,
+                         opened_dataset.B.attrs.get('grid_mapping'))
+
+        gm = GridMapping.from_dataset(opened_dataset)
+        self.assertIn("LAEA Europe", gm.crs.srs)

xcube/core/gridmapping/cfconv.py

@@ -22,8 +22,11 @@
 import warnings
 from typing import Optional, Dict, Any, Hashable, Union, Set, List, Tuple
 
+import numpy as np
 import pyproj
 import xarray as xr
+import zarr
+import zarr.convenience
 
 from xcube.core.schema import get_dataset_chunks
 
@@ -285,3 +288,38 @@ def _find_dataset_tile_size(dataset: xr.Dataset,
     if tile_width is not None and tile_height is not None:
         return tile_width, tile_height
     return None
+
+
+def add_spatial_ref(group_store: zarr.convenience.StoreLike,
+                    crs: pyproj.CRS,
+                    crs_var_name: Optional[str] = 'spatial_ref',
+                    xy_dim_names: Optional[Tuple[str, str]] = None):
+    """
+    Helper function that allows adding a spatial reference to an
+    existing Zarr dataset.
+
+    :param group_store: The dataset's existing Zarr store or path.
+    :param crs: The coordinate reference system.
+    :param crs_var_name: The name of the variable that will hold the
+        spatial reference. Defaults to "spatial_ref".
+    :param xy_dim_names: The names of the x and y dimensions.
+        Defaults to ("x", "y").
+    """
+    spatial_attrs = crs.to_cf()
+    spatial_attrs['_ARRAY_DIMENSIONS'] = []  # Required by xarray
+    group = zarr.open(group_store, mode='r+')
+    spatial_ref = group.array(crs_var_name,
+                              0,
+                              shape=(),
+                              dtype=np.uint8,
+                              fill_value=0)
+    spatial_ref.attrs.update(**spatial_attrs)
+
+    yx_dim_names = list(reversed(xy_dim_names or ('x', 'y')))
+    for item_name, item in group.items():
+        if item_name != crs_var_name:
+            dims = item.attrs.get('_ARRAY_DIMENSIONS')
+            if dims and len(dims) >= 2 and dims[-2:] == yx_dim_names:
+                item.attrs['grid_mapping'] = crs_var_name
+
+    zarr.convenience.consolidate_metadata(group_store)