diff --git a/cordex/__init__.py b/cordex/__init__.py index 532d739..639a83c 100644 --- a/cordex/__init__.py +++ b/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", ] diff --git a/cordex/cf.py b/cordex/cf.py index 0ee3ea0..512fe67 100644 --- a/cordex/cf.py +++ b/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", }, }, diff --git a/cordex/domain.py b/cordex/domain.py index f990fd8..78a7f2d 100644 --- a/cordex/domain.py +++ b/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"]) diff --git a/cordex/preprocessing/preprocessing.py b/cordex/preprocessing/preprocessing.py index 125878f..f1835a6 100644 --- a/cordex/preprocessing/preprocessing.py +++ b/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: