tweaked grib to_xarray

climetlab/sources/readers/grib.py

@@ -14,6 +14,7 @@
 
 from . import Reader
 from climetlab.utils.bbox import BoundingBox
+from climetlab.utils.grib import post_xarray_open_dataset_hook, open_dataset_params
 
 LOG = logging.getLogger(__name__)
 
@@ -197,4 +198,6 @@ def __len__(self):
     def to_xarray(self):
         import xarray as xr
 
-        return xr.open_dataset(self.path, engine="cfgrib")
+        params = open_dataset_params()
+        ds = xr.open_dataset(self.path, engine="cfgrib", **params)
+        return post_xarray_open_dataset_hook(ds)

climetlab/sources/zarr_s3.py

@@ -62,6 +62,7 @@ def url_to_store(url):
 
         #        self._ds = xr.open_zarr(cache)
         self._ds = xr.open_mfdataset(stores, engine="zarr")
+        # self._ds = xr.open_mfdataset(stores, engine="zarr", combine='nested')
 
     def to_xarray(self):
         return self._ds

climetlab/utils/grib.py

@@ -0,0 +1,54 @@
+def post_xarray_open_dataset_hook(ds):
+    # we may want to do this too :
+    # import cf2cdm
+    # ds = cf2cdm.translate_coords(ds, cf2cdm.CDS)
+    # or
+    # ds = cf2cdm.translate_coords(ds, cf2cdm.ECMWF)
+
+    # ensure our internal conventions
+    if "number" in list(ds.coords):
+        ds = ds.rename({"number": "realization"})
+    if "time" in list(ds.coords):
+        ds = ds.rename({"time": "forecast_time"})
+    if "valid_time" in list(ds.coords):
+        ds = ds.rename({"valid_time": "time"})
+    if "heightAboveGround" in list(ds.coords):
+        # if we decide to keep it, rename it.
+        # ds = ds.rename({'heightAboveGround':'height_above_ground'})
+        ds = ds.squeeze("heightAboveGround")
+        ds = ds.drop("heightAboveGround")
+    if "surface" in list(ds.coords):
+        ds = ds.squeeze("surface")
+        ds = ds.drop("surface")
+    return ds
+
+def open_dataset_params(time_convention="withstep"):
+    params = {}
+    if time_convention == "withstep":
+        time_dims = ["time", "step"]  # this is the default of engine='cfgrib'
+        chunk_sizes_in = {
+            "time": 1,
+            "latitude": None,
+            "longitude": None,
+            "number": 1,
+            "step": 1,
+        }
+
+    elif time_convention == "nostep":
+        time_dims = ["time", "valid_time"]
+        chunk_sizes_in = {
+            "time": 1,
+            "latitude": None,
+            "longitude": None,
+            "number": 1,
+            "valid_time": 1,
+        }
+
+    else:
+        raise Exception()
+
+    # params['engine'] = 'cfgrib'
+    params["chunks"] = chunk_sizes_in
+    params["backend_kwargs"] = dict(squeeze=False, time_dims=time_dims)
+
+    return params