generalize dims as X/Y when lat/lon or rlat/rlon

.zenodo.json

@@ -34,6 +34,11 @@
       "name": "Caron, Louis-Philippe",
       "affiliation": "Ouranos",
       "orcid": "0000-0001-5221-0147"
+    },
+    {
+      "name": "Braun, Marco",
+      "affiliation": "Ouranos",
+      "orcid": "0000-0001-5061-3217"
     }
   ],
   "keywords": [

AUTHORS.rst

@@ -21,3 +21,4 @@ Contributors
 * Louis-Philippe Caron <caron.louis-philippe@ouranos.ca>
 * Sarah Gammon <gammon.sarah@ouranos.ca> `@SarahG-579462 <https://github.com/SarahG-579462>`_
 * Yannick Rousseau
+* Marco Braun <Braun.Marco@ouranos.ca> `@vindelico <https://github.com/vindelico>`_

HISTORY.rst

@@ -16,6 +16,7 @@ New features and enhancements
 * New function `get_warming_level` to search within the IPCC CMIP global temperatures CSV without requiring data. (:issue:`208`, :pull:`210`).
 * File re-structuration from catalogs with ``xscen.catutils.build_path``. (:pull:`205`).
 * New scripting functions `save_and_update` and `move_and_delete`. (:pull:`214`).
+* Spatial dimensions can be generalized as X/Y when rechunking and will be mapped to rlon/rlat or lon/lat accordingly. (:pull: `221`).
 
 Breaking changes
 ^^^^^^^^^^^^^^^^

xscen/io.py

@@ -304,22 +304,7 @@ def save_to_netcdf(
     xarray.Dataset.to_netcdf
     """
     if rechunk:
-        for rechunk_var in ds.data_vars:
-            # Support for chunks varying per variable
-            if rechunk_var in rechunk:
-                rechunk_dims = rechunk[rechunk_var]
-            else:
-                rechunk_dims = rechunk
-
-            ds[rechunk_var] = ds[rechunk_var].chunk(
-                {
-                    d: chnks
-                    for d, chnks in rechunk_dims.items()
-                    if d in ds[rechunk_var].dims
-                }
-            )
-            ds[rechunk_var].encoding.pop("chunksizes", None)
-            ds[rechunk_var].encoding.pop("chunks", None)
+        ds = _rechunk_for_saving(ds, rechunk)
 
     path = Path(filename)
     path.parent.mkdir(parents=True, exist_ok=True)
@@ -404,22 +389,7 @@ def save_to_zarr(
             ds[v].encoding.clear()
 
     if rechunk:
-        for rechunk_var in ds.data_vars:
-            # Support for chunks varying per variable
-            if rechunk_var in rechunk:
-                rechunk_dims = rechunk[rechunk_var]
-            else:
-                rechunk_dims = rechunk
-
-            ds[rechunk_var] = ds[rechunk_var].chunk(
-                {
-                    d: chnks
-                    for d, chnks in rechunk_dims.items()
-                    if d in ds[rechunk_var].dims
-                }
-            )
-            ds[rechunk_var].encoding.pop("chunksizes", None)
-            ds[rechunk_var].encoding.pop("chunks", None)
+        ds = _rechunk_for_saving(ds, rechunk)
 
     path = Path(filename)
     path.parent.mkdir(parents=True, exist_ok=True)
@@ -518,6 +488,30 @@ def coerce_attrs(attrs):
             raise
 
 
+def _rechunk_for_saving(ds, rechunk):
+    """Rechunk before saving to .zarr or .nc, generalized as Y/X for different axes lat/lon, rlat/rlon."""
+    for rechunk_var in ds.data_vars:
+        # Support for chunks varying per variable
+        if rechunk_var in rechunk:
+            rechunk_dims = rechunk[rechunk_var].copy()
+        else:
+            rechunk_dims = rechunk.copy()
+
+        # get actual axes labels
+        if "X" in rechunk_dims and "X" not in ds.dims:
+            rechunk_dims[ds.cf.axes["X"][0]] = rechunk_dims.pop("X")
+        if "Y" in rechunk_dims and "Y" not in ds.dims:
+            rechunk_dims[ds.cf.axes["Y"][0]] = rechunk_dims.pop("Y")
+
+        ds[rechunk_var] = ds[rechunk_var].chunk(
+            {d: chnks for d, chnks in rechunk_dims.items() if d in ds[rechunk_var].dims}
+        )
+        ds[rechunk_var].encoding.pop("chunksizes", None)
+        ds[rechunk_var].encoding.pop("chunks", None)
+
+    return ds
+
+
 @parse_config
 def rechunk(
     path_in: Union[os.PathLike, str, xr.Dataset],