Add guess_coord_axis. (#67)

* Add guess_coord_axis.

* small doc updates

Author: dcherian

cf_xarray/accessor.py

@@ -92,17 +92,6 @@
             "degreesE",
         ),
     },
-    # "regular_expression": {
-    #     "time": r"time[0-9]*",
-    #     "vertical": (
-    #         r"(lv_|bottom_top|sigma|h(ei)?ght|altitude|depth|isobaric|pres|"
-    #         r"isotherm)[a-z_]*[0-9]*"
-    #     ),
-    #     "y": r"y",
-    #     "latitude": r"x?lat[a-z0-9]*",
-    #     "x": r"x",
-    #     "longitude": r"x?lon[a-z0-9]*",
-    # },
 }
 
 # "vertical" is just an alias for "Z"
@@ -112,6 +101,54 @@
 # "long_name" and "standard_name" criteria are the same. For convenience.
 coordinate_criteria["long_name"] = coordinate_criteria["standard_name"]
 
+
+#: regular expressions for guess_coord_axis
+regex = {
+    "time": "time[0-9]*",
+    "vertical": (
+        "(lv_|bottom_top|sigma|h(ei)?ght|altitude|depth|isobaric|pres|"
+        "isotherm)[a-z_]*[0-9]*"
+    ),
+    "Y": "y",
+    "latitude": "y?lat[a-z0-9]*",
+    "X": "x",
+    "longitude": "x?lon[a-z0-9]*",
+}
+regex["Z"] = regex["vertical"]
+regex["T"] = regex["time"]
+
+
+attrs = {
+    "X": {"axis": "X"},
+    "T": {"axis": "T", "standard_name": "time"},
+    "Y": {"axis": "Y"},
+    "Z": {"axis": "Z"},
+    "latitude": {"units": "degrees_north", "standard_name": "latitude"},
+    "longitude": {"units": "degrees_east", "standard_name": "longitude"},
+}
+attrs["time"] = attrs["T"]
+attrs["vertical"] = attrs["Z"]
+
+
+def _is_datetime_like(da: DataArray) -> bool:
+    import numpy as np
+
+    if np.issubdtype(da.dtype, np.datetime64) or np.issubdtype(
+        da.dtype, np.timedelta64
+    ):
+        return True
+
+    try:
+        import cftime
+
+        if isinstance(da.data[0], cftime.datetime):
+            return True
+    except ImportError:
+        pass
+
+    return False
+
+
 # Type for Mapper functions
 Mapper = Callable[[Union[DataArray, Dataset], str], List[str]]
 
@@ -1090,6 +1127,45 @@ def rename_like(
                 variable.attrs["coordinates"] = coordinates
         return self._maybe_to_dataarray(ds)
 
+    def guess_coord_axis(self, verbose: bool = False) -> Union[DataArray, Dataset]:
+        """
+        Automagically guesses X, Y, Z, T, latitude, longitude, and adds
+        appropriate attributes. Uses regexes from Metpy and inspired by Iris
+        function of same name.
+
+        Existing attributes will not be modified.
+
+        Parameters
+        ----------
+        verbose: bool
+            Print extra info to screen
+
+        Returns
+        -------
+        DataArray or Dataset with appropriate attributes added
+        """
+        import re
+
+        obj = self._obj.copy(deep=True)
+        for dim in obj.dims:
+            if _is_datetime_like(obj[dim]):
+                if verbose:
+                    print(
+                        f"I think {dim!r} is of type 'time' since it has a datetime-like type."
+                    )
+                obj[dim].attrs = dict(ChainMap(obj[dim].attrs, attrs["time"]))
+                continue  # prevent second detection
+
+            for axis, pattern in regex.items():
+                # match variable names
+                if re.match(pattern, dim.lower()):
+                    if verbose:
+                        print(
+                            f"I think {dim!r} is of type {axis!r} since it matched {pattern!r}"
+                        )
+                    obj[dim].attrs = dict(ChainMap(obj[dim].attrs, attrs[axis]))
+        return obj
+
 
 @xr.register_dataset_accessor("cf")
 class CFDatasetAccessor(CFAccessor):

cf_xarray/tests/test_accessor.py

@@ -402,3 +402,25 @@ def test_add_bounds(obj, dims):
 def test_docstring():
     assert "One of ('X'" in airds.cf.groupby.__doc__
     assert "One or more of ('X'" in airds.cf.mean.__doc__
+
+
+def test_guess_axis_coord():
+    ds = xr.Dataset()
+    ds["time"] = ("time", pd.date_range("2001-01-01", "2001-04-01"))
+    ds["lon_rho"] = ("lon_rho", [1, 2, 3, 4, 5])
+    ds["lat_rho"] = ("lat_rho", [1, 2, 3, 4, 5])
+    ds["x1"] = ("x1", [1, 2, 3, 4, 5])
+    ds["y1"] = ("y1", [1, 2, 3, 4, 5])
+
+    dsnew = ds.cf.guess_coord_axis()
+    assert dsnew.time.attrs == {"standard_name": "time", "axis": "T"}
+    assert dsnew.lon_rho.attrs == {
+        "standard_name": "longitude",
+        "units": "degrees_east",
+    }
+    assert dsnew.lat_rho.attrs == {
+        "standard_name": "latitude",
+        "units": "degrees_north",
+    }
+    assert dsnew.x1.attrs == {"axis": "X"}
+    assert dsnew.y1.attrs == {"axis": "Y"}

doc/api.rst

@@ -20,6 +20,7 @@ DataArray
     DataArray.cf.describe
     DataArray.cf.get_standard_names
     DataArray.cf.get_valid_keys
+    DataArray.cf.guess_coord_axis
     DataArray.cf.rename_like
 
 Dataset
@@ -37,4 +38,5 @@ Dataset
     Dataset.cf.describe
     Dataset.cf.get_standard_names
     Dataset.cf.get_valid_keys
+    Dataset.cf.guess_coord_axis
     Dataset.cf.rename_like

doc/whats-new.rst

@@ -1,9 +1,16 @@
 What's New
 ----------
+v0.2.1
+======
+
+- Add ``.cf.guess_coord_axis`` to automagically guess axis and coord names, and add
+  appropriate attributes. (:pr:`67`, :issue:`46`). `Deepak Cherian`_.
 
 v0.2.0 (Jul 28, 2020)
 =====================
 
+``cf_xarray`` is now available on conda-forge. Thanks to `Anderson Banihirwe`_ and `Filipe Fernandes`_
+
 - Remap datetime accessor syntax for groupby. E.g. ``.cf.groupby("T.month")`` → ``.cf.groupby("ocean_time.month")``.
   (:pr:`64`, :issue:`6`). `Julia Kent`_.
 - Added ``.cf.rename_like`` to rename matching variables. Only coordinate variables
@@ -33,5 +40,7 @@ v0.1.3
 
 - Support expanding key to multiple dimension names.
 
+.. _`Anderson Banihirwe`: https://github.com/andersy005
 .. _`Deepak Cherian`: https://github.com/dcherian
+.. _`Filipe Fernandes`: https://github.com/ocefpaf
 .. _`Julia Kent`: https://github.com/jukent