Improve renaming logic to prevent million warnings (#280)

* Improve renaming logic to prevent million warnings

* Reactivate warnings

* Eliminate warnings about renaming dims

* Dont rename dim coordinates again

* Testing

* Testing

* remove dims after renaming

* another try

* fix

* Remove prints

* Fix warning dataset_id

tests/test_preprocessing.py

@@ -472,19 +472,51 @@ def test_combined_preprocessing_dropped_coords(add_coords, shift):
     assert "bnds" not in ds.coords
 
 
-def test_combined_preprocessing_mislabeled_coords():
-    """Test if the renaming is applied to datavariables and then if they are moved to the coords."""
+def test_rename_mislabeled_coords():
+    """Test if the renaming is applied to datavariables"""
     # create a 2d dataset
-    xlen, ylen, zlen = (10, 5, 1)
-    ds = (
-        create_test_ds("x", "y", "dummy", xlen, ylen, zlen).squeeze().drop_vars("dummy")
+    xlen, ylen, zlen = (10, 5, 3)
+    ds = create_test_ds("x", "y", "z", xlen, ylen, zlen).squeeze()
+    ds["nav_lon"] = ds.lon  # assign longitude as data variable
+    ds = ds.drop_vars(["lon"])
+
+    ds_pp = rename_cmip6(ds)
+    np.testing.assert_allclose(ds.nav_lon.data, ds_pp.lon.data)
+
+
+def test_duplicate_renamed_coordinates():
+    # create a 2d dataset
+    xlen, ylen, zlen = (10, 5, 3)
+    ds = create_test_ds("x", "y", "lev", xlen, ylen, zlen)
+    ds = ds.drop_vars("lon")  # drop the original longitude
+    # assign two coordinates which should both be renamed according to the renaming dict
+    coord_da_1 = xr.DataArray(np.random.rand(xlen, ylen), dims=["x", "y"])
+    coord_da_2 = xr.DataArray(np.random.rand(xlen, ylen), dims=["x", "y"])
+    ds = ds.assign_coords(longitude=coord_da_1, nav_lon=coord_da_2)
+    print(ds)
+    with pytest.warns(
+        match="While renaming to target `lon`, more than one candidate was found"
+    ):
+        ds_pp = rename_cmip6(ds)
+
+    assert "nav_lon" in ds_pp.coords
+    xr.testing.assert_allclose(
+        ds_pp.lon.reset_coords(drop=True).drop(["x", "y"]), coord_da_1
     )
-    ds = ds.assign(depth=5.0)
-    ds.depth.attrs["units"] = "m"  # otherwise pint complains.
 
-    ds_pp = combined_preprocessing(ds)
-    assert "lev" in ds_pp.coords
-    np.testing.assert_allclose(ds.depth.data, ds_pp.lev.data)
+
+def test_renamed_coordinate_exists():
+    # create a 2d dataset
+    xlen, ylen, zlen = (10, 5, 3)
+    ds = create_test_ds("x", "y", "lev", xlen, ylen, zlen)
+    # assign two coordinates which should both be renamed according to the renaming dict
+    coord_da = xr.DataArray(np.random.rand(xlen, ylen), dims=["x", "y"])
+    ds = ds.assign_coords(longitude=coord_da)
+
+    ds_pp = rename_cmip6(ds)
+    # make sure the original lon is intact
+    xr.testing.assert_allclose(ds_pp.lon, ds.lon)
+    assert "longitude" in ds_pp
 
 
 def test_preserve_attrs():

xmip/preprocessing.py

@@ -7,7 +7,7 @@
 import pint_xarray  # noqa: F401
 import xarray as xr
 
-from xmip.utils import _maybe_make_list, cmip6_dataset_id
+from xmip.utils import cmip6_dataset_id
 
 
 # global object for units
@@ -57,58 +57,54 @@ def cmip6_renaming_dict():
     return rename_dict
 
 
-def _invert_dict(rdict):
-    exploded_dict = {}
-    # there is probably a more effective way to 'invert' a dictionary
-    for k, v in rdict.items():
-        v = _maybe_make_list(v)
-        for vv in v:
-            exploded_dict[vv] = k
-    return exploded_dict
-
-
 def rename_cmip6(ds, rename_dict=None):
     """Homogenizes cmip6 dataasets to common naming"""
-    ds = ds.copy()
     attrs = {k: v for k, v in ds.attrs.items()}
+    ds_id = cmip6_dataset_id(ds)
 
     if rename_dict is None:
         rename_dict = cmip6_renaming_dict()
 
-    inverted_rename_dict = _invert_dict(rename_dict)
+    # TODO: Be even stricter here and reset every variable except the one given in the attr
+    # as variable_id
+    # ds_reset = ds.reset_coords()
 
-    ds_reset = ds.reset_coords()
-
-    def _maybe_rename(obj, rdict):
-        return obj.rename({kk: vv for kk, vv in rdict.items() if kk in obj.dims})
+    def _maybe_rename_dims(da, rdict):
+        for di in da.dims:
+            for target, candidates in rdict.items():
+                if di in candidates:
+                    da = da.swap_dims({di: target})
+                    if di in da.coords:
+                        da = da.drop_vars(di)
+        return da
 
     # first take care of the dims and reconstruct a clean ds
     ds = xr.Dataset(
         {
-            k: _maybe_rename(ds_reset[k], inverted_rename_dict)
-            for k in ds_reset.data_vars
+            k: _maybe_rename_dims(ds[k], rename_dict)
+            for k in list(ds.data_vars) + list(set(ds.coords) - set(ds.dims))
         }
     )
 
+    rename_vars = list(set(ds.variables) - set(ds.dims))
+
+    for target, candidates in rename_dict.items():
+        if target not in ds:
+            matching_candidates = [ca for ca in candidates if ca in rename_vars]
+            if len(matching_candidates) > 0:
+                if len(matching_candidates) > 1:
+                    warnings.warn(
+                        f"{ds_id}:While renaming to target `{target}`, more than one candidate was found {matching_candidates}. Renaming {matching_candidates[0]} to {target}. Please double check results."
+                    )
+                ds = ds.rename({matching_candidates[0]: target})
+
     # special treatment for 'lon'/'lat' if there is no 'x'/'y' after renaming process
     for di, co in [("x", "lon"), ("y", "lat")]:
         if di not in ds.dims and co in ds.dims:
             ds = ds.rename({co: di})
 
-    # now rename the variables
-    # try and pass here, cause some of the datasets (MIROC) have like 3 times the same info
-    # e.g. lev/sigma/zlev...not sure this is the best way to handle this with
-    # a silent fail here though...
-    for va in ds.data_vars:
-        try:
-            ds = ds.rename({va: inverted_rename_dict[va]})
-        except Exception as e:
-            warnings.warn(f"Renaming failed with {e}")
-            pass
-
     # restore attributes
     ds.attrs = attrs
-
     return ds