added rewrite_coords

cordex/__init__.py

@@ -2,7 +2,7 @@
 
 from . import regions, tables, tutorial
 from .accessor import CordexDataArrayAccessor, CordexDatasetAccessor  # noqa
-from .domain import cordex_domain, create_dataset, domain, domain_info, vertices
+from .domain import cordex_domain, create_dataset, domain, domain_info, vertices, rewrite_coords
 from .tables import domains, ecmwf
 from .transform import (
     map_crs,
@@ -41,4 +41,5 @@
     "domains",
     "ecmwf",
     "cell_area",
+    "rewrite_coords",
 ]

cordex/cf.py

@@ -186,20 +186,20 @@
                 "long_name": "latitude in rotated pole grid",
                 "units": "degrees",
             },
-            "longitude": {
+            "lon": {
                 "standard_name": "longitude",
                 "long_name": "longitude",
                 "units": "degrees_east",
             },
-            "latitude": {
+            "lat": {
                 "standard_name": "latitude",
                 "long_name": "latitude",
                 "units": "degrees_north",
             },
-            "lon_vertices": {
+            "vertices_lon": {
                 "units": "degrees_east",
             },
-            "lat_vertices": {
+            "vertices_lat": {
                 "units": "degrees_north",
             },
         },

cordex/domain.py

@@ -9,7 +9,7 @@
 from . import cf
 from .config import nround
 from .tables import domains
-from .transform import grid_mapping, transform, transform_bounds
+from .transform import grid_mapping, transform, transform_coords, transform_bounds
 from .utils import cell_area, get_tempfile
 
 
@@ -566,6 +566,78 @@ def vertices(rlon, rlat, src_crs, trg_crs=None):
     return xr.merge([lat_vertices, lon_vertices])
 
 
+def rewrite_coords(ds, coords="xy", domain_id=None, mip_era="CMIP5", method="nearest"):
+    """
+    Rewrite the linear coordinates (X and Y axes) in a dataset to correct rounding errors.
+
+    This function is useful for ensuring that the coordinates in a dataset are consistent and can be compared to other datasets. It can reindex the dataset based on specified coordinates or domain information.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset containing the grid to be rewritten.
+    coords : str, optional
+        Specifies which coordinates to rewrite. Options are:
+        - "xy": Rewrite only the X and Y coordinates.
+        - "lonlat": Rewrite only the longitude and latitude coordinates.
+        - "all": Rewrite both X, Y, longitude, and latitude coordinates.
+        Default is "xy".
+    domain_id : str, optional
+        The domain identifier used to obtain grid information. If not provided, the function will attempt to use the grid mapping information from the dataset.
+    mip_era : str, optional
+        The MIP era (e.g., "CMIP5", "CMIP6") used to determine coordinate attributes. Default is "CMIP5".
+    method : str, optional
+        The method used for reindexing. Options include "nearest", "linear", etc. Default is "nearest".
+
+    Returns
+    -------
+    ds : xr.Dataset
+        The dataset with rewritten coordinates.
+    """
+    if domain_id is None and ds.cf["grid_mapping"].grid_mapping_name == "rotated_latitude_longitude":
+        domain_id = ds.cx.domain_id 
+    if domain_id:
+        grid_info = domain_info(domain_id)
+        dx = grid_info["dlon"]
+        dy = grid_info["dlat"]
+        x0 = grid_info["ll_lon"]
+        y0 = grid_info["ll_lat"]
+        nx = grid_info["nlon"]
+        ny = grid_info["nlat"]
+    else:
+        x = ds.cf["X"].data
+        y = ds.cf["Y"].data
+        nx = x.size
+        ny = y.size
+        dx = (x[1] - x[0]).round(5)
+        dy = (y[1] - y[0]).round(5)
+        x0 = x[0].round(nround)
+        y0 = y[0].round(nround)
+
+    xn = _lin_coord(nx, dx, x0)
+    yn = _lin_coord(ny, dy, y0)
+
+    if coords == "xy" or coords == "all":
+        ds = ds.cf.reindex(X=xn, Y=yn, method=method)
+
+    if coords == "lonlat" or coords == "all":
+        try:
+            trg_dims=(ds.cf["longitude"].name, ds.cf["latitude"].name)
+            overwrite=True
+        except KeyError:
+            trg_dims=("lon", "lat")
+            overwrite=False
+        dst = transform_coords(ds, trg_dims=trg_dims)
+        if overwrite is False:
+            ds = ds.assign_coords({trg_dims[0]: dst[trg_dims[0]], trg_dims[1]: dst[trg_dims[1]]})
+            ds[trg_dims[0]].attrs = cf.vocabulary[mip_era]["coords"][trg_dims[0]]
+            ds[trg_dims[1]].attrs = cf.vocabulary[mip_era]["coords"][trg_dims[1]]
+        else:
+            ds[trg_dims[0]][:] = dst[trg_dims[0]]
+            ds[trg_dims[1]][:] = dst[trg_dims[1]]
+
+    return ds
+
 def _crop_to_domain(ds, domain_id, drop=True):
     domain = cordex_domain(domain_id)
     x_mask = ds.cf["X"].round(8).isin(domain.cf["X"])

cordex/preprocessing/preprocessing.py

@@ -271,7 +271,7 @@ def replace_rlon_rlat(ds, domain=None):
     """
     ds = ds.copy()
     if domain is None:
-        domain = ds.attrs.get("CORDEX_domain", None)
+        domain = ds.cx.domain_id
     dm = cordex_domain(domain)
     for coord in ["rlon", "rlat"]:
         if coord in ds.coords:
@@ -327,7 +327,7 @@ def replace_lon_lat(ds, domain=None):
     """
     ds = ds.copy()
     if domain is None:
-        domain = ds.attrs.get("CORDEX_domain", None)
+        domain = ds.cx.domain_id
     dm = cordex_domain(domain)
     for coord in ["lon", "lat"]:
         if coord in ds.coords: