Expand more functions that accept DataArrayOrName

Two new functions `emsarray.utils.name_to_data_array` and
`emsarray.utils.data_array_to_name` have been added to accomplish this.
`Convention._get_data_array()` has been deprecaterd in favour of
`name_to_data_array`.

Author: mx-moth

docs/api/types.rst

@@ -1,9 +1,6 @@
-.. module:: emsarray.types
-
 ==============
 emsarray.types
 ==============
 
 .. automodule:: emsarray.types
-   :noindex:
    :members:

docs/conf.py

@@ -74,6 +74,7 @@
 napoleon_type_aliases = {
     'xarray.core.dataset.Dataset': ':class:`~xarray.Dataset',
     'xarray.core.dataarray.DataArray': ':class:`~xarray.DataArray',
+    'DataArrayOrName': ':data:`data array or name <emsarray.types.DataArrayOrName>`',
 }
 
 

docs/releases/development.rst

@@ -30,3 +30,6 @@ Next release (in development)
   (:pr:`137`).
 * Lint Python code in `docs/` and `scripts/`
   (:pr:`141`).
+* Add :func:`emsarray.utils.name_to_data_array()` and :func:`~emsarray.utils.data_array_to_name()` functions.
+  Allow more functions to interchangeably take either a data array or the name of a data array
+  (:pr:`142`).

src/emsarray/conventions/_base.py

@@ -25,7 +25,7 @@
     polygons_to_collection
 )
 from emsarray.state import State
-from emsarray.types import Bounds, Pathish
+from emsarray.types import Bounds, DataArrayOrName, Pathish
 
 if TYPE_CHECKING:
     # Import these optional dependencies only during type checking
@@ -39,8 +39,6 @@
 logger = logging.getLogger(__name__)
 
 
-DataArrayOrName = Union[Hashable, xarray.DataArray]
-
 #: Some type that can enumerate the different :ref:`grid types <grids>`
 #: present in a dataset.
 #: This can be an :class:`enum.Enum` listing each different kind of grid.
@@ -266,19 +264,15 @@ def bind(self) -> None:
                 "cannot assign a new convention.")
         state.bind_convention(self)
 
-    @abc.abstractmethod
+    @utils.deprecated(
+        (
+            "Convention._get_data_array() has been deprecated. "
+            "Use emsarray.utils.name_to_data_array() instead."
+        ),
+        DeprecationWarning,
+    )
     def _get_data_array(self, data_array: DataArrayOrName) -> xarray.DataArray:
-        """
-        Utility to help get a data array for this dataset.
-        If a string is passed in, the matching data array is fetched from the dataset.
-        If a data array is passed in,
-        it is inspected to ensure the surface dimensions align
-        before being returned as-is.
-
-        This is useful for methods that support being passed either
-        the name of a data array or a data array instance.
-        """
-        pass
+        return utils.name_to_data_array(self.dataset, data_array)
 
     @cached_property
     def time_coordinate(self) -> xarray.DataArray:
