Merge pull request #92 from csiro-coasts/convention-coordinates

Convention coordinates

Author: mx-moth

docs/releases/development.rst

@@ -9,3 +9,7 @@ Next release (in development)
   are dropped from the returned dataset,
   or filled with a sensible fill value
   (:pr:`90`).
+* Align automatic coordinate detection of time and depth with CF Conventions.
+  Add :attr:`.Convention.time_coordinate` and :attr:`.Convention.depth_coordinate`,
+  deprecate :meth:`.Convention.get_times()` and :meth:`.Convention.get_depths()`
+  (:pr:`92`).

src/emsarray/conventions/_base.py

@@ -19,6 +19,7 @@
 
 from emsarray import utils
 from emsarray.compat.shapely import SpatialIndex
+from emsarray.exceptions import NoSuchCoordinateError
 from emsarray.operations import depth
 from emsarray.plot import (
     _requires_plot, animate_on_figure, plot_on_figure, polygons_to_collection
@@ -282,22 +283,122 @@ def _get_data_array(self, data_array: DataArrayOrName) -> xr.DataArray:
         else:
             return self.dataset[data_array]
 
+    @cached_property
+    def time_coordinate(self) -> xr.DataArray:
+        """The time coordinate for this dataset.
+
+        Returns
+        -------
+        xarray.DataArray
+            The variable for the time coordinate for this dataset.
+
+        Raises
+        ------
+        exceptions.NoSuchCoordinateError
+            If no time coordinate was found
+
+        See Also
+        --------
+        get_time_name
+        """
+        return self.dataset[self.get_time_name()]
+
+    @cached_property
+    def depth_coordinate(self) -> xr.DataArray:
+        """The depth coordinate for this dataset.
+
+        Returns
+        -------
+        xarray.DataArray
+            The variable for the depth coordinate for this dataset.
+
+        Raises
+        ------
+        exceptions.NoSuchCoordinateError
+            If no depth coordinate was found
+
+        See Also
+        --------
+        get_depth_name
+        """
+        return self.dataset[self.get_depth_name()]
+
     def get_time_name(self) -> Hashable:
-        """Get the name of the time variable in this dataset."""
+        """Get the name of the time variable in this dataset.
+
+        Returns
+        -------
+        Hashable
+            The name of the time coordinate.
+
+        Raises
+        ------
+        exceptions.NoSuchCoordinateError
+            If no time coordinate was found
+
+        Notes
+        -----
+        The CF Conventions state that
+        a time variable is defined by having a `units` attribute
+        formatted according to the UDUNITS package [1]_.
+
+        xarray will find all time variables and convert them to numpy datetimes.
+
+        References
+        ----------
+        .. [1] `CF Conventions v1.10, 4.4 Time Coordinate <https://cfconventions.org/Data/cf-conventions/cf-conventions-1.10/cf-conventions.html#time-coordinate>`_
+        """
         for name, variable in self.dataset.variables.items():
-            if variable.attrs.get('standard_name') == 'time':
-                return name
-        raise KeyError("Dataset does not have a time dimension")
+            # xarray will automatically decode all time variables
+            # and move the 'units' attribute over to encoding to store this change.
+            if 'units' in variable.encoding:
+                units = variable.encoding['units']
+                # A time variable must have units of the form '<units> since <epoc>'
+                if 'since' in units:
+                    # The variable must now be a numpy datetime
+                    if variable.dtype.type == np.datetime64:
+                        return name
+        raise NoSuchCoordinateError("Could not find time coordinate in dataset")
 
     def get_depth_name(self) -> Hashable:
         """Get the name of the layer depth coordinate variable.
+
         For datasets with multiple depth variables, this should be the one that
         represents the centre of the layer, not the bounds.
 
         Note that this is the name of the coordinate variable,
         not the name of the dimension, for datasets where these differ.
+
+        Returns
+        -------
+        Hashable
+            The name of the depth coordinate.
+
+        Raises
+        ------
+        exceptions.NoSuchCoordinateError
+            If no time coordinate was found
+
+        Notes
+        -----
+        The CF Conventions state that
+        a depth variable is identifiable by units of pressure; or
+        the presence of the ``positive`` attribute with value of ``up`` or ``down``
+        [2]_.
+
+        In practice, many datasets do not follow this convention.
+        In addition to checking for the ``positive`` attribute,
+        all coordinates are checked for a ``standard_name: "depth"``,
+        ``coordinate_type: "Z"``, or ``axiz: "Z"``.
+
+        References
+        ----------
+        .. [2] `CF Conventions v1.10, 4.3 Vertical (Height or Depth) Coordinate <https://cfconventions.org/Data/cf-conventions/cf-conventions-1.10/cf-conventions.html#vertical-coordinate>`_
         """
-        return self.get_all_depth_names()[0]
+        try:
+            return self.get_all_depth_names()[0]
+        except IndexError:
+            raise NoSuchCoordinateError("Could not find depth coordinate in dataset")
 
     def get_all_depth_names(self) -> List[Hashable]:
         """Get the names of all depth layers.
@@ -312,7 +413,8 @@ def get_all_depth_names(self) -> List[Hashable]:
             data_array = self.dataset[name]
 
             if not (
-                data_array.attrs.get('axis') == 'Z'
+                data_array.attrs.get('positive', '').lower() in {'up', 'down'}
+                or data_array.attrs.get('axis') == 'Z'
                 or data_array.attrs.get('cartesian_axis') == 'Z'
                 or data_array.attrs.get('coordinate_type') == 'Z'
                 or data_array.attrs.get('standard_name') == 'depth'
@@ -333,27 +435,49 @@ def get_all_depth_names(self) -> List[Hashable]:
 
         return depth_names
 
+    @utils.deprecated(
+        (
+            "Convention.get_depths() is deprecated. "
+            "Use Convention.depth_coordinate.values instead."
+        ),
+        DeprecationWarning,
+    )
     def get_depths(self) -> np.ndarray:
         """Get the depth of each vertical layer in this dataset.
 
+        .. deprecated:: 0.5.0
+            This method is replaced by
+            :attr:`Convention.depth_coordinate.values <Convention.depth_coordinate>`.
+
         Returns
         -------
         :class:`numpy.ndarray`
             An array of depths, one per vertical layer in the dataset.
         """
-        return cast(np.ndarray, self.dataset.variables[self.get_depth_name()].values)
-
+        return cast(np.ndarray, self.depth_coordinate.values)
+
+    @utils.deprecated(
+        (
+            "Convention.get_times() is deprecated. "
+            "Use Convention.time_coordinate.values instead."
+        ),
+        DeprecationWarning,
+    )
     def get_times(self) -> np.ndarray:
         """Get all timesteps in this dataset.
 
+        .. deprecated:: 0.5.0
+            This method is replaced by
+            :attr:`Convention.time_coordinate.values <Convention.time_coordinate>`.
+
         Returns
         -------
         :class:`numpy.ndarray`
             An array of datetimes.
             The datetimes will be whatever native format the dataset uses,
             likely :class:`numpy.datetime64`.
         """
-        return cast(np.ndarray, self.dataset.variables[self.get_time_name()].values)
+        return cast(np.ndarray, self.time_coordinate.values)
 
     @abc.abstractmethod
     def ravel_index(self, index: Index) -> int:

src/emsarray/conventions/shoc.py

@@ -22,6 +22,8 @@
 
 import xarray as xr
 
+from emsarray.exceptions import NoSuchCoordinateError
+
 from ._base import Specificity
 from .arakawa_c import ArakawaC, ArakawaCGridKind
 from .grid import CFGrid2D, CFGrid2DTopology
@@ -48,13 +50,21 @@ class ShocStandard(ArakawaC):
     }
 
     def get_depth_name(self) -> Hashable:
-        return 'z_centre'
+        name = 'z_centre'
+        if name not in self.dataset.variables:
+            raise NoSuchCoordinateError(
+                f"SHOC dataset did not have expected depth coordinate {name!r}")
+        return name
 
     def get_all_depth_names(self) -> List[Hashable]:
         return ['z_centre', 'z_grid']
 
     def get_time_name(self) -> Hashable:
-        return 't'
+        name = 't'
+        if name not in self.dataset.variables:
+            raise NoSuchCoordinateError(
+                f"SHOC dataset did not have expected time coordinate {name!r}")
+        return name
 
     def drop_geometry(self) -> xr.Dataset:
         dataset = super().drop_geometry()
@@ -99,10 +109,18 @@ def check_dataset(cls, dataset: xr.Dataset) -> Optional[int]:
         return Specificity.HIGH
 
     def get_time_name(self) -> Hashable:
-        return 'time'
+        name = 'time'
+        if name not in self.dataset.variables:
+            raise NoSuchCoordinateError(
+                f"SHOC dataset did not have expected time coordinate {name!r}")
+        return name
 
     def get_depth_name(self) -> Hashable:
-        return 'zc'
+        name = 'zc'
+        if name not in self.dataset.variables:
+            raise NoSuchCoordinateError(
+                f"SHOC dataset did not have expected depth coordinate {name!r}")
+        return name
 
     def get_all_depth_names(self) -> List[Hashable]:
         return [self.get_depth_name()]

src/emsarray/exceptions.py

@@ -22,3 +22,11 @@ class ConventionViolationWarning(UserWarning):
     For example, an attribute has an invalid type,
     but is still interpretable.
     """
+
+
+class NoSuchCoordinateError(KeyError, EmsarrayError):
+    """
+    Raised when a dataset does not have a particular coordinate,
+    such as in :attr:`.Convention.time_coordinate` and
+    :attr:`.Convention.depth_coordinate`.
+    """

src/emsarray/utils.py

@@ -15,6 +15,7 @@
 import logging
 import textwrap
 import time
+import warnings
 from types import TracebackType
 from typing import (
     Any, Callable, Hashable, Iterable, List, Literal, Mapping, MutableMapping,
@@ -721,3 +722,13 @@ def make_polygons_with_holes(
         indices=complete_row_indices,
         out=out)
     return out
+
+
+def deprecated(message: str, category: Type[Warning] = DeprecationWarning) -> Callable:
+    def decorator(fn: Callable) -> Callable:
+        @functools.wraps(fn)
+        def wrapped(*args: Any, **kwargs: Any) -> Any:
+            warnings.warn(message, category=category, stacklevel=2)
+            return fn(*args, **kwargs)
+        return wrapped
+    return decorator

tests/conventions/test_base.py

@@ -2,8 +2,9 @@
 
 import dataclasses
 import enum
+import pathlib
 from functools import cached_property
-from typing import Dict, Hashable, List, Optional, Tuple
+from typing import Dict, Hashable, Optional, Tuple
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -14,6 +15,7 @@
 
 from emsarray import masking, utils
 from emsarray.conventions import Convention, SpatialIndexItem
+from emsarray.exceptions import NoSuchCoordinateError
 from emsarray.types import Pathish
 
 
@@ -41,15 +43,6 @@ class SimpleConvention(Convention[SimpleGridKind, SimpleGridIndex]):
     def check_dataset(cls, dataset: xr.Dataset) -> Optional[int]:
         return None
 
-    def get_time_name(self) -> Hashable:
-        return 't'
-
-    def get_depth_name(self) -> Hashable:
-        return 'z'
-
-    def get_all_depth_names(self) -> List[Hashable]:
-        return [self.get_depth_name()]
-
     @cached_property
     def shape(self) -> Tuple[int, int]:
         y, x = map(int, self.dataset['botz'].shape)
@@ -114,6 +107,42 @@ def apply_clip_mask(self, clip_mask: xr.Dataset, work_dir: Pathish) -> xr.Datase
         return masking.mask_grid_dataset(self.dataset, clip_mask, work_dir)
 
 
+def test_get_time_name(datasets: pathlib.Path) -> None:
+    dataset = xr.open_dataset(datasets / 'times.nc')
+    SimpleConvention(dataset).bind()
+    assert dataset.ems.get_time_name() == 'time'
+    xr.testing.assert_equal(dataset.ems.time_coordinate, dataset['time'])
+
+
+def test_get_time_name_missing() -> None:
+    dataset = xr.Dataset()
+    SimpleConvention(dataset).bind()
+    with pytest.raises(NoSuchCoordinateError):
+        dataset.ems.get_time_name()
+
+
+@pytest.mark.parametrize('attrs', [
+    {'positive': 'up'},
+    {'positive': 'DOWN'},
+    {'standard_name': 'depth'},
+    {'axis': 'Z'},
+], ids=lambda a: '{}:{}'.format(*next(iter(a.items()))))
+def test_get_depth_name(attrs: dict) -> None:
+    dataset = xr.Dataset({
+        'name': (['dim'], [0, 1, 2], attrs),
+    })
+    SimpleConvention(dataset).bind()
+    assert dataset.ems.get_depth_name() == 'name'
+    xr.testing.assert_equal(dataset.ems.depth_coordinate, dataset['name'])
+
+
+def test_get_depth_name_missing() -> None:
+    dataset = xr.Dataset()
+    SimpleConvention(dataset).bind()
+    with pytest.raises(NoSuchCoordinateError):
+        dataset.ems.get_depth_name()
+
+
 def test_mask():
     dataset = xr.Dataset({
         'values': (['z', 'y', 'x'], np.random.standard_normal((5, 10, 20))),