Merge branch 'main' into Illviljan-dataset_line_plot

Author: Illviljan

ci/requirements/environment-windows.yml

@@ -10,6 +10,7 @@ dependencies:
   - cftime
   - dask
   - distributed
+  - fsspec!=2021.7.0
   - h5netcdf
   - h5py
   - hdf5

ci/requirements/environment.yml

@@ -12,6 +12,7 @@ dependencies:
   - cftime
   - dask
   - distributed
+  - fsspec!=2021.7.0
   - h5netcdf
   - h5py
   - hdf5

doc/api.rst

@@ -686,6 +686,7 @@ Dataset methods
    open_zarr
    Dataset.to_netcdf
    Dataset.to_pandas
+   Dataset.as_numpy
    Dataset.to_zarr
    save_mfdataset
    Dataset.to_array
@@ -716,6 +717,8 @@ DataArray methods
    DataArray.to_pandas
    DataArray.to_series
    DataArray.to_dataframe
+   DataArray.to_numpy
+   DataArray.as_numpy
    DataArray.to_index
    DataArray.to_masked_array
    DataArray.to_cdms2

doc/whats-new.rst

@@ -56,6 +56,10 @@ New Features
 - Allow removal of the coordinate attribute ``coordinates`` on variables by setting ``.attrs['coordinates']= None``
   (:issue:`5510`).
   By `Elle Smith <https://github.com/ellesmith88>`_.
+- Added :py:meth:`DataArray.to_numpy`, :py:meth:`DataArray.as_numpy`, and :py:meth:`Dataset.as_numpy`. (:pull:`5568`).
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- Units in plot labels are now automatically inferred from wrapped :py:meth:`pint.Quantity` arrays. (:pull:`5561`).
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/core/dataarray.py

@@ -426,12 +426,12 @@ def __init__(
         self._close = None
 
     def _replace(
-        self,
+        self: T_DataArray,
         variable: Variable = None,
         coords=None,
         name: Union[Hashable, None, Default] = _default,
         indexes=None,
-    ) -> "DataArray":
+    ) -> T_DataArray:
         if variable is None:
             variable = self.variable
         if coords is None:
@@ -623,7 +623,16 @@ def __len__(self) -> int:
 
     @property
     def data(self) -> Any:
-        """The array's data as a dask or numpy array"""
+        """
+        The DataArray's data as an array. The underlying array type
+        (e.g. dask, sparse, pint) is preserved.
+
+        See Also
+        --------
+        DataArray.to_numpy
+        DataArray.as_numpy
+        DataArray.values
+        """
         return self.variable.data
 
     @data.setter
@@ -632,13 +641,46 @@ def data(self, value: Any) -> None:
 
     @property
     def values(self) -> np.ndarray:
-        """The array's data as a numpy.ndarray"""
+        """
+        The array's data as a numpy.ndarray.
+
+        If the array's data is not a numpy.ndarray this will attempt to convert
+        it naively using np.array(), which will raise an error if the array
+        type does not support coercion like this (e.g. cupy).
+        """
         return self.variable.values
 
     @values.setter
     def values(self, value: Any) -> None:
         self.variable.values = value
 
+    def to_numpy(self) -> np.ndarray:
+        """
+        Coerces wrapped data to numpy and returns a numpy.ndarray.
+
+        See also
+        --------
+        DataArray.as_numpy : Same but returns the surrounding DataArray instead.
+        Dataset.as_numpy
+        DataArray.values
+        DataArray.data
+        """
+        return self.variable.to_numpy()
+
+    def as_numpy(self: T_DataArray) -> T_DataArray:
+        """
+        Coerces wrapped data and coordinates into numpy arrays, returning a DataArray.
+
+        See also
+        --------
+        DataArray.to_numpy : Same but returns only the data as a numpy.ndarray object.
+        Dataset.as_numpy : Converts all variables in a Dataset.
+        DataArray.values
+        DataArray.data
+        """
+        coords = {k: v.as_numpy() for k, v in self._coords.items()}
+        return self._replace(self.variable.as_numpy(), coords, indexes=self._indexes)
+
     @property
     def _in_memory(self) -> bool:
         return self.variable._in_memory
@@ -931,7 +973,7 @@ def persist(self, **kwargs) -> "DataArray":
         ds = self._to_temp_dataset().persist(**kwargs)
         return self._from_temp_dataset(ds)
 
-    def copy(self, deep: bool = True, data: Any = None) -> "DataArray":
+    def copy(self: T_DataArray, deep: bool = True, data: Any = None) -> T_DataArray:
         """Returns a copy of this array.
 
         If `deep=True`, a deep copy is made of the data array.
@@ -2742,7 +2784,7 @@ def to_masked_array(self, copy: bool = True) -> np.ma.MaskedArray:
         result : MaskedArray
             Masked where invalid values (nan or inf) occur.
         """
-        values = self.values  # only compute lazy arrays once
+        values = self.to_numpy()  # only compute lazy arrays once
         isnull = pd.isnull(values)
         return np.ma.MaskedArray(data=values, mask=isnull, copy=copy)
 

xarray/core/dataset.py

@@ -1323,6 +1323,18 @@ def copy(self, deep: bool = False, data: Mapping = None) -> "Dataset":
 
         return self._replace(variables, attrs=attrs)
 
+    def as_numpy(self: "Dataset") -> "Dataset":
+        """
+        Coerces wrapped data and coordinates into numpy arrays, returning a Dataset.
+
+        See also
+        --------
+        DataArray.as_numpy
+        DataArray.to_numpy : Returns only the data as a numpy.ndarray object.
+        """
+        numpy_variables = {k: v.as_numpy() for k, v in self.variables.items()}
+        return self._replace(variables=numpy_variables)
+
     @property
     def _level_coords(self) -> Dict[str, Hashable]:
         """Return a mapping of all MultiIndex levels and their corresponding

xarray/core/pycompat.py

@@ -1,47 +1,63 @@
 from distutils.version import LooseVersion
+from importlib import import_module
 
 import numpy as np
 
 from .utils import is_duck_array
 
 integer_types = (int, np.integer)
 
-try:
-    import dask
-    import dask.array
-    from dask.base import is_dask_collection
 
-    dask_version = LooseVersion(dask.__version__)
+class DuckArrayModule:
+    """
+    Solely for internal isinstance and version checks.
 
-    # solely for isinstance checks
-    dask_array_type = (dask.array.Array,)
+    Motivated by having to only import pint when required (as pint currently imports xarray)
+    https://github.com/pydata/xarray/pull/5561#discussion_r664815718
+    """
 
-    def is_duck_dask_array(x):
-        return is_duck_array(x) and is_dask_collection(x)
+    def __init__(self, mod):
+        try:
+            duck_array_module = import_module(mod)
+            duck_array_version = LooseVersion(duck_array_module.__version__)
+
+            if mod == "dask":
+                duck_array_type = (import_module("dask.array").Array,)
+            elif mod == "pint":
+                duck_array_type = (duck_array_module.Quantity,)
+            elif mod == "cupy":
+                duck_array_type = (duck_array_module.ndarray,)
+            elif mod == "sparse":
+                duck_array_type = (duck_array_module.SparseArray,)
+            else:
+                raise NotImplementedError
+
+        except ImportError:  # pragma: no cover
+            duck_array_module = None
+            duck_array_version = LooseVersion("0.0.0")
+            duck_array_type = ()
 
+        self.module = duck_array_module
+        self.version = duck_array_version
+        self.type = duck_array_type
+        self.available = duck_array_module is not None
 
-except ImportError:  # pragma: no cover
-    dask_version = LooseVersion("0.0.0")
-    dask_array_type = ()
-    is_duck_dask_array = lambda _: False
-    is_dask_collection = lambda _: False
 
-try:
-    # solely for isinstance checks
-    import sparse
+def is_duck_dask_array(x):
+    if DuckArrayModule("dask").available:
+        from dask.base import is_dask_collection
+
+        return is_duck_array(x) and is_dask_collection(x)
+    else:
+        return False
+
 
-    sparse_version = LooseVersion(sparse.__version__)
-    sparse_array_type = (sparse.SparseArray,)
-except ImportError:  # pragma: no cover
-    sparse_version = LooseVersion("0.0.0")
-    sparse_array_type = ()
+dsk = DuckArrayModule("dask")
+dask_version = dsk.version
+dask_array_type = dsk.type
 
-try:
-    # solely for isinstance checks
-    import cupy
+sp = DuckArrayModule("sparse")
+sparse_array_type = sp.type
+sparse_version = sp.version
 
-    cupy_version = LooseVersion(cupy.__version__)
-    cupy_array_type = (cupy.ndarray,)
-except ImportError:  # pragma: no cover
-    cupy_version = LooseVersion("0.0.0")
-    cupy_array_type = ()
+cupy_array_type = DuckArrayModule("cupy").type

xarray/core/variable.py

@@ -29,10 +29,12 @@
 from .indexing import BasicIndexer, OuterIndexer, VectorizedIndexer, as_indexable
 from .options import _get_keep_attrs
 from .pycompat import (
+    DuckArrayModule,
     cupy_array_type,
     dask_array_type,
     integer_types,
     is_duck_dask_array,
+    sparse_array_type,
 )
 from .utils import (
     NdimSizeLenMixin,
@@ -259,7 +261,7 @@ def _as_array_or_item(data):
 
     TODO: remove this (replace with np.asarray) once these issues are fixed
     """
-    data = data.get() if isinstance(data, cupy_array_type) else np.asarray(data)
+    data = np.asarray(data)
     if data.ndim == 0:
         if data.dtype.kind == "M":
             data = np.datetime64(data, "ns")
@@ -1069,6 +1071,30 @@ def chunk(self, chunks={}, name=None, lock=False):
 
         return self._replace(data=data)
 
+    def to_numpy(self) -> np.ndarray:
+        """Coerces wrapped data to numpy and returns a numpy.ndarray"""
+        # TODO an entrypoint so array libraries can choose coercion method?
+        data = self.data
+
+        # TODO first attempt to call .to_numpy() once some libraries implement it
+        if isinstance(data, dask_array_type):
+            data = data.compute()
+        if isinstance(data, cupy_array_type):
+            data = data.get()
+        # pint has to be imported dynamically as pint imports xarray
+        pint_array_type = DuckArrayModule("pint").type
+        if isinstance(data, pint_array_type):
+            data = data.magnitude
+        if isinstance(data, sparse_array_type):
+            data = data.todense()
+        data = np.asarray(data)
+
+        return data
+
+    def as_numpy(self: VariableType) -> VariableType:
+        """Coerces wrapped data into a numpy array, returning a Variable."""
+        return self._replace(data=self.to_numpy())
+
     def _as_sparse(self, sparse_format=_default, fill_value=dtypes.NA):
         """
         use sparse-array as backend.

xarray/plot/plot.py

@@ -430,7 +430,7 @@ def line(
 
     # Remove pd.Intervals if contained in xplt.values and/or yplt.values.
     xplt_val, yplt_val, x_suffix, y_suffix, kwargs = _resolve_intervals_1dplot(
-        xplt.values, yplt.values, kwargs
+        xplt.to_numpy(), yplt.to_numpy(), kwargs
     )
     xlabel = label_from_attrs(xplt, extra=x_suffix)
     ylabel = label_from_attrs(yplt, extra=y_suffix)
@@ -449,7 +449,7 @@ def line(
         ax.set_title(darray._title_for_slice())
 
     if darray.ndim == 2 and add_legend:
-        ax.legend(handles=primitive, labels=list(hueplt.values), title=hue_label)
+        ax.legend(handles=primitive, labels=list(hueplt.to_numpy()), title=hue_label)
 
     # Rotate dates on xlabels
     # Do this without calling autofmt_xdate so that x-axes ticks
@@ -551,7 +551,7 @@ def hist(
     """
     ax = get_axis(figsize, size, aspect, ax)
 
-    no_nan = np.ravel(darray.values)
+    no_nan = np.ravel(darray.to_numpy())
     no_nan = no_nan[pd.notnull(no_nan)]
 
     primitive = ax.hist(no_nan, **kwargs)
@@ -1153,8 +1153,8 @@ def newplotfunc(
             dims = (yval.dims[0], xval.dims[0])
 
         # better to pass the ndarrays directly to plotting functions
-        xval = xval.values
-        yval = yval.values
+        xval = xval.to_numpy()
+        yval = yval.to_numpy()
 
         # May need to transpose for correct x, y labels
         # xlab may be the name of a coord, we have to check for dim names

xarray/plot/utils.py

@@ -9,6 +9,7 @@
 import pandas as pd
 
 from ..core.options import OPTIONS
+from ..core.pycompat import DuckArrayModule
 from ..core.utils import is_scalar
 
 try:
@@ -474,12 +475,20 @@ def label_from_attrs(da, extra=""):
     else:
         name = ""
 
-    if da.attrs.get("units"):
-        units = " [{}]".format(da.attrs["units"])
-    elif da.attrs.get("unit"):
-        units = " [{}]".format(da.attrs["unit"])
+    def _get_units_from_attrs(da):
+        if da.attrs.get("units"):
+            units = " [{}]".format(da.attrs["units"])
+        elif da.attrs.get("unit"):
+            units = " [{}]".format(da.attrs["unit"])
+        else:
+            units = ""
+        return units
+
+    pint_array_type = DuckArrayModule("pint").type
+    if isinstance(da.data, pint_array_type):
+        units = " [{}]".format(str(da.data.units))
     else:
-        units = ""
+        units = _get_units_from_attrs(da)
 
     return "\n".join(textwrap.wrap(name + extra + units, 30))
 
@@ -896,7 +905,7 @@ def _get_nice_quiver_magnitude(u, v):
     import matplotlib as mpl
 
     ticker = mpl.ticker.MaxNLocator(3)
-    mean = np.mean(np.hypot(u.values, v.values))
+    mean = np.mean(np.hypot(u.to_numpy(), v.to_numpy()))
     magnitude = ticker.tick_values(0, mean)[-2]
     return magnitude
 

xarray/tests/__init__.py

@@ -83,6 +83,7 @@ def LooseVersion(vstring):
 has_numbagg, requires_numbagg = _importorskip("numbagg")
 has_seaborn, requires_seaborn = _importorskip("seaborn")
 has_sparse, requires_sparse = _importorskip("sparse")
+has_cupy, requires_cupy = _importorskip("cupy")
 has_cartopy, requires_cartopy = _importorskip("cartopy")
 # Need Pint 0.15 for __dask_tokenize__ tests for Quantity wrapped Dask Arrays
 has_pint_0_15, requires_pint_0_15 = _importorskip("pint", minversion="0.15")