@@ -902,10 +896,10 @@ def plot_on_figure(
         ----------
         figure : matplotlib.figure.Figure
             The :class:`~matplotlib.figure.Figure` instance to plot this on.
-        scalar : xarray.DataArray or str
+        scalar : DataArrayOrName
             The :class:`~xarray.DataArray` to plot,
             or the name of an existing DataArray in this Dataset.
-        vector : tuple of xarray.DataArray or str
+        vector : tuple of DataArrayOrName
             A tuple of the *u* and *v* components of a vector.
             The components should be a :class:`~xarray.DataArray`,
             or the name of an existing DataArray in this Dataset.
@@ -918,10 +912,10 @@ def plot_on_figure(
         :func:`.plot.plot_on_figure` : The underlying implementation
         """
         if scalar is not None:
-            kwargs['scalar'] = self._get_data_array(scalar)
+            kwargs['scalar'] = utils.name_to_data_array(self.dataset, scalar)
 
         if vector is not None:
-            kwargs['vector'] = tuple(map(self._get_data_array, vector))
+            kwargs['vector'] = tuple(utils.name_to_data_array(self.dataset, v) for v in vector)
 
         if title is not None:
             kwargs['title'] = title
@@ -977,7 +971,7 @@ def animate_on_figure(
         ----------
         figure : matplotlib.figure.Figure
             The :class:`matplotlib.figure.Figure` to plot the animation on
-        data_array : Hashable or xarray.DataArray
+        data_array : DataArrayOrName
             The :class:`xarray.DataArray` to plot.
             If a string is passed in,
             the variable with that name is taken from :attr:`dataset`.
@@ -1006,26 +1000,26 @@ def animate_on_figure(
 
         if coordinate is None:
             # Assume the user wants to plot along the time axis by default.
-            coordinate = self.get_time_name()
-        if isinstance(coordinate, xarray.DataArray):
-            utils.check_data_array_dimensions_match(
-                self.dataset, coordinate, dimensions=coordinate.dims)
-
-        coordinate = self._get_data_array(coordinate)
+            coordinate = self.time_coordinate
+        else:
+            coordinate = utils.name_to_data_array(self.dataset, coordinate)
 
         if len(coordinate.dims) != 1:
             raise ValueError("Coordinate variable must be one dimensional")
 
         coordinate_dim = coordinate.dims[0]
 
         if scalar is not None:
-            scalar = self._get_data_array(scalar)
+            scalar = utils.name_to_data_array(self.dataset, scalar)
             if coordinate_dim not in scalar.dims:
                 raise ValueError("Scalar dimensions do not match coordinate axis to animate along")
             kwargs['scalar'] = scalar
 
         if vector is not None:
-            vector = (self._get_data_array(vector[0]), self._get_data_array(vector[1]))
+            vector = (
+                utils.name_to_data_array(self.dataset, vector[0]),
+                utils.name_to_data_array(self.dataset, vector[1]),
+            )
             if not all(coordinate_dim in component.dims for component in vector):
                 raise ValueError("Vector dimensions do not match coordinate axis to animate along")
             kwargs['vector'] = vector
@@ -1119,7 +1113,7 @@ def make_poly_collection(
                     "Can not pass both `data_array` and `array` to make_poly_collection"
                 )
 
-            data_array = self._get_data_array(data_array)
+            data_array = utils.name_to_data_array(self.dataset, data_array)
 
             data_array = self.ravel(data_array)
             if len(data_array.dims) > 1:
@@ -1189,7 +1183,8 @@ def make_quiver(
         values = numpy.nan, numpy.nan
 
         if u is not None and v is not None:
-            u, v = self._get_data_array(u), self._get_data_array(v)
+            u = utils.name_to_data_array(self.dataset, u)
+            v = utils.name_to_data_array(self.dataset, v)
 
             if u.dims != v.dims:
                 raise ValueError(
@@ -1539,15 +1534,15 @@ def drop_geometry(self) -> xarray.Dataset:
         """
         return self.dataset.drop_vars(self.get_all_geometry_names())
 
-    def select_variables(self, variables: Iterable[Hashable]) -> xarray.Dataset:
+    def select_variables(self, variables: Iterable[DataArrayOrName]) -> xarray.Dataset:
         """Select only a subset of the variables in this dataset, dropping all others.
 
         This will keep all coordinate variables and all geometry variables.
 
         Parameters
         ----------
-        variables : iterable of Hashable
-            The names of all data variables to select.
+        variables : iterable of DataArrayOrName
+            The data variables to select.
 
         Returns
         -------
@@ -1570,6 +1565,7 @@ def select_variables(self, variables: Iterable[Hashable]) -> xarray.Dataset:
             keep_vars.add(self.get_time_name())
         except NoSuchCoordinateError:
             pass
+        keep_vars = {utils.data_array_to_name(self.dataset, v) for v in keep_vars}
         return self.dataset.drop_vars(all_vars - keep_vars)
 
     @abc.abstractmethod
@@ -1701,7 +1697,7 @@ def ocean_floor(self) -> xarray.Dataset:
         """An alias for :func:`emsarray.operations.depth.ocean_floor`"""
         return depth.ocean_floor(
             self.dataset, self.get_all_depth_names(),
-            non_spatial_variables=[self.get_time_name()])
+            non_spatial_variables=[self.time_coordinate])
 
     def normalize_depth_variables(
         self,
@@ -1786,23 +1782,6 @@ def get_grid_kind(self, data_array: xarray.DataArray) -> GridKind:
                 return kind
         raise ValueError("Unknown grid kind")
 
-    def _get_data_array(self, data_array: DataArrayOrName) -> xarray.DataArray:
-        if isinstance(data_array, xarray.DataArray):
-            grid_kind = self.get_grid_kind(data_array)
-            grid_dimensions = self.grid_dimensions[grid_kind]
-            for dimension in grid_dimensions:
-                # The data array already has matching dimension names
-                # as we found the grid kind using `Convention.get_grid_kind()`.
-                if self.dataset.sizes[dimension] != data_array.sizes[dimension]:
-                    raise ValueError(
-                        f"Mismatched dimension {dimension!r}, "
-                        "dataset has size {self.dataset.sizes[dimension]} but "
-                        "data array has size {data_array.sizes[dimension]}!"
-                    )
-            return data_array
-        else:
-            return self.dataset[data_array]
-
     @abc.abstractmethod
     def unpack_index(self, index: Index) -> tuple[GridKind, Sequence[int]]:
         """Convert a native index in to a grid kind and dimension indices.

src/emsarray/operations/depth.py

@@ -5,19 +5,20 @@
 import warnings
 from collections import defaultdict
 from collections.abc import Hashable
-from typing import Optional, cast
+from typing import Iterable, Optional, cast
 
 import numpy
 import xarray
 
 from emsarray import utils
+from emsarray.types import DataArrayOrName
 
 
 def ocean_floor(
     dataset: xarray.Dataset,
-    depth_variables: list[Hashable],
+    depth_coordinates: Iterable[DataArrayOrName],
     *,
-    non_spatial_variables: Optional[list[Hashable]] = None,
+    non_spatial_variables: Optional[Iterable[DataArrayOrName]] = None,
 ) -> xarray.Dataset:
     """Make a new :class:`xarray.Dataset` reduced along the given depth
     coordinates to only contain values along the ocean floor.
@@ -26,12 +27,12 @@ def ocean_floor(
     ----------
     dataset
         The dataset to reduce.
-    depth_variables
-        The names of depth coordinate variables.
+    depth_coordinates : iterable of DataArrayOrName
+        The depth coordinate variables.
         For supported conventions, use :meth:`.Convention.get_all_depth_names()`.
-    non_spatial_variables
+    non_spatial_variables : iterable of DataArrayOrName
         Optional.
-        A list of the names of any non-spatial coordinate variables, such as time.
+        A list of any non-spatial coordinate variables, such as time.
         The ocean floor is assumed to be static across non-spatial dimensions.
         For supported conventions, use :meth:`.Convention.get_time_name()`.
 
@@ -74,7 +75,7 @@ def ocean_floor(
 
         >>> operations.ocean_floor(
         ...     big_dataset['temp'].isel(record=0).to_dataset(),
-        ...     depth_variables=big_dataset.ems.get_all_depth_names())
+        ...     depth_coordinates=big_dataset.ems.get_all_depth_names())
         <xarray.Dataset>
         Dimensions:  (y: 5, x: 5)
         Coordinates:
@@ -106,15 +107,17 @@ def ocean_floor(
     # and that for a combination of depth dimension and spatial dimensions,
     # the ocean floor is static.
 
+    depth_coordinates = list(depth_coordinates)
+
     dataset = normalize_depth_variables(
-        dataset, depth_variables,
+        dataset, depth_coordinates,
         positive_down=True, deep_to_shallow=False)
 
     if non_spatial_variables is None:
         non_spatial_variables = []
 
     # The name of all the relevant _dimensions_, not _coordinates_
-    depth_dimensions = utils.dimensions_from_coords(dataset, depth_variables)
+    depth_dimensions = utils.dimensions_from_coords(dataset, depth_coordinates)
     non_spatial_dimensions = utils.dimensions_from_coords(dataset, non_spatial_variables)
 
     for depth_dimension in sorted(depth_dimensions, key=hash):
@@ -196,7 +199,7 @@ def _find_ocean_floor_indices(
 
 def normalize_depth_variables(
     dataset: xarray.Dataset,
-    depth_variables: list[Hashable],
+    depth_variables: Iterable[DataArrayOrName],
     *,
     positive_down: Optional[bool] = None,
     deep_to_shallow: Optional[bool] = None,
@@ -216,10 +219,8 @@ def normalize_depth_variables(
     ----------
     dataset : xarray.Dataset
         The dataset to normalize
-    depth_variables : list of Hashable
-        The names of the depth coordinate variables.
-        This should be the names of the variables, not the dimensions,
-        for datasets where these differ.
+    depth_variables : iterable of DataArrayOrName
+        The depth coordinate variables.
     positive_down : bool, optional
         If True, positive values will indicate depth below the surface.
         If False, negative values indicate depth below the surface.
@@ -240,8 +241,9 @@ def normalize_depth_variables(
     :meth:`.Convention.get_all_depth_names`
     """
     new_dataset = dataset.copy()
-    for name in depth_variables:
-        variable = dataset[name]
+    for variable in depth_variables:
+        variable = utils.name_to_data_array(dataset, variable)
+        name = variable.name
         if len(variable.dims) != 1:
             raise ValueError(
                 f"Can't normalize multidimensional depth variable {name!r} "

src/emsarray/transect.py

@@ -1,5 +1,5 @@
 import dataclasses
-from collections.abc import Hashable, Iterable
+from collections.abc import Iterable
 from functools import cached_property
 from typing import Any, Callable, Generic, Optional, Union, cast
 
@@ -18,8 +18,8 @@
 
 from emsarray.conventions import Convention, Index
 from emsarray.plot import _requires_plot, make_plot_title
-from emsarray.types import Landmark
-from emsarray.utils import move_dimensions_to_end
+from emsarray.types import DataArrayOrName, Landmark
+from emsarray.utils import move_dimensions_to_end, name_to_data_array
 
 # Useful for calculating distances in a AzimuthalEquidistant projection
 # centred on some point:
@@ -116,13 +116,13 @@ def __init__(
         self,
         dataset: xarray.Dataset,
         line: shapely.LineString,
-        depth: Optional[Union[Hashable, xarray.DataArray]] = None,
+        depth: Optional[DataArrayOrName] = None,
     ):
         self.dataset = dataset
         self.convention = dataset.ems
         self.line = line
         if depth is not None:
-            self.depth = self.convention._get_data_array(depth)
+            self.depth = name_to_data_array(dataset, depth)
         else:
             self.depth = self.convention.depth_coordinate
 

src/emsarray/types.py

@@ -3,9 +3,11 @@
 """
 
 import os
+from collections.abc import Hashable
 from typing import Union
 
 import shapely
+import xarray
 
 #: Something that can be used as a path.
 Pathish = Union[os.PathLike, str]
@@ -17,3 +19,6 @@
 #: A landmark for a plot.
 #: This is a tuple of the landmark name and and its location.
 Landmark = tuple[str, shapely.Point]
+
+#: Either an :class:`xarray.DataArray`, or the name of a data array in a dataset.
+DataArrayOrName = Union[Hashable, xarray.DataArray]