infer dimension ordering within nc_specs using cf-xarray

CHANGELOG.rst

@@ -6,15 +6,16 @@ Changelog
     `Unreleased <https://github.com/CSHS-CWRA/RavenPy>`_ (latest)
     -------------------------------------------------------------
 
-    Contributors:
+    Contributors: David Huard (:user:`huard`).
 
     Changes
     ^^^^^^^
     * No change.
 
     Fixes
     ^^^^^
-    * No change.
+    * In `nc_specs`, set `dim_names_nc` in the order expected by Raven (x, y, t). Previously, we only made sure that `time` was the last dimension, but did not ensure x and y were in the right order. (PR #533)
+
 
 .. _changes_0.19.1:
 

src/ravenpy/config/commands.py

@@ -853,7 +853,12 @@ class ReadFromNetCDF(FlatCommand):
     @field_validator("dim_names_nc")
     @classmethod
     def reorder_time(cls, v):
-        """TODO: Return dimensions as x, y, t. Currently only puts time at the end."""
+        """
+        Return dimensions as x, y, t.
+
+        This is a fail safe because if input files are CF-compliant, dimensions should already
+        have been ordered by `nc_specs`.
+        """
         dims = list(v)
         for time_dim in ("t", "time"):
             if time_dim in dims:

src/ravenpy/config/utils.py

@@ -78,7 +78,7 @@ def nc_specs(
             if v in ds.data_vars:
                 nc_var = ds[v]
                 attrs["var_name_nc"] = v
-                attrs["dim_names_nc"] = nc_var.dims
+                attrs["dim_names_nc"] = infer_dim_names(nc_var)
                 attrs["_time_dim_name_nc"] = ds.cf["time"].name
                 attrs["_dim_size_nc"] = dict(zip(nc_var.dims, nc_var.shape))
                 attrs["units"] = nc_var.attrs.get("units")
@@ -185,3 +185,34 @@ def get_annotations(a):
             yield from get_annotations(arg)
         else:
             yield arg
+
+
+def infer_dim_names(da: xr.DataArray) -> tuple:
+    """
+    Return names of dimensions in dataset in order expected by Raven.
+
+    If 3D, return X, Y, T axes names if they can be inferred from CF conventions.
+    If 2D, return STATION, T
+    """
+    try:
+        if da.ndim == 1:
+            dims = da.cf.axes["T"]
+            if len(dims) != 1:
+                raise ValueError("Should have exactly 1 dimension.")
+
+        elif da.ndim == 2:
+            dims = list(da.dims)
+            tdim = da.cf.axes["T"][0]
+            dims.remove(tdim)
+            dims.append(tdim)
+
+        elif da.ndim == 3:
+            dims = da.cf.axes["X"] + da.cf.axes["Y"] + da.cf.axes["T"]
+            if len(dims) != 3:
+                raise ValueError("Should have exactly 3 dimensions.")
+
+    except:
+        # In case CF inference fails, return the original dims.
+        return da.dims
+
+    return tuple(dims)

tests/test_utils.py

@@ -14,6 +14,24 @@ def test_nc_specs(yangtze):
     f = yangtze.fetch("raven-gr4j-cemaneige/Salmon-River-Near-Prince-George_meteo_daily.nc")
     attrs = nc_specs(f, "PRECIP", station_idx=1, alt_names=("rain",))
     assert "file_name_nc" in attrs
+    assert attrs["dim_names_nc"] == ("time",)
+
+    # 2D with station dimension
+    f = get_local_testdata("raven-gr4j-cemaneige/Salmon-River-Near-Prince-George_meteo_daily_2d.nc")
+    attrs = nc_specs(f, "PRECIP", station_idx=1, alt_names=("rain",))
+    assert attrs["dim_names_nc"] == (
+        "region",
+        "time",
+    )
+
+    # 3D - Since this file is not CF compliant, nc_specs cannot infer the correct dimension order
+    f = get_local_testdata("raven-gr4j-cemaneige/Salmon-River-Near-Prince-George_meteo_daily_3d.nc")
+    attrs = nc_specs(f, "PRECIP", station_idx=1, alt_names=("rain",))
+    assert attrs["dim_names_nc"] == ("time", "lon", "lat")
+
+    f = get_local_testdata("cmip5/tas_Amon_CanESM2_rcp85_r1i1p1_200601-210012_subset.nc")
+    attrs = nc_specs(f, "TEMP_AVE", station_idx=1, engine="netcdf4")
+    assert attrs["dim_names_nc"] == ("lon", "lat", "time")
 
 
 def test_nc_specs_bad(bad_netcdf